2022 Day 11

use std::env;
use std::io::{self};

extern crate regex;
use regex::Regex;

#[derive(Debug)]
enum Operation {
    Add(i64),
    Multiply(i64),
    Square,
}
#[derive(Debug)]
struct Monkey {
    pub items: Vec<i64>,
    pub operation: Operation,
    pub test_value: i64,
    pub target_true: i64,
    pub target_false: i64,
    pub inspections: i64,
}
impl From<&str> for Monkey {
    fn from(s: &str) -> Self {
        let lines: Vec<_> = s.split("\n").collect();

        // Regexes
        let re_starting  = Regex::new(r"Starting items: ([0-9 ,]+)").unwrap();
        let re_operation = Regex::new(r"Operation: new = old (\+|\*) (old|[\d]+)").unwrap();
        let re_test      = Regex::new(r"Test: divisible by ([\d]+)").unwrap();
        let re_true      = Regex::new(r"If true: throw to monkey ([\d]+)").unwrap();
        let re_false     = Regex::new(r"If false: throw to monkey ([\d]+)").unwrap();

        // Captures
        let starting_captures = re_starting.captures(&lines[1]).unwrap();
        let operate_captures = re_operation.captures(&lines[2]).unwrap();
        let test_captures = re_test.captures(&lines[3]).unwrap();
        let true_captures = re_true.captures(&lines[4]).unwrap();
        let false_captures = re_false.captures(&lines[5]).unwrap();

        // Data extractions
        let starting_items: Vec<i64> = starting_captures[1].split(",").map(|x| x.trim().parse().unwrap()).collect();
        let operation = if &operate_captures[2] == "old" {
            Operation::Square
        } else {
            match &operate_captures[1] {
                "+" => Operation::Add     (operate_captures[2].parse().unwrap()),
                "*" => Operation::Multiply(operate_captures[2].parse().unwrap()),
                other => panic!("Error: Unknown operation {other}"),
            }
        };

        let test_value: i64 = test_captures[1].parse().unwrap();
        let target_true  = true_captures[1].parse().unwrap();
        let target_false = false_captures[1].parse().unwrap();

        // Construct object
        Self {
            items: starting_items,
            operation: operation,
            test_value: test_value,
            target_true: target_true,
            target_false: target_false,
            inspections: 0,
        }
    }
}

fn solve(input: &str, part2: bool) -> io::Result<()> {
    let input_str = std::fs::read_to_string(input).unwrap();
    let input_str = input_str.trim();
    let input: Vec<_> = input_str.split("\n\n").collect();

    // Initialize
    let mut monkeys: Vec<_> = input.iter().map(|x| Monkey::from(*x)).collect();
    let num_monkeys = &monkeys.len();
    let lcm = &monkeys.iter().map(|x| x.test_value).product::<i64>();

    // Turns
    let lcm = &monkeys.iter().map(|x| x.test_value).product::<i64>();
    let rounds = if part2 { 10_000 } else { 20 };
    for round in 0..rounds {
        for m_id in 0..*num_monkeys {
            for item_id in 0..monkeys[m_id].items.len() {
                monkeys[m_id].inspections += 1;
                match monkeys[m_id].operation {
                    Operation::Add(v) => monkeys[m_id].items[item_id] += v,
                    Operation::Multiply(v) => monkeys[m_id].items[item_id] *= v,
                    Operation::Square => monkeys[m_id].items[item_id] *= monkeys[m_id].items[item_id],
                }
                if !part2 { monkeys[m_id].items[item_id] /= 3; }
                let target_monkey = if monkeys[m_id].items[item_id] % monkeys[m_id].test_value == 0 {
                    monkeys[m_id].target_true
                } else {
                    monkeys[m_id].target_false
                };
                if part2 { monkeys[m_id].items[item_id] %= lcm; }
                let item_transferred = monkeys[m_id].items[item_id];
                monkeys[target_monkey as usize].items.push(item_transferred);
            }
            monkeys[m_id].items.clear();
        }
    }

    // Monkey Business
    monkeys.sort_by(|a, b| b.inspections.cmp(&a.inspections));
    let top_two: i64 = monkeys.iter().take(2).map(|x| x.inspections).product();
    let part = if part2 { 2 } else { 1 };
    println!("Part {part}: {top_two}");
    // Part 1: 120384
    // Part 2: 32059801242

    Ok(())
}

fn main() {
    let args: Vec<String> = env::args().collect();
    let filename = &args[1];
    solve(&filename,false).unwrap();
    solve(&filename,true).unwrap();
}