import datetime
from collections import defaultdict

from tools.aoc import AOCDay
from typing import Any


class Day(AOCDay):
    inputs = [
        [
            (240, "test_input4_1"),
            (95199, "input4")
        ],
        [
            (4455, "test_input4_1"),
            (7887, "input4")
        ]
    ]

    def get_guards(self) -> dict:
        active_guard = 0
        sleep_start = 0
        guards = {}

        for event in sorted(self.getInput()):
            time, msg = event.split("] ")
            timestamp = datetime.datetime.strptime(time, "[%Y-%m-%d %H:%M")
            if msg.startswith("Guard"):
                active_guard = int(msg.split()[1][1:])
                if active_guard not in guards:
                    guards[active_guard] = {
                        'total': 0,
                        'minutes': defaultdict(int)
                    }
            elif msg.startswith("falls asleep"):
                if timestamp.hour == 0:
                    sleep_start = timestamp.minute
            elif msg.startswith("wakes up"):
                guards[active_guard]['total'] += timestamp.minute - sleep_start
                for m in range(sleep_start, timestamp.minute):
                    guards[active_guard]['minutes'][m] += 1

        return guards

    def part1(self) -> Any:
        guards = self.get_guards()
        max_guard = 0
        max_minutes = 0
        most_minute = 0
        for guard, data in guards.items():
            if data['total'] < max_minutes:
                continue
            max_guard = guard
            max_minutes = data['total']
            most_minute = sorted(data['minutes'], key=lambda k: data['minutes'][k], reverse=True)[0]

        return max_guard * most_minute

    def part2(self) -> Any:
        guards = self.get_guards()
        max_guard = 0
        max_minute = 0
        max_time = 0
        for guard, data in guards.items():
            if not data['minutes']:
                continue
            most_minute = sorted(data['minutes'], key=lambda k: data['minutes'][k], reverse=True)[0]
            if data['minutes'][most_minute] < max_time:
                continue
            max_guard = guard
            max_time = data['minutes'][most_minute]
            max_minute = most_minute

        return max_guard * max_minute


if __name__ == '__main__':
    day = Day(2018, 4)
    day.run(verbose=True)