use crate::coordinates::Coordinates;
use crate::qube::{Dimension, NodeIdx, Qube};
use std::collections::BTreeMap;
use std::collections::HashMap;
use std::sync::atomic::Ordering;
use tiny_vec::TinyVec;

impl Qube {
    /// Creates a hash map where the keys are structural hashes of child nodes
    /// and the values are vectors of node indices that share the same hash.
    fn children_hash_map(
        &mut self,
        children: &BTreeMap<Dimension, TinyVec<NodeIdx, 4>>,
    ) -> HashMap<u64, Vec<NodeIdx>> {
        let mut map: HashMap<u64, Vec<NodeIdx>> = HashMap::new();

        for (_dim, kids) in children.iter() {
            for &c in kids.iter() {
                let h = self.compute_structural_hash(c);
                map.entry(h).or_default().push(c);
            }
        }
        map
    }

    /// Checks if a node is a leaf node (i.e., it has no children).
    fn is_leaf(&self, id: NodeIdx) -> bool {
        self.node_ref(id).expect("valid node").children().is_empty()
    }

    /// Recursively prunes empty nodes from the tree.
    fn prune_empty_nodes_recursively(&mut self, node_id: NodeIdx) {
        let children: Vec<NodeIdx> = {
            let node = self.node_ref(node_id).unwrap();
            node.children().values().flat_map(|v| v.iter().copied()).collect()
        };

        for child in &children {
            self.prune_empty_nodes_recursively(*child);
        }

        let keep: std::collections::HashSet<NodeIdx> = children
            .into_iter()
            .filter(|&child| !matches!(self.node_ref(child).unwrap().coords(), Coordinates::Empty))
            .collect();

        let parent = self.node_mut(node_id).unwrap();
        for kids in parent.children_mut().values_mut() {
            kids.retain(|id| keep.contains(id));
        }
    }

    /// Invalidates the cached structural hash of a node.
    fn invalidate_structural_hash(&mut self, id: NodeIdx) {
        let node = self.node_mut(id).unwrap();
        node.structural_hash().store(0, Ordering::Release);
    }

    /// Deduplicates the children of a node by merging nodes with identical structural hashes.
    fn dedup_children_locally(&mut self, parent: NodeIdx) {
        let snapshot = {
            let node = self.node_ref(parent).unwrap();
            node.children().clone()
        };

        for (dim, kids) in snapshot {
            let mut seen: HashMap<u64, NodeIdx> = HashMap::new();
            let mut unique: Vec<NodeIdx> = Vec::new();

            for &child in &kids {
                let h = self.compute_structural_hash(child);

                if seen.insert(h, child).is_none() {
                    unique.push(child);
                }
            }

            let parent_node = self.node_mut(parent).unwrap();
            parent_node.children_mut().insert(dim, unique.into());
        }

        self.invalidate_structural_hash(parent);
    }

    /// Recursively deduplicates nodes in the tree, starting from the given node.
    fn dedup_recursively(&mut self, node_id: NodeIdx) {
        let children: Vec<NodeIdx> = {
            let node = self.node_ref(node_id).unwrap();
            node.children().values().flat_map(|v| v.iter().copied()).collect()
        };

        for child in children {
            self.dedup_recursively(child);
        }

        self.dedup_children_locally(node_id);
    }

    /// Merges two subtrees by merging their coordinates and children.
    #[allow(dead_code)]
    fn merge_subtrees(&mut self, target_id: NodeIdx, source_id: NodeIdx) {
        // Merge the coordinates of the source node into the target node
        {
            let mut target_coords = self.node_ref(target_id).unwrap().coords().clone();
            let source_coords = self.node_ref(source_id).unwrap().coords().clone();

            let merged_coords = target_coords.merge_coords(&source_coords);
            let target_node = self.node_mut(target_id).unwrap();
            *target_node.coords_mut() = merged_coords;
        }

        // Recursively merge the children of the source node into the target node
        let source_children = {
            let source_node = self.node_ref(source_id).unwrap();
            source_node.children().clone()
        };

        for (dim, source_child_ids) in source_children {
            for source_child_id in source_child_ids {
                let target_children = {
                    let target_node = self.node_ref(target_id).unwrap();
                    target_node.children().get(&dim).cloned().unwrap_or_default()
                };

                let mut merged_children = target_children.clone();
                merged_children.push(source_child_id);

                let target_node = self.node_mut(target_id).unwrap();
                target_node.children_mut().insert(dim, merged_children.into());
            }
        }

        // Invalidate the structural hash of the target node
        self.invalidate_structural_hash(target_id);
    }

    /// Compresses the tree by merging nodes, pruning empty nodes, and deduplicating nodes.
    pub fn compress(&mut self) {
        // This method performs the following steps:
        // 1. Compresses nodes recursively.
        // 2. Prunes empty nodes from the tree.
        // 3. Deduplicates nodes that may have become identical after compression.

        let root = self.root();
        self.compress_recursively(root);
        self.prune_empty_nodes_recursively(root);
        self.dedup_recursively(root);
    }

    /// Recursively compresses the tree, merging coordinates of child nodes where possible.
    fn compress_recursively(&mut self, node_id: NodeIdx) {
        let children: Vec<NodeIdx> = {
            let node = self.node_ref(node_id).expect("Valid nodeIdx in tree");
            node.children().values().flat_map(|v| v.iter().copied()).collect()
        };

        if children.is_empty() {
            return;
        }

        let all_children_are_leaves = children.iter().all(|&id| self.is_leaf(id));

        if all_children_are_leaves {
            // group by dimension
            let mut by_dim: HashMap<Dimension, Vec<NodeIdx>> = HashMap::new();

            for &child in &children {
                let dim = *self.node_ref(child).unwrap().dim();
                by_dim.entry(dim).or_default().push(child);
            }

            for group in by_dim.values() {
                if group.len() > 1 {
                    self.merge_coords(group.to_vec());
                }
            }

            return;
        }

        for child in children {
            self.compress_recursively(child);
        }

        // children are fully compressed so we can hash & merge them
        let children_map = {
            let children = {
                let node = self.node_ref(node_id).expect("Valid nodeIdx in tree");
                node.children().clone()
            };
            self.children_hash_map(&children)
        };

        for group in children_map.values() {
            if group.len() <= 1 {
                continue; // nothing to merge
            }

            self.merge_coords(group.clone());
        }
    }

    /// Merges the coordinates of a group of nodes into the first node in the group.
    fn merge_coords(&mut self, group: Vec<NodeIdx>) {
        // The coordinates of all other nodes in the group are set to `Coordinates::Empty`.

        assert!(!group.is_empty());

        let mut merged: Coordinates = { self.node_ref(group[0]).unwrap().coords().clone() };

        for &id in group.iter().skip(1) {
            let coords = self.node_ref(id).unwrap().coords();
            merged.extend(coords);
        }

        {
            let node = self.node_mut(group[0]).unwrap();
            *node.coords_mut() = merged;
        }

        for &id in group.iter().skip(1) {
            let node = self.node_mut(id).unwrap();
            *node.coords_mut() = Coordinates::Empty;
        }
    }
}
