use microlp::{ComparisonOp, OptimizationDirection, Problem};
use std::collections::VecDeque;

fn main() {
    let input = include_str!("../../input/10");
    println!("Part 1: {}", part1(input));
    println!("Part 2: {}", part2(input));
}

#[test]
fn example() {
    let input = include_str!("../../input/10-test");
    assert_eq!(part1(input), 7);
    assert_eq!(part2(input), 33);
}

fn part1(input: &str) -> usize {
    parse_input(input).iter().map(minimum_light_presses).sum()
}

fn part2(input: &str) -> usize {
    parse_input(input).iter().map(minimum_joltage_presses).sum()
}

#[derive(Default, Debug)]
struct Machine {
    required_lights: Vec<bool>,
    buttons: Vec<Button>,
    required_joltages: Vec<usize>,
}

impl Machine {
    fn parse(input: &str) -> Self {
        let mut m = Machine::default();
        for part in input.split_whitespace() {
            if part.starts_with('[') {
                m.required_lights = parse_lights(part);
            } else if part.starts_with('(') {
                m.buttons.push(parse_button(part));
            } else if part.starts_with('{') {
                m.required_joltages = parse_joltages(part);
            } else {
                panic!("Unknown part '{}'", part);
            }
        }
        m
    }
}

fn minimum_light_presses(machine: &Machine) -> usize {
    let mut work = VecDeque::new();
    work.push_back((vec![false; machine.required_lights.len()], 0usize));

    while let Some((lights, presses)) = work.pop_front() {
        let matches = lights.iter().eq(&machine.required_lights);
        if matches {
            return presses;
        }

        let new_lights = machine
            .buttons
            .iter()
            .map(|b| (press_light(&lights, b), presses + 1));
        work.extend(new_lights);
    }

    unreachable!()
}

fn minimum_joltage_presses(machine: &Machine) -> usize {
    let mut problem = Problem::new(OptimizationDirection::Minimize);
    let mut buttons_vars = Vec::new();
    for _ in 0..machine.buttons.len() {
        let var = problem.add_integer_var(1.0, (0, i32::MAX));
        buttons_vars.push(var);
    }

    for (i, joltage) in machine.required_joltages.iter().copied().enumerate() {
        let vars = buttons_vars
            .iter()
            .copied()
            .zip(&machine.buttons)
            .filter(|(_var, button)| button.contains(&i))
            .map(|(var, _button)| (var, 1.0));
        problem.add_constraint(vars, ComparisonOp::Eq, joltage as f64);
    }

    let solution = problem.solve().unwrap();
    buttons_vars
        .into_iter()
        .map(|var| solution[var].round() as usize)
        .sum()
}

fn press_light(lights: &[bool], button: &Button) -> Vec<bool> {
    let mut lights = lights.to_vec();
    for &b in button {
        lights[b] ^= true;
    }
    lights
}

fn parse_lights(input: &str) -> Vec<bool> {
    input
        .strip_prefix('[')
        .unwrap()
        .strip_suffix(']')
        .unwrap()
        .chars()
        .map(|c| match c {
            '.' => false,
            '#' => true,
            other => panic!("unknown light char '{other}'"),
        })
        .collect()
}

type Button = Vec<usize>;

fn parse_button(input: &str) -> Button {
    input
        .strip_prefix('(')
        .unwrap()
        .strip_suffix(')')
        .unwrap()
        .split(',')
        .map(|c| c.parse().unwrap())
        .collect()
}

fn parse_joltages(input: &str) -> Vec<usize> {
    input
        .strip_prefix('{')
        .unwrap()
        .strip_suffix('}')
        .unwrap()
        .split(',')
        .map(|c| c.parse().unwrap())
        .collect()
}

fn parse_input(input: &str) -> Vec<Machine> {
    input.lines().map(Machine::parse).collect()
}