//! ZX-diagram: nodes (Z/X spiders, H-boxes) and wires.
//!
//! The ZX-calculus is a graphical language for quantum computing.
//! Spiders are the fundamental building blocks:
//! - Z-spider (green): |0⟩^⊗n + e^{iα}|1⟩^⊗n  
//! - X-spider (red):   |+⟩^⊗n + e^{iα}|−⟩^⊗n
//! - H-box: Hadamard gate

use std::collections::HashMap;
use std::f64::consts::PI;

use serde::{Deserialize, Serialize};

/// Identifier for a node in the ZX diagram.
pub type NodeId = usize;
/// Identifier for a wire (edge) in the ZX diagram.
pub type WireId = usize;

/// Node type in the ZX-calculus.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum ZXNode {
    /// Z-spider (green dot) with phase α (radians).
    Z { phase: f64 },
    /// X-spider (red dot) with phase α (radians).
    X { phase: f64 },
    /// Hadamard box (yellow square).
    H,
    /// Input boundary (numbered wire from left).
    Input(usize),
    /// Output boundary (numbered wire from right).
    Output(usize),
}

impl ZXNode {
    /// Phase of the spider, 0 for H/boundary nodes.
    #[must_use]
    pub fn phase(&self) -> f64 {
        match self {
            Self::Z { phase } | Self::X { phase } => *phase,
            _ => 0.0,
        }
    }

    /// True if this is a Z-spider.
    #[must_use]
    pub fn is_z(&self) -> bool {
        matches!(self, Self::Z { .. })
    }
    /// True if this is an X-spider.
    #[must_use]
    pub fn is_x(&self) -> bool {
        matches!(self, Self::X { .. })
    }
    /// True if this is an H-box.
    #[must_use]
    pub fn is_h(&self) -> bool {
        matches!(self, Self::H)
    }
    /// True if this is a boundary (input or output).
    #[must_use]
    pub fn is_boundary(&self) -> bool {
        matches!(self, Self::Input(_) | Self::Output(_))
    }

    /// True if this is a zero-phase spider (identity-like).
    #[must_use]
    pub fn is_zero_phase(&self) -> bool {
        match self {
            Self::Z { phase } | Self::X { phase } => {
                phase.abs() < 1e-12 || (phase - 2.0 * PI).abs() < 1e-12
            }
            _ => false,
        }
    }

    /// Same colour as another spider?
    #[must_use]
    pub fn same_colour(&self, other: &Self) -> bool {
        matches!(
            (self, other),
            (Self::Z { .. }, Self::Z { .. }) | (Self::X { .. }, Self::X { .. })
        )
    }
}

impl std::fmt::Display for ZXNode {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Z { phase } => write!(f, "Z({phase:.3})"),
            Self::X { phase } => write!(f, "X({phase:.3})"),
            Self::H => write!(f, "H"),
            Self::Input(i) => write!(f, "In({i})"),
            Self::Output(i) => write!(f, "Out({i})"),
        }
    }
}

/// Wire type: regular (Clifford) or Hadamard-decorated.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum WireType {
    /// Standard wire.
    Regular,
    /// Hadamard wire (equivalent to inserting an H-box).
    Hadamard,
}

/// A ZX-diagram: undirected hypergraph of spiders and wires.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ZXDiagram {
    /// Node storage: node_id → ZXNode.
    pub nodes: HashMap<NodeId, ZXNode>,
    /// Wire storage: wire_id → (node_a, node_b, wire_type).
    pub wires: HashMap<WireId, (NodeId, NodeId, WireType)>,
    /// Ordered input node ids.
    pub inputs: Vec<NodeId>,
    /// Ordered output node ids.
    pub outputs: Vec<NodeId>,
    next_node: NodeId,
    next_wire: WireId,
}

impl ZXDiagram {
    // ── Construction ──────────────────────────────────────────────

    /// Empty diagram.
    #[must_use]
    pub fn new() -> Self {
        Self {
            nodes: HashMap::new(),
            wires: HashMap::new(),
            inputs: Vec::new(),
            outputs: Vec::new(),
            next_node: 0,
            next_wire: 0,
        }
    }

    /// Add a node, returning its id.
    pub fn add_node(&mut self, node: ZXNode) -> NodeId {
        let id = self.next_node;
        self.next_node += 1;
        self.nodes.insert(id, node);
        id
    }

    /// Add an input boundary, returning its node id.
    pub fn add_input(&mut self) -> NodeId {
        let idx = self.inputs.len();
        let id = self.add_node(ZXNode::Input(idx));
        self.inputs.push(id);
        id
    }

    /// Add an output boundary, returning its node id.
    pub fn add_output(&mut self) -> NodeId {
        let idx = self.outputs.len();
        let id = self.add_node(ZXNode::Output(idx));
        self.outputs.push(id);
        id
    }

    /// Add a wire between two nodes.
    pub fn add_wire(&mut self, a: NodeId, b: NodeId) -> WireId {
        let id = self.next_wire;
        self.next_wire += 1;
        self.wires.insert(id, (a, b, WireType::Regular));
        id
    }

    /// Add a Hadamard-decorated wire.
    pub fn add_h_wire(&mut self, a: NodeId, b: NodeId) -> WireId {
        let id = self.next_wire;
        self.next_wire += 1;
        self.wires.insert(id, (a, b, WireType::Hadamard));
        id
    }

    // ── Accessors ─────────────────────────────────────────────────

    /// Neighbours of a node (connected node ids).
    #[must_use]
    pub fn neighbours(&self, node: NodeId) -> Vec<(NodeId, WireId, WireType)> {
        self.wires
            .iter()
            .filter_map(|(&wid, &(a, b, wt))| {
                if a == node {
                    Some((b, wid, wt))
                } else if b == node {
                    Some((a, wid, wt))
                } else {
                    None
                }
            })
            .collect()
    }

    /// Degree of a node.
    #[must_use]
    pub fn degree(&self, node: NodeId) -> usize {
        self.neighbours(node).len()
    }

    /// Number of nodes.
    #[must_use]
    pub fn node_count(&self) -> usize {
        self.nodes.len()
    }

    /// Number of wires.
    #[must_use]
    pub fn wire_count(&self) -> usize {
        self.wires.len()
    }

    /// Number of input/output qubits.
    #[must_use]
    pub fn qubit_count(&self) -> usize {
        self.inputs.len()
    }

    // ── Diagram mutation ──────────────────────────────────────────

    /// Remove a node and all its wires.
    pub fn remove_node(&mut self, id: NodeId) {
        self.nodes.remove(&id);
        self.wires.retain(|_, &mut (a, b, _)| a != id && b != id);
        self.inputs.retain(|&i| i != id);
        self.outputs.retain(|&o| o != id);
    }

    /// Remove a wire by id.
    pub fn remove_wire(&mut self, wid: WireId) {
        self.wires.remove(&wid);
    }

    /// Update the phase of a spider node.
    pub fn set_phase(&mut self, id: NodeId, phase: f64) {
        if let Some(ZXNode::Z { phase: p } | ZXNode::X { phase: p }) = self.nodes.get_mut(&id) {
            *p = phase;
        }
    }

    // ── Rewrite helpers ───────────────────────────────────────────

    /// True if nodes `a` and `b` are directly connected.
    #[must_use]
    pub fn connected(&self, a: NodeId, b: NodeId) -> bool {
        self.wires
            .values()
            .any(|&(x, y, _)| (x == a && y == b) || (x == b && y == a))
    }

    /// Find the wire connecting `a` and `b`, if any.
    #[must_use]
    pub fn wire_between(&self, a: NodeId, b: NodeId) -> Option<WireId> {
        self.wires.iter().find_map(|(&wid, &(x, y, _))| {
            if (x == a && y == b) || (x == b && y == a) {
                Some(wid)
            } else {
                None
            }
        })
    }

    /// Merge two same-colour adjacent spiders (spider fusion rule).
    /// Returns the id of the merged spider (reuses `a`), removes `b`.
    pub fn fuse_spiders(&mut self, a: NodeId, b: NodeId) -> Option<NodeId> {
        let node_a = self.nodes.get(&a)?.clone();
        let node_b = self.nodes.get(&b)?.clone();
        if !node_a.same_colour(&node_b) {
            return None;
        }

        // New phase = sum of phases (mod 2π)
        let new_phase = (node_a.phase() + node_b.phase()).rem_euclid(2.0 * PI);
        self.set_phase(a, new_phase);

        // Find wire between a and b and remove it
        if let Some(wid) = self.wire_between(a, b) {
            self.remove_wire(wid);
        }

        // Redirect all wires from b to a (except the one we removed)
        let b_wires: Vec<(WireId, NodeId, WireType)> = self
            .wires
            .iter()
            .filter_map(|(&wid, &(x, y, wt))| {
                if x == b {
                    Some((wid, y, wt))
                } else if y == b {
                    Some((wid, x, wt))
                } else {
                    None
                }
            })
            .collect();

        for (wid, other, wt) in b_wires {
            self.wires.remove(&wid);
            if other != a {
                self.wires.insert(wid, (a, other, wt));
            }
        }

        self.remove_node(b);
        Some(a)
    }
}

impl Default for ZXDiagram {
    fn default() -> Self {
        Self::new()
    }
}

/// Build the ZX representation of a Hadamard gate (1 qubit).
/// H = Z(0) —H— Z(0) with Hadamard wire.
#[must_use]
pub fn hadamard_diagram() -> ZXDiagram {
    let mut d = ZXDiagram::new();
    let inp = d.add_input();
    let z = d.add_node(ZXNode::Z { phase: 0.0 });
    let out = d.add_output();
    d.add_wire(inp, z);
    d.add_h_wire(z, out);
    d
}

/// Build the ZX representation of a CNOT gate (2 qubits).
#[must_use]
pub fn cnot_diagram() -> ZXDiagram {
    let mut d = ZXDiagram::new();
    let in0 = d.add_input();
    let in1 = d.add_input();
    let z = d.add_node(ZXNode::Z { phase: 0.0 }); // control
    let x = d.add_node(ZXNode::X { phase: 0.0 }); // target
    let out0 = d.add_output();
    let out1 = d.add_output();
    d.add_wire(in0, z);
    d.add_wire(z, out0);
    d.add_wire(in1, x);
    d.add_wire(x, out1);
    d.add_wire(z, x); // entangling wire
    d
}

/// Build the ZX representation of a T gate.
/// T = Z(π/4)
#[must_use]
pub fn t_gate_diagram() -> ZXDiagram {
    let mut d = ZXDiagram::new();
    let inp = d.add_input();
    let z = d.add_node(ZXNode::Z { phase: PI / 4.0 });
    let out = d.add_output();
    d.add_wire(inp, z);
    d.add_wire(z, out);
    d
}

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

    #[test]
    fn cnot_zx_representation() {
        let d = cnot_diagram();
        assert_eq!(d.qubit_count(), 2);
        assert_eq!(d.inputs.len(), 2);
        assert_eq!(d.outputs.len(), 2);
        // Z (control) and X (target) nodes plus 4 boundary nodes = 6
        assert_eq!(d.node_count(), 6);
    }

    #[test]
    fn hadamard_zx_representation() {
        let d = hadamard_diagram();
        assert_eq!(d.qubit_count(), 1);
        assert_eq!(d.wire_count(), 2); // regular + hadamard wire
    }

    #[test]
    fn t_gate_zx_representation() {
        let d = t_gate_diagram();
        // Should have a Z(π/4) node
        let has_t = d
            .nodes
            .values()
            .any(|n| matches!(n, ZXNode::Z { phase } if (phase - PI/4.0).abs() < 1e-9));
        assert!(has_t, "T gate should contain Z(π/4) node");
    }

    #[test]
    fn spider_fusion_z() {
        let mut d = ZXDiagram::new();
        let a = d.add_node(ZXNode::Z { phase: PI / 4.0 });
        let b = d.add_node(ZXNode::Z { phase: PI / 4.0 });
        d.add_wire(a, b);
        let merged = d.fuse_spiders(a, b);
        assert!(merged.is_some());
        // Fused phase = π/4 + π/4 = π/2
        let new_phase = d.nodes[&a].phase();
        assert!(
            (new_phase - PI / 2.0).abs() < 1e-9,
            "fused phase = {new_phase}"
        );
    }

    #[test]
    fn spider_fusion_x() {
        let mut d = ZXDiagram::new();
        let a = d.add_node(ZXNode::X { phase: PI / 2.0 });
        let b = d.add_node(ZXNode::X { phase: PI / 2.0 });
        d.add_wire(a, b);
        d.fuse_spiders(a, b);
        let new_phase = d.nodes[&a].phase();
        assert!((new_phase - PI).abs() < 1e-9);
    }

    #[test]
    fn spider_fusion_different_colour_fails() {
        let mut d = ZXDiagram::new();
        let a = d.add_node(ZXNode::Z { phase: 0.0 });
        let b = d.add_node(ZXNode::X { phase: 0.0 });
        d.add_wire(a, b);
        assert!(d.fuse_spiders(a, b).is_none());
    }

    #[test]
    fn identity_removal_candidate() {
        // A zero-phase Z-spider with exactly 2 legs is an identity wire
        let mut d = ZXDiagram::new();
        let inp = d.add_input();
        let z = d.add_node(ZXNode::Z { phase: 0.0 });
        let out = d.add_output();
        d.add_wire(inp, z);
        d.add_wire(z, out);
        assert_eq!(d.degree(z), 2);
        assert!(d.nodes[&z].is_zero_phase());
    }

    #[test]
    fn node_degree() {
        let d = cnot_diagram();
        // The Z (control) node connects to: in0, out0, X = degree 3
        let z_id = d
            .nodes
            .iter()
            .find_map(|(&id, n)| if n.is_z() { Some(id) } else { None })
            .unwrap();
        assert_eq!(d.degree(z_id), 3);
    }
}