day-04.rs

use chrono::{offset::TimeZone, DateTime, Duration, Timelike, Utc};
use hashbrown::HashMap;
use logos::{Lexer, Logos};

#[derive(Copy, Clone, PartialEq, Logos)]
enum Token {
    #[end]
    End,
    #[error]
    Error,
    #[token = "["]
    DateStart,
    #[regex = "[0-9][0-9][0-9][0-9]-[0-1][0-9]-[0-3][0-9] [0-2][0-9]:[0-5][0-9]"]
    Date,
    #[token = "]"]
    DateEnd,
    #[token = "wakes up"]
    WakesUp,
    #[token = "falls asleep"]
    FallsAsleep,
    #[regex = "Guard"]
    GuardPre,
    #[regex = "#[0-9]+"]
    GuardNr,
    #[regex = "begins shift"]
    GuardPost,
}

#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
pub struct Action {
    pub date: DateTime<Utc>,
    pub kind: ActionKind,
}

#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
pub enum ActionKind {
    WakesUp,
    FallsAsleep,
    BeginsShift(u16),
}

pub struct ActionParser<'a> {
    lexer: Lexer<Token, &'a str>,
}

impl<'a> ActionParser<'a> {
    fn consume<T>(&mut self, expected: Token, f: impl FnOnce(&'a str) -> Option<T>) -> Option<T> {
        if self.lexer.token != expected {
            return None;
        }
        let val = f(self.lexer.slice());
        self.lexer.advance();
        val
    }
    fn skip(&mut self, expected: Token) -> Option<()> {
        self.consume(expected, |_| Some(()))
    }
}

impl<'a> Iterator for ActionParser<'a> {
    type Item = Action;

    fn next(&mut self) -> Option<Action> {
        self.skip(Token::DateStart)?;
        let date = self.consume(Token::Date, |s| Utc.datetime_from_str(s, "%F %R").ok())?;
        self.skip(Token::DateEnd)?;
        let token = self.lexer.token;
        self.lexer.advance();
        let kind = match token {
            Token::WakesUp => ActionKind::WakesUp,
            Token::FallsAsleep => ActionKind::FallsAsleep,
            Token::GuardPre => {
                let id = self.consume(Token::GuardNr, |s| s[1..].parse().ok())?;
                self.skip(Token::GuardPost)?;
                ActionKind::BeginsShift(id)
            }
            _ => return None,
        };
        Some(Action { date, kind })
    }
}

pub fn actions(input: &str) -> impl Iterator<Item = Action> + '_ {
    ActionParser {
        lexer: Lexer::new(input),
    }
}

pub fn actions_sorted(input: &str) -> Vec<Action> {
    let mut actions = actions(input).collect::<Vec<_>>();
    actions.sort_unstable();
    actions
}

pub struct AsleepTime {
    pub guard_id: u16,
    pub start: DateTime<Utc>,
    pub end: DateTime<Utc>,
}

pub fn asleep_times(actions_sorted: &[Action]) -> impl Iterator<Item = AsleepTime> + '_ {
    let mut guard_id = 0;
    actions_sorted.windows(2).filter_map(move |window| {
        match (window[0].kind, window[1].kind) {
            (ActionKind::BeginsShift(id), _) => guard_id = id,
            (ActionKind::FallsAsleep, ActionKind::WakesUp) => {
                return Some(AsleepTime {
                    guard_id,
                    start: window[0].date,
                    end: window[1].date,
                })
            }
            _ => {}
        }
        None
    })
}

struct GuardSchedule {
    total: Duration,
    minutes: [u8; 60],
}

impl Default for GuardSchedule {
    fn default() -> Self {
        GuardSchedule {
            total: Duration::zero(),
            minutes: [0; 60],
        }
    }
}

struct MinuteMostAsleep {
    minute: usize,
    count: u8,
}

impl GuardSchedule {
    fn push_asleep_time(&mut self, asleep_time: AsleepTime) {
        let duration = asleep_time.end - asleep_time.start;
        self.total = self.total + duration;

        let (start, end) = (asleep_time.start.minute(), asleep_time.end.minute());
        if start > end {
            for counter in &mut self.minutes[start as usize..] {
                *counter += 1;
            }
            for counter in &mut self.minutes[..end as usize] {
                *counter += 1;
            }
        } else {
            for counter in &mut self.minutes[start as usize..end as usize] {
                *counter += 1;
            }
        }
    }

    fn minute_most_asleep(&self) -> MinuteMostAsleep {
        let (minute, count) = self
            .minutes
            .iter()
            .cloned()
            .enumerate()
            .max_by_key(|&(_, count)| count)
            .unwrap();
        MinuteMostAsleep { minute, count }
    }
}

#[derive(Default)]
struct CompleteSchedule {
    schedule: HashMap<u16, GuardSchedule>,
}

impl CompleteSchedule {
    fn push_asleep_time(&mut self, asleep_time: AsleepTime) {
        self.schedule
            .entry(asleep_time.guard_id)
            .or_default()
            .push_asleep_time(asleep_time);
    }

    fn from_asleep_times(asleep_times: impl IntoIterator<Item = AsleepTime>) -> Self {
        let mut schedule = Self::default();
        for asleep_time in asleep_times {
            schedule.push_asleep_time(asleep_time);
        }
        schedule
    }

    fn guard_with_most_time_asleep(&self) -> Option<(u16, &GuardSchedule)> {
        let (&id, schedule) = self.schedule.iter().max_by_key(|(_, s)| s.total)?;
        Some((id, schedule))
    }

    fn guard_with_minute_most_asleep(&self) -> Option<(u16, MinuteMostAsleep)> {
        self.schedule
            .iter()
            .map(|(&id, schedule)| (id, schedule.minute_most_asleep()))
            .max_by_key(|(_, minute)| minute.count)
    }
}

pub fn part1(actions_sorted: &[Action]) -> Option<u16> {
    let schedule = CompleteSchedule::from_asleep_times(asleep_times(actions_sorted));

    let (guard_id, guard_schedule) = schedule.guard_with_most_time_asleep()?;
    let minute = guard_schedule.minute_most_asleep().minute as u16;

    Some(guard_id * minute)
}

pub fn part2(actions_sorted: &[Action]) -> Option<u16> {
    let schedule = CompleteSchedule::from_asleep_times(asleep_times(actions_sorted));

    let (guard_id, minute) = schedule.guard_with_minute_most_asleep()?;
    let minute = minute.minute as u16;

    Some(guard_id * minute)
}