from collections import defaultdict, Counter
from datetime import datetime
import re


with open('data/04.txt') as f:
    regex = re.compile(r'\[(\d{4}-\d{2}-\d{2} \d{2}:\d{2})\](?: Guard #(\d+)){0,1} (.+ .+)')
    events = []
    for line in f:
        dt, guard_id, event_type = regex.match(line).groups()
        dt = datetime.strptime(dt, '%Y-%m-%d %H:%M')
        guard_id = 0 if guard_id is None else int(guard_id)
        events.append([dt, guard_id, event_type])
    events.sort(key=lambda e: e[0])


spans = []
sleep_start = None
current_guard_id = None
for dt, guard_id, event_type in events:
    if event_type == 'begins shift':
        current_guard_id = guard_id
    elif event_type == 'falls asleep':
        sleep_start = dt
    elif event_type == 'wakes up':
        spans.append((current_guard_id, sleep_start, dt))

by_minute = defaultdict(lambda: defaultdict(list))
by_guard = defaultdict(lambda: defaultdict(list))
total_time_asleep = Counter()
for guard_id, start, end in spans:
    for i in range((end - start).seconds // 60):
        minute = start.minute + i
        by_minute[minute][guard_id].append(start.date())
        by_guard[guard_id][minute].append(start.date())
        total_time_asleep[guard_id] += 1

sleepiest_guard = total_time_asleep.most_common()[0][0]
sleepiest_minute, dates = sorted(
    (m for m in by_guard[sleepiest_guard].items()), 
    key=lambda x: len(x[1]), 
    reverse=True
)[0]

print(f'Part 1: {sleepiest_guard * sleepiest_minute}\n')


sleepiest_minute = defaultdict(lambda: (0, 0))
for guard_id, minutes in by_guard.items():
    for minute, dates in minutes.items():
        times_asleep = len(dates)
        if times_asleep > sleepiest_minute[minute][1]:
            sleepiest_minute[minute] = (guard_id, times_asleep)

minute, (guard_id, num_times_asleep) = sorted(sleepiest_minute.items(), key=lambda x: x[1][1], reverse=True)[0]

print(f'Part 2: {minute * guard_id}\n')