use std::{cmp::max, collections::HashMap, ops::Index};

pub struct ComboNode {
    combo: Vec<u32>,
    weight: u32,
}

impl Index<usize> for ComboNode {
    type Output = u32;
    fn index(&self, index: usize) -> &Self::Output {
        self.combo.index(index)
    }
}

impl ComboNode {
    pub fn weight(&self) -> u32 {
        self.weight
    }

    pub fn new(combo: Vec<u32>) -> Self {
        let sum = combo.iter().sum();
        Self { combo, weight: sum }
    }

    pub fn get_neighbours(&self) -> impl Iterator<Item = ComboNode> {
        (0..self.combo.len()).map(|i| {
            let mut combo = self.combo.clone();
            combo[i] += 1;
            Self {
                combo,
                weight: self.weight + 1,
            }
        })
    }
}

pub struct ComboGraph {
    nodes: HashMap<u32, ComboNode>,
    values: Vec<u32>,
    target: u32,
    queue: Vec<ComboNode>,
}

impl ComboGraph {
    pub fn new(mut values: Vec<u32>, target: u32) -> Self {
        // ensure values is sorted
        values.sort_unstable();
        let num_values = values.len();
        Self {
            nodes: HashMap::new(),
            values,
            target,
            queue: vec![ComboNode::new(vec![0; num_values])],
        }
    }

    pub fn closest_node_target(&self) -> Option<&ComboNode> {
        let closest = self
            .nodes
            .iter()
            // Treat a zero difference as a difference of 1 to accomodate for rounding
            .map(|(pos, combo)| (max(self.target.abs_diff(*pos), 1), combo.weight(), pos))
            .min()?;

        self.nodes.get(closest.2)
    }

    /// Visit a combo and calculate its position and weight
    pub fn visit(&mut self, node: ComboNode) -> bool {
        let mut pos = 0;

        for i in 0..self.values.len() {
            pos += self.values[i] * node[i];
        }

        // Check if pos already found with less weight
        let current = self.nodes.get(&pos);
        if let Some(combo) = current {
            if combo.weight() > node.weight() {
                // should never happen
                panic!("Found better path second");
            }
            return false;
        } else if pos == self.target {
            self.nodes.insert(pos, node);
            return true;
        } else if pos < self.target {
            node.get_neighbours().for_each(|n| {
                self.queue.push(n);
            });
        }
        self.nodes.insert(pos, node);

        false
    }

    /// grows the graph to find the optimal path
    pub fn grow(&mut self) -> bool {
        let mut next = self.queue.pop();
        while let Some(node) = next {
            if self.visit(node) {
                return true;
            }
            next = self.queue.pop();
        }
        false
    }

    pub fn target_combo(&self) -> Option<Vec<u32>> {
        let combo = self.nodes.get(&self.target)?;
        Some(
            (0..self.values.len())
                .flat_map(|i| vec![self.values[i]; combo[i] as usize])
                .collect(),
        )
    }

    pub fn closest_combo(&self) -> Option<Vec<u32>> {
        let combo = self.closest_node_target()?;
        Some(
            (0..self.values.len())
                .flat_map(|i| vec![self.values[i]; combo[i] as usize])
                .collect(),
        )
    }
}

pub fn find_optimal_combo(strengths: Vec<u32>, target: u32, rounded: bool) -> Vec<u32> {
    let mut graph = ComboGraph::new(strengths, target);
    if graph.grow() && !rounded {
        graph.target_combo().expect("grow signalled target hit")
    } else {
        graph.closest_combo().expect("non empty graph")
    }
}