from pprint import pprint
from datetime import timedelta
import numpy
import dateutil.parser
import re
from collections import defaultdict

guards_asleep = set()
guard_sleep_duration = defaultdict(int)
guard_sleep_times = defaultdict(lambda: numpy.zeros(24 * 60, dtype=numpy.int))

with open("Day4.in", "r") as f:
    lines = sorted(f.readlines())

RE_TIME = re.compile(r"\[(.*)\]")
RE_SHIFT = re.compile(r".*#(\d+).*")

current_guard = None
last_minute = None
for line in lines:
    time_match = RE_TIME.match(line)

    time = dateutil.parser.parse(time_match.group(1))

    minute = time.time().minute
    if last_minute is None:
        last_minute = minute
    delta = minute - last_minute

    if current_guard in guards_asleep:
        guard_sleep_duration[current_guard] += delta
        for i in range(last_minute, minute):
            guard_sleep_times[current_guard][i] += 1

    last_minute = minute

    match = RE_SHIFT.match(line)
    if match:
        current_guard = int(match.group(1))

    if "wakes up" in line:
        guards_asleep.discard(current_guard)

    if "falls asleep" in line:
        guards_asleep.add(current_guard)

a = numpy.array([numpy.insert(value, 0, key) for (key, value) in guard_sleep_times.items()])
x, y = numpy.unravel_index(a[:, 1:].argmax(), a[:, 1:].shape)

guard = a[x, 0]
print(guard, y, guard * y)