use std::{
    collections::{BinaryHeap, VecDeque},
    hash::Hash,
    iter::{repeat, repeat_n},
};

use aoc_runner_derive::{aoc, aoc_generator};
use indicatif::{ProgressBar, ProgressStyle};
use itertools::Itertools;
use regex::Regex;
use wide::{CmpGt, i16x16};

#[derive(Clone, Debug, Default)]
struct MachineDefinition {
    desired: Vec<bool>,
    buttons: Vec<Vec<usize>>,
    buttons2: Vec<i16x16>,
    buttons_max: i16x16,
    joltages: i16x16,
}

impl MachineDefinition {
    fn create<'a>(&'a self) -> Machine<'a> {
        Machine {
            d: self,
            lights: Vec::from_iter(repeat_n(false, self.desired.len())),
        }
    }

    fn create2<'a>(&'a self) -> JoltMachine<'a> {
        JoltMachine {
            d: self,
            joltages: i16x16::splat(0),
        }
    }
}

impl From<&str> for MachineDefinition {
    fn from(value: &str) -> Self {
        let parse_re = Regex::new(
            r##"\[(?<desired>[.#]+)\] (?<buttons>(\([0-9,]+\) ?)+) \{(?<joltages>[0-9,]+)\}"##,
        )
        .unwrap();

        let parts = parse_re.captures(value).unwrap();
        let joltages: [i16; 16] = parts["joltages"]
            .split(',')
            .map(|n| n.parse().unwrap())
            .chain(repeat(0))
            .take(16)
            .collect_array()
            .unwrap();

        let buttons = parts["buttons"]
            .split_ascii_whitespace()
            .map(|s| {
                s[1..s.len() - 1]
                    .split(',')
                    .map(|n| n.parse().unwrap())
                    .collect()
            })
            .sorted_unstable_by_key(|s: &Vec<usize>| s.len())
            .rev()
            .collect_vec();

        let mut buttons2 = Vec::new();
        let mut buttons_max = [0i16; 16];

        for (i, b) in buttons.iter().enumerate() {
            let mut but = [0i16; 16];
            for i in b {
                but[*i] = 1;
            }
            buttons2.push(i16x16::new(but));

            // find the joltage this button affects with the lowest value
            // it is the max number of presses for this button
            buttons_max[i] = b.iter().map(|idx| joltages[*idx]).min().unwrap();
        }

        MachineDefinition {
            desired: parts["desired"].chars().map(|c| c == '#').collect(),
            buttons: buttons,
            buttons2: buttons2,
            buttons_max: i16x16::new(buttons_max),
            joltages: i16x16::new(joltages),
        }
    }
}

#[derive(Clone, Debug)]
struct Machine<'a> {
    d: &'a MachineDefinition,
    lights: Vec<bool>,
}

#[derive(Clone, Debug)]
struct JoltMachine<'a> {
    d: &'a MachineDefinition,
    joltages: i16x16,
}

impl<'a> Machine<'a> {
    /// Get the state after pressing `button`, returns None if the state is as desired.
    fn press(&self, button: usize) -> Option<Self> {
        let mut new_state = self.lights.clone();
        for light in &self.d.buttons[button] {
            new_state[*light] = !new_state[*light];
        }
        if new_state == self.d.desired {
            None
        } else {
            Some(Machine {
                d: self.d,
                lights: new_state,
            })
        }
    }
    /// Get the possible states from the current position
    fn next_states(&self) -> Vec<(usize, Option<Self>)> {
        self.d
            .buttons
            .iter()
            .enumerate()
            .map(|(i, _but)| (i, self.press(i)))
            .collect()
    }
}

impl<'a> JoltMachine<'a> {
    fn press_jolts(&self, button: usize, presses: &i16x16) -> (i16x16, Option<Self>) {
        // let mut new_joltage = self.joltages.clone();
        // // for jolt in &self.d.buttons[button] {
        // //     new_joltage[*jolt] += 1;
        // // }
        let new_joltage = self.joltages + self.d.buttons2[button];
        let mut new_presses = presses.clone();
        new_presses.as_mut_array()[button] += 1;
        if new_joltage == self.d.joltages {
            (new_presses, None)
        } else {
            (
                new_presses,
                Some(Self {
                    d: self.d,
                    joltages: new_joltage,
                }),
            )
        }
    }

    fn next_states_jolt(&self, presses: &i16x16) -> Vec<(i16x16, Option<Self>)> {
        self.d
            .buttons
            .iter()
            .enumerate()
            .map(|(i, _but)| self.press_jolts(i, &presses))
            // .inspect(|(p, o)| println!(" {p:?} {o:?}\n"))
            // joltages monotonically increase, so cull any where a joltage is higher than needed
            .filter(|(presses, candidate)| {
                !presses.simd_gt(self.d.buttons_max).any()
                    && candidate.as_ref().is_none_or(|c| {
                        !c.joltages.simd_gt(self.d.joltages).any()
                        // !c.joltages
                        //     .iter()
                        //     .zip(self.d.joltages.iter())
                        //     .any(|(candidate, expected)| candidate > expected)
                    })
            })
            .collect()
    }
}

#[derive(Debug, Clone)]
struct PressSet<'a> {
    machine: Machine<'a>,
    presses: usize,
}

#[derive(Debug, Clone)]
struct PressSet2<'a> {
    machine: JoltMachine<'a>,
    presses: i16x16,
}

// NOTE: All compares are reversed so our max heap becomes a min heap
impl<'a> Eq for PressSet<'a> {}
impl<'a> Eq for PressSet2<'a> {}

impl<'a> PartialEq for PressSet<'a> {
    fn eq(&self, other: &Self) -> bool {
        other.presses.eq(&self.presses)
    }
}

impl<'a> PartialEq for PressSet2<'a> {
    fn eq(&self, other: &Self) -> bool {
        other.presses.eq(&self.presses)
    }
}

impl<'a> PartialOrd for PressSet<'a> {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

impl<'a> PartialOrd for PressSet2<'a> {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

impl<'a> Ord for PressSet<'a> {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        other.presses.cmp(&self.presses)
    }
}

impl<'a> Ord for PressSet2<'a> {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        other.presses.reduce_add().cmp(&self.presses.reduce_add())
    }
}

fn find_best(md: &MachineDefinition) -> usize {
    let m = md.create();
    let mut to_check = BinaryHeap::new();
    to_check.push(PressSet {
        presses: 0,
        machine: m,
    });

    while let Some(candidate) = to_check.pop() {
        let cm = candidate.machine.clone();
        for next in cm.next_states() {
            let presses = candidate.presses + 1;
            if let Some(new_m) = next.1 {
                to_check.push(PressSet {
                    presses,
                    machine: new_m.clone(),
                })
            } else {
                return presses;
            }
        }
    }
    panic!()
}

fn find_best_jolts(md: &MachineDefinition) -> usize {
    let m = md.create2();
    let mut to_check = VecDeque::new();
    to_check.push_back(PressSet2 {
        presses: i16x16::splat(0),
        machine: m,
    });

    let mut pb = ProgressBar::no_length()
        .with_style(
            ProgressStyle::with_template(
                "[{elapsed_precise}/{eta_precise}] {bar:40.cyan/blue} {pos:>7}/{len:7} {per_sec}",
            )
            .unwrap(),
        )
        .with_finish(indicatif::ProgressFinish::AndLeave);

    while let Some(candidate) = to_check.pop_front() {
        pb.inc(1);
        pb.set_length(to_check.len() as u64);
        let cm = candidate.machine.clone();
        for (presses, next) in cm.next_states_jolt(&candidate.presses) {
            if let Some(new_m) = next {
                to_check.push_back(PressSet2 {
                    presses,
                    machine: new_m.clone(),
                })
            } else {
                return presses.reduce_add() as usize;
            }
        }
    }
    panic!()
}

#[aoc_generator(day10)]
fn parse(input: &str) -> Vec<MachineDefinition> {
    input.lines().map(|l| l.into()).collect()
}

#[aoc(day10, part1)]
fn part1(input: &[MachineDefinition]) -> u64 {
    input
        .iter()
        .map(find_best)
        // .inspect(|sol| println!("  [{sol:?}]"))
        .map(|sol| sol as u64)
        .sum()
}

#[aoc(day10, part2)]
fn part2(input: &[MachineDefinition]) -> u64 {
    input
        .iter()
        .map(find_best_jolts)
        .map(|sol| sol as u64)
        .sum()
}

#[cfg(test)]
mod tests {
    use rstest::rstest;

    use super::*;

    const EXAMPLE: &str = "[.##.] (3) (1,3) (2) (2,3) (0,2) (0,1) {3,5,4,7}
[...#.] (0,2,3,4) (2,3) (0,4) (0,1,2) (1,2,3,4) {7,5,12,7,2}
[.###.#] (0,1,2,3,4) (0,3,4) (0,1,2,4,5) (1,2) {10,11,11,5,10,5}";

    #[rstest]
    #[case(EXAMPLE, 7)]
    fn part1_example(#[case] input: &str, #[case] expected: u64) {
        assert_eq!(part1(&parse(input)), expected);
    }

    #[rstest]
    #[case(EXAMPLE, 33)]
    fn part2_example(#[case] input: &str, #[case] expected: u64) {
        assert_eq!(part2(&parse(input)), expected);
    }
}