enki 0.0.1

use crate::element::Element;
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};

/// Weisfeiler–Lehman hashing parameters
const WL_INIT_MUL: u64 = 0x9E3779B1; // golden ratio constant (Knuth)
const WL_STEP_MUL: u64 = 0x1000193; // FNV prime
const WL_ROUNDS: usize = 3; // number of refinement rounds

#[derive(Debug, Clone)]
pub struct Node {
    pub atom: Element,
}

#[derive(Debug, Clone)]
pub struct Edge {
    pub destination: usize,
}

#[derive(Debug, Clone)]
pub struct Molecule {
    pub nodes: Vec<Node>,
    pub edges: Vec<Vec<Edge>>,
}

impl Molecule {
    pub fn new() -> Self {
        Self {
            nodes: Vec::new(),
            edges: Vec::new(),
        }
    }

    // Helper to build a molecule with no edges (single atom)
    #[allow(dead_code)]
    fn atom(a: Element) -> Molecule {
        Molecule {
            nodes: vec![Node { atom: a }],
            edges: vec![vec![]],
        }
    }

    pub fn add_node(&mut self, atom: Element) -> usize {
        let id = self.nodes.len();
        self.nodes.push(Node { atom });
        self.edges.push(Vec::new());
        id
    }

    pub fn degree(&self, node_id: usize) -> u8 {
        self.edges[node_id].len() as u8
    }

    pub fn open_ports(&self, node_id: usize) -> u8 {
        self.nodes[node_id].atom.max_bonds() - self.degree(node_id)
    }

    pub fn add_bond(&mut self, a: usize, b: usize) -> Result<(), String> {
        if a == b {
            return Err("Cannot bond atom to itself".into());
        }

        if self.open_ports(a) == 0 {
            return Err(format!(
                "Atom {} ({:?}) has no free ports",
                a, self.nodes[a].atom
            ));
        }

        if self.open_ports(b) == 0 {
            return Err(format!(
                "Atom {} ({:?}) has no free ports",
                b, self.nodes[b].atom
            ));
        }

        // prevent duplicate bonds
        if self.edges[a].iter().any(|e| e.destination == b) {
            return Err(format!("Bond between {} and {} already exists", a, b));
        }

        // undirected graph: add edge in both directions
        self.edges[a].push(Edge { destination: b });
        self.edges[b].push(Edge { destination: a });

        Ok(())
    }

    /// Weisfeiler–Lehman (WL) hashing approach
    /// Compute an isomorphism-invariant hash for the molecule.
    /// Isomorphic molecules produce identical hashes.
    pub fn wl_hash(&self) -> u64 {
        // --- Step 1: initial labels from atom type ---
        let mut labels: Vec<u64> = self
            .nodes
            .iter()
            .map(|n| {
                // simple stable mapping from enum discriminant
                let mut h = DefaultHasher::new();
                n.atom.hash(&mut h);
                h.finish()
            })
            .collect();

        // --- Step 2: iterative refinement (propagate neighborhood structure) ---
        for _ in 0..WL_ROUNDS {
            labels = (0..self.nodes.len())
                .map(|i| {
                    let mut neigh: Vec<u64> = self.edges[i]
                        .iter()
                        .map(|e| labels[e.destination])
                        .collect();
                    neigh.sort_unstable();

                    let mut h = labels[i].wrapping_mul(WL_INIT_MUL);
                    for v in neigh {
                        h = h.wrapping_mul(WL_STEP_MUL) ^ v;
                    }
                    h
                })
                .collect();
        }

        // --- Step 3: order-independent combination ---
        let mut sorted = labels;
        sorted.sort_unstable();
        let mut hasher = DefaultHasher::new();
        sorted.hash(&mut hasher);
        hasher.finish()
    }
}

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

    #[test]
    fn identical_atoms_hash_the_same() {
        let m1 = Molecule::atom(Element::I);
        let m2 = Molecule::atom(Element::I);
        assert_eq!(m1.wl_hash(), m2.wl_hash());
    }

    #[test]
    fn different_atoms_hash_differently() {
        let m1 = Molecule::atom(Element::S);
        let m2 = Molecule::atom(Element::K);
        assert_ne!(m1.wl_hash(), m2.wl_hash());
    }

    #[test]
    fn molecules_with_permuted_nodes_hash_the_same() {
        // Represents (K I)
        // Node 0 = A (application)
        // Node 1 = K
        // Node 2 = I
        let mol1 = Molecule {
            nodes: vec![
                Node { atom: Element::A },
                Node { atom: Element::K },
                Node { atom: Element::I },
            ],
            edges: vec![
                vec![Edge { destination: 1 }, Edge { destination: 2 }], // A -> K, I
                vec![Edge { destination: 0 }],                          // K -> A
                vec![Edge { destination: 0 }],                          // I -> A
            ],
        };

        // Same term but permuted node IDs
        let mol2 = Molecule {
            nodes: vec![
                Node { atom: Element::K },
                Node { atom: Element::I },
                Node { atom: Element::A },
            ],
            edges: vec![
                vec![Edge { destination: 2 }],                          // K -> A
                vec![Edge { destination: 2 }],                          // I -> A
                vec![Edge { destination: 0 }, Edge { destination: 1 }], // A -> K, I
            ],
        };

        assert_eq!(mol1.wl_hash(), mol2.wl_hash());
    }

    #[test]
    fn hash_non_trivial_molecule() {
        // Represents (S K) K  (a common SKI identity encoding)
        //
        //       A
        //      / \
        //     A   K
        //    / \
        //   S   K
        //
        // Ensure that:
        // A nontrivial SKI structure (not just single atoms) can be hashed.
        // The hashing algorithm runs without panic, even on a deeper application tree.
        // The result is stable and deterministic (if you run it twice, you get the same hash).
        // The hash isn’t degenerate (always zero).
        let mol = Molecule {
            nodes: vec![
                Node { atom: Element::A }, // 0
                Node { atom: Element::A }, // 1
                Node { atom: Element::S }, // 2
                Node { atom: Element::K }, // 3
                Node { atom: Element::K }, // 4
            ],
            edges: vec![
                vec![Edge { destination: 1 }, Edge { destination: 4 }], // A0 -> A1, K4
                vec![
                    Edge { destination: 0 },
                    Edge { destination: 2 },
                    Edge { destination: 3 },
                ], // A1 -> A0, S2, K3
                vec![Edge { destination: 1 }],                          // S -> A1
                vec![Edge { destination: 1 }],                          // K3 -> A1
                vec![Edge { destination: 0 }],                          // K4 -> A0
            ],
        };

        // Just check it hashes deterministically
        let h = mol.wl_hash();
        assert!(h != 0);
    }

    #[test]
    fn hash_empty_molecule() {
        let mol = Molecule {
            nodes: vec![],
            edges: vec![],
        };
        assert_eq!(mol.wl_hash(), mol.wl_hash()); // consistent
    }

    #[test]
    fn hash_invalid_molecule() {
        let mol = Molecule {
            nodes: vec![Node { atom: Element::A }],
            edges: vec![vec![]],
        };
        assert!(mol.wl_hash() != 0);
    }
}