geographdb-core 0.5.4

//! Matrix Product State (MPS) backed by the 4D graph.
//!
//! An MPS represents a quantum state of N sites as a chain of rank-3 tensors:
//!
//!   A¹[χ_left, d, χ_right] ── A²[χ_left, d, χ_right] ── … ── Aᴺ[…]
//!
//! where `d` is the physical dimension (2 for qubits) and `χ` is the bond
//! dimension controlling how much entanglement the state can represent.
//!
//! Each site maps onto a `GraphNode4D`:
//!   - position x = site index, y = z = 0
//!   - `properties["shape"]` = [χ_left, d, χ_right]
//!   - `properties["data"]`  = flat f32 tensor in row-major order
//!   - outgoing `TemporalEdge` to site i+1 with `weight = χ_right` (bond dim)

use crate::algorithms::four_d::{GraphNode4D, GraphProperties, TemporalEdge};
use serde_json::Value;

// ── Tensor helpers ────────────────────────────────────────────────────────────

/// Store a rank-3 tensor on `node`.
///
/// `shape` is `[χ_left, d, χ_right]`. `data` must have length
/// `shape[0] * shape[1] * shape[2]`.
pub fn set_tensor(node: &mut GraphNode4D, shape: [usize; 3], data: Vec<f32>) {
    node.properties.insert("shape".into(), shape_to_json(shape));
    node.properties.insert("data".into(), data_to_json(&data));
}

/// Retrieve the rank-3 tensor stored on `node`.
///
/// Returns `None` if the node has no tensor data.
pub fn get_tensor(node: &GraphNode4D) -> Option<([usize; 3], Vec<f32>)> {
    let shape = json_to_shape(node.properties.get("shape")?)?;
    let data = json_to_data(node.properties.get("data")?)?;
    Some((shape, data))
}

// ── MPS construction ──────────────────────────────────────────────────────────

/// Build an MPS chain from a list of `(shape, data)` pairs.
///
/// Each pair describes one site tensor: `shape = [χ_left, d, χ_right]`.
/// Returns a `Vec<GraphNode4D>` wired as a left-to-right chain; node `i` has
/// an outgoing edge to node `i+1` with `weight = χ_right` of site `i`.
pub fn build_mps(tensors: &[(&[usize], &[f32])]) -> Vec<GraphNode4D> {
    let mut nodes: Vec<GraphNode4D> = tensors
        .iter()
        .enumerate()
        .map(|(i, (shape, data))| {
            let shape3 = [shape[0], shape[1], shape[2]];
            let mut node = GraphNode4D {
                id: i as u64,
                x: i as f32,
                y: 0.0,
                z: 0.0,
                begin_ts: 0,
                end_ts: 1,
                properties: GraphProperties::new(),
                successors: vec![],
            };
            set_tensor(&mut node, shape3, data.to_vec());
            node
        })
        .collect();

    // Wire edges: site i → site i+1, weight = χ_right of site i
    for i in 0..nodes.len().saturating_sub(1) {
        let chi_right = if let Some((shape, _)) = get_tensor(&nodes[i]) {
            shape[2] as f32
        } else {
            1.0
        };
        nodes[i].successors.push(TemporalEdge {
            dst: (i + 1) as u64,
            weight: chi_right,
            begin_ts: 0,
            end_ts: 1,
        });
    }

    nodes
}

// ── MPS operations ────────────────────────────────────────────────────────────

/// Compute ⟨ψ|ψ⟩ by contracting the MPS chain left-to-right.
///
/// Returns the squared norm of the state. For a normalised state this is 1.0.
pub fn mps_norm_sq(nodes: &[GraphNode4D]) -> f64 {
    if nodes.is_empty() {
        return 0.0;
    }

    // Left boundary: contract first tensor with its conjugate.
    // L[χ_r, χ_r'] = Σ_s A[1, s, χ_r] * A*[1, s, χ_r']
    let (shape0, data0) = match get_tensor(&nodes[0]) {
        Some(t) => t,
        None => return 0.0,
    };
    // chi_l0 is always 1 at the left boundary; index = s * chi_r0 + chi_r_idx
    let (_chi_l0, d0, chi_r0) = (shape0[0], shape0[1], shape0[2]);
    let mut l = vec![0.0f64; chi_r0 * chi_r0];
    for s in 0..d0 {
        for a in 0..chi_r0 {
            for b in 0..chi_r0 {
                let idx_a = s * chi_r0 + a;
                let idx_b = s * chi_r0 + b;
                l[a * chi_r0 + b] += data0[idx_a] as f64 * data0[idx_b] as f64;
            }
        }
    }

    // Sweep: contract L with each subsequent tensor.
    for node in &nodes[1..] {
        let (shape, data) = match get_tensor(node) {
            Some(t) => t,
            None => return 0.0,
        };
        let (chi_l, d, chi_r) = (shape[0], shape[1], shape[2]);
        let mut l_new = vec![0.0f64; chi_r * chi_r];
        for s in 0..d {
            for a_new in 0..chi_r {
                for b_new in 0..chi_r {
                    let mut val = 0.0f64;
                    for a_old in 0..chi_l {
                        for b_old in 0..chi_l {
                            let idx_a = a_old * d * chi_r + s * chi_r + a_new;
                            let idx_b = b_old * d * chi_r + s * chi_r + b_new;
                            val +=
                                l[a_old * chi_l + b_old] * data[idx_a] as f64 * data[idx_b] as f64;
                        }
                    }
                    l_new[a_new * chi_r + b_new] += val;
                }
            }
        }
        l = l_new;
    }

    // Final trace: L is [1×1] for a valid MPS.
    l[0]
}

/// Apply a single-site gate to site `site`.
///
/// `gate` is a `d×d` unitary in row-major order. The physical index of the
/// site tensor is updated: A'[χ_l, s', χ_r] = Σ_s gate[s', s] · A[χ_l, s, χ_r].
pub fn mps_apply_gate(nodes: &mut [GraphNode4D], site: usize, gate: &[f32]) {
    let (shape, data) = match get_tensor(&nodes[site]) {
        Some(t) => t,
        None => return,
    };
    let (chi_l, d, chi_r) = (shape[0], shape[1], shape[2]);

    // A'[χ_l, s', χ_r] = Σ_s gate[s', s] * A[χ_l, s, χ_r]
    let mut new_data = vec![0.0f32; chi_l * d * chi_r];
    for a in 0..chi_l {
        for s_out in 0..d {
            for b in 0..chi_r {
                let mut val = 0.0f32;
                for s_in in 0..d {
                    val += gate[s_out * d + s_in] * data[a * d * chi_r + s_in * chi_r + b];
                }
                new_data[a * d * chi_r + s_out * chi_r + b] = val;
            }
        }
    }

    set_tensor(&mut nodes[site], shape, new_data);
}

/// Apply a two-site gate to sites `site_a` and `site_b`.
///
/// The gate is a `d²×d²` unitary in row-major order indexed as
/// `gate[(s0'*d + s1') * d² + (s0*d + s1)]`.
///
/// **Restrictions:** both sites must have `χ_in = 1` (product state) and `d = 2`
/// (qubits). For general bond dimensions a full PEPS contraction engine is required.
///
/// Returns the actual bond dimension `χ_new` kept after SVD truncation (≤ `chi_max`).
/// Singular values below `1e-7` are treated as zero and not kept.
pub fn mps_apply_gate_2site(
    nodes: &mut [GraphNode4D],
    site_a: usize,
    site_b: usize,
    gate: &[f32],
    chi_max: usize,
) -> usize {
    let (shape_a, data_a) = match get_tensor(&nodes[site_a]) {
        Some(t) => t,
        None => return 0,
    };
    let (shape_b, data_b) = match get_tensor(&nodes[site_b]) {
        Some(t) => t,
        None => return 0,
    };
    let (chi_la, d, chi_m) = (shape_a[0], shape_a[1], shape_a[2]);
    let (chi_mb, d_b, chi_rb) = (shape_b[0], shape_b[1], shape_b[2]);
    assert_eq!(d, 2, "only d=2 (qubit) supported");
    assert_eq!(d_b, 2, "only d=2 (qubit) supported");
    assert_eq!(chi_la, 1, "site_a chi_l must be 1");
    assert_eq!(chi_m, 1, "shared bond must be chi=1");
    assert_eq!(chi_mb, 1, "site_b chi_l must be 1");
    assert_eq!(chi_rb, 1, "site_b chi_r must be 1");
    assert_eq!(gate.len(), d * d * d * d, "gate must be d²×d²");

    // Step 1: contract Θ[s0, s1] = A[0, s0, 0] * B[0, s1, 0]  (chi=1 outer product)
    let mut theta = [0.0f32; 4];
    for s0 in 0..d {
        for s1 in 0..d {
            theta[s0 * d + s1] = data_a[s0] * data_b[s1];
        }
    }

    // Step 2: apply gate  Θ'[s0'*d+s1'] = Σ_{s0,s1} gate[...] * Θ[s0*d+s1]
    let mut theta_prime = [0.0f32; 4];
    for s_out in 0..(d * d) {
        let mut val = 0.0f32;
        for s_in in 0..(d * d) {
            val += gate[s_out * d * d + s_in] * theta[s_in];
        }
        theta_prime[s_out] = val;
    }

    // Step 3: SVD of Θ' as a (d×d) matrix M[s0', s1'] = Θ'[s0'*d + s1']
    let (u, sigma, vt) = svd_2x2(theta_prime);

    // Step 4: keep singular values above threshold, up to chi_max
    const SVD_EPS: f32 = 1e-7;
    let chi_new = (0..d)
        .filter(|&k| sigma[k] > SVD_EPS)
        .count()
        .min(chi_max)
        .max(1);

    // A'[0, s0', k] = U[s0', k] * sigma[k]   shape [1, d, chi_new]
    let mut data_a_new = vec![0.0f32; d * chi_new];
    for s0 in 0..d {
        for k in 0..chi_new {
            data_a_new[s0 * chi_new + k] = u[s0 * d + k] * sigma[k];
        }
    }

    // B'[k, s1', 0] = V^T[k, s1']            shape [chi_new, d, 1]
    let mut data_b_new = vec![0.0f32; chi_new * d];
    for k in 0..chi_new {
        for s1 in 0..d {
            data_b_new[k * d + s1] = vt[k * d + s1];
        }
    }

    set_tensor(&mut nodes[site_a], [1, d, chi_new], data_a_new);
    set_tensor(&mut nodes[site_b], [chi_new, d, 1], data_b_new);

    // Update the edge weight between site_a → site_b to reflect new bond dim.
    let dst_id = nodes[site_b].id;
    if let Some(e) = nodes[site_a]
        .successors
        .iter_mut()
        .find(|e| e.dst == dst_id)
    {
        e.weight = chi_new as f32;
    }

    chi_new
}

// ── Internal helpers ──────────────────────────────────────────────────────────

/// Analytic SVD of a 2×2 real matrix (row-major input).
///
/// Returns `(u, sigma, vt)`:
/// - `u[row * 2 + col]` — left singular vectors as columns (row-major)
/// - `sigma[k]`         — singular values in descending order
/// - `vt[k * 2 + col]`  — right singular vectors as rows (V^T, row-major)
///
/// Only correct for 2×2 matrices; used by `mps_apply_gate_2site`.
fn svd_2x2(m: [f32; 4]) -> ([f32; 4], [f32; 2], [f32; 4]) {
    let (a, b, c, d) = (m[0], m[1], m[2], m[3]);

    // B = M^T M  (symmetric, PSD)
    let b00 = a * a + c * c;
    let b01 = a * b + c * d;
    let b11 = b * b + d * d;

    // Eigenvalues of B via quadratic formula
    let tr = b00 + b11;
    let disc = ((b00 - b11).powi(2) + 4.0 * b01 * b01).sqrt();
    let lam1 = (tr + disc) * 0.5;
    let lam2 = (tr - disc).max(0.0) * 0.5;
    let sig1 = lam1.max(0.0).sqrt();
    let sig2 = lam2.max(0.0).sqrt();

    // Right singular vectors (eigenvectors of B)
    let (v0, v1) = if b01.abs() < 1e-7 {
        // B already diagonal — eigenvectors are standard basis with correct ordering
        if b00 >= b11 {
            ([1.0f32, 0.0f32], [0.0f32, 1.0f32])
        } else {
            ([0.0f32, 1.0f32], [1.0f32, 0.0f32])
        }
    } else {
        let v0x = lam1 - b11;
        let v0y = b01;
        let n = (v0x * v0x + v0y * v0y).sqrt();
        let (v0x, v0y) = (v0x / n, v0y / n);
        ([v0x, v0y], [-v0y, v0x])
    };

    // V^T rows are right singular vectors
    let vt = [v0[0], v0[1], v1[0], v1[1]];

    // Left singular vectors: u_k = M v_k / sigma_k
    let (u0, u1) = {
        let compute = |vx: f32, vy: f32, sig: f32| -> [f32; 2] {
            if sig > 1e-10 {
                [(a * vx + b * vy) / sig, (c * vx + d * vy) / sig]
            } else {
                [0.0, 0.0]
            }
        };
        let u0 = compute(v0[0], v0[1], sig1);
        let u1 = if sig2 > 1e-10 {
            compute(v1[0], v1[1], sig2)
        } else {
            [-u0[1], u0[0]] // null-space: rotate u0 by 90°
        };
        (u0, u1)
    };

    // U as row-major matrix (columns are left singular vectors)
    let u = [u0[0], u1[0], u0[1], u1[1]];

    (u, [sig1, sig2], vt)
}

fn shape_to_json(shape: [usize; 3]) -> Value {
    Value::Array(shape.iter().map(|&x| Value::from(x as u64)).collect())
}

fn data_to_json(data: &[f32]) -> Value {
    Value::Array(data.iter().map(|&x| Value::from(x as f64)).collect())
}

fn json_to_shape(v: &Value) -> Option<[usize; 3]> {
    let arr = v.as_array()?;
    if arr.len() != 3 {
        return None;
    }
    Some([
        arr[0].as_u64()? as usize,
        arr[1].as_u64()? as usize,
        arr[2].as_u64()? as usize,
    ])
}

fn json_to_data(v: &Value) -> Option<Vec<f32>> {
    let arr = v.as_array()?;
    arr.iter().map(|x| x.as_f64().map(|f| f as f32)).collect()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::algorithms::four_d::GraphProperties;

    // ── Helpers ──────────────────────────────────────────────────────────────

    fn blank_node(id: u64, x: f32) -> GraphNode4D {
        GraphNode4D {
            id,
            x,
            y: 0.0,
            z: 0.0,
            begin_ts: 0,
            end_ts: 1,
            properties: GraphProperties::new(),
            successors: vec![],
        }
    }

    /// |0⟩ tensor: shape [1, 2, 1], data [1.0, 0.0]  (up = 1, down = 0)
    fn ket0_tensor() -> ([usize; 3], Vec<f32>) {
        ([1, 2, 1], vec![1.0, 0.0])
    }

    /// |1⟩ tensor: shape [1, 2, 1], data [0.0, 1.0]
    fn ket1_tensor() -> ([usize; 3], Vec<f32>) {
        ([1, 2, 1], vec![0.0, 1.0])
    }

    fn cnot_gate() -> Vec<f32> {
        vec![
            1.0, 0.0, 0.0, 0.0, // |00⟩→|00⟩
            0.0, 1.0, 0.0, 0.0, // |01⟩→|01⟩
            0.0, 0.0, 0.0, 1.0, // |10⟩→|11⟩
            0.0, 0.0, 1.0, 0.0, // |11⟩→|10⟩
        ]
    }

    fn cz_gate() -> Vec<f32> {
        vec![
            1.0, 0.0, 0.0, 0.0, // |00⟩→|00⟩
            0.0, 1.0, 0.0, 0.0, // |01⟩→|01⟩
            0.0, 0.0, 1.0, 0.0, // |10⟩→|10⟩
            0.0, 0.0, 0.0, -1.0, // |11⟩→-|11⟩
        ]
    }

    /// Hadamard gate: [[1/√2, 1/√2], [1/√2, -1/√2]]
    fn hadamard() -> Vec<f32> {
        let s = std::f32::consts::FRAC_1_SQRT_2;
        vec![s, s, s, -s]
    }

    // ── Test 1: tensor round-trip ─────────────────────────────────────────────

    #[test]
    fn test_set_get_tensor_roundtrip() {
        let mut node = blank_node(0, 0.0);
        let shape = [2, 3, 4];
        let data: Vec<f32> = (0..24).map(|i| i as f32 * 0.5).collect();
        set_tensor(&mut node, shape, data.clone());
        let (got_shape, got_data) = get_tensor(&node).expect("tensor should be present");
        assert_eq!(got_shape, shape);
        assert_eq!(got_data.len(), 24);
        for (a, b) in got_data.iter().zip(data.iter()) {
            assert!((a - b).abs() < 1e-5, "data mismatch: {a} vs {b}");
        }
    }

    // ── Test 2: single-site |0⟩ norm ─────────────────────────────────────────

    #[test]
    fn test_single_site_ket0_norm_is_one() {
        let (shape, data) = ket0_tensor();
        let nodes = build_mps(&[(&shape, &data)]);
        let norm = mps_norm_sq(&nodes);
        assert!((norm - 1.0).abs() < 1e-5, "norm = {norm}");
    }

    // ── Test 3: 4-site |0000⟩ norm ───────────────────────────────────────────

    #[test]
    fn test_four_site_product_state_norm_is_one() {
        let (s, d) = ket0_tensor();
        let nodes = build_mps(&[(&s, &d), (&s, &d), (&s, &d), (&s, &d)]);
        let norm = mps_norm_sq(&nodes);
        assert!((norm - 1.0).abs() < 1e-5, "norm = {norm}");
    }

    // ── Test 4: Hadamard preserves norm ──────────────────────────────────────

    #[test]
    fn test_hadamard_gate_preserves_norm() {
        let (s, d) = ket0_tensor();
        let mut nodes = build_mps(&[(&s, &d), (&s, &d), (&s, &d), (&s, &d)]);
        mps_apply_gate(&mut nodes, 0, &hadamard());
        let norm = mps_norm_sq(&nodes);
        assert!((norm - 1.0).abs() < 1e-5, "norm after Hadamard = {norm}");
    }

    // ── Test 5: Bell state has bond dim 2 ────────────────────────────────────

    #[test]
    fn test_bell_state_bond_dim_is_two() {
        // |Φ+⟩ = (|00⟩ + |11⟩)/√2
        // Site 0: shape [1, 2, 2], data = [[1/√2, 0], [0, 1/√2]] row-major
        let s = std::f32::consts::FRAC_1_SQRT_2;
        let a0_shape = [1usize, 2, 2];
        let a0_data: Vec<f32> = vec![s, 0.0, 0.0, s]; // [χl=0,d=0,χr=0],[χl=0,d=1,χr=1]
                                                      // Site 1: shape [2, 2, 1], data = [[1,0],[0,1]] (identity per bond slot)
        let a1_shape = [2usize, 2, 1];
        let a1_data: Vec<f32> = vec![1.0, 0.0, 0.0, 1.0];
        let nodes = build_mps(&[(&a0_shape, &a0_data), (&a1_shape, &a1_data)]);

        // Edge from site 0 to site 1 should carry bond dim = 2
        let bond_dim = nodes[0]
            .successors
            .first()
            .map(|e| e.weight as usize)
            .expect("edge must exist");
        assert_eq!(bond_dim, 2, "bond dim = {bond_dim}");
    }

    // ── Test 6: zero tensor has norm 0 ───────────────────────────────────────

    #[test]
    fn test_zero_tensor_norm_is_zero() {
        let shape = [1usize, 2, 1];
        let data = vec![0.0f32, 0.0];
        let nodes = build_mps(&[(&shape, &data)]);
        let norm = mps_norm_sq(&nodes);
        assert!(norm.abs() < 1e-10, "norm = {norm}");
    }

    // ── Test 8: CNOT on |00⟩ stays product (bond = 1) ────────────────────────

    #[test]
    fn test_2site_cnot_on_00_stays_product() {
        let (s, d) = ket0_tensor();
        let mut nodes = build_mps(&[(&s, &d), (&s, &d)]);
        let chi_new = mps_apply_gate_2site(&mut nodes, 0, 1, &cnot_gate(), 2);
        let norm = mps_norm_sq(&nodes);
        assert!((norm - 1.0).abs() < 1e-5, "norm = {norm}");
        assert_eq!(chi_new, 1, "CNOT|00⟩ is product — bond must stay 1");
    }

    // ── Test 9: CNOT on |10⟩ gives |11⟩ (uses ket1) ─────────────────────────

    #[test]
    fn test_2site_cnot_on_10_gives_11_norm1() {
        let (s0, d0) = ket1_tensor();
        let (s1, d1) = ket0_tensor();
        let mut nodes = build_mps(&[(&s0, &d0), (&s1, &d1)]);
        let chi_new = mps_apply_gate_2site(&mut nodes, 0, 1, &cnot_gate(), 2);
        let norm = mps_norm_sq(&nodes);
        assert!((norm - 1.0).abs() < 1e-5, "CNOT|10⟩ norm = {norm}");
        assert_eq!(chi_new, 1, "CNOT|10⟩=|11⟩ is still product");
    }

    // ── Test 10: CNOT on |+0⟩ creates Bell state (bond = 2, norm = 1) ────────

    #[test]
    fn test_2site_cnot_plus0_creates_bell_norm1() {
        let (s, d) = ket0_tensor();
        let mut nodes = build_mps(&[(&s, &d), (&s, &d)]);
        mps_apply_gate(&mut nodes, 0, &hadamard());
        let chi_new = mps_apply_gate_2site(&mut nodes, 0, 1, &cnot_gate(), 2);
        let norm = mps_norm_sq(&nodes);
        assert!((norm - 1.0).abs() < 1e-4, "Bell state norm = {norm}");
        assert_eq!(chi_new, 2, "Bell state requires bond dim 2");
    }

    // ── Test 11: CZ on |++⟩ preserves norm ───────────────────────────────────

    #[test]
    fn test_2site_cz_plus_plus_preserves_norm() {
        let (s, d) = ket0_tensor();
        let mut nodes = build_mps(&[(&s, &d), (&s, &d)]);
        mps_apply_gate(&mut nodes, 0, &hadamard());
        mps_apply_gate(&mut nodes, 1, &hadamard());
        let _chi_new = mps_apply_gate_2site(&mut nodes, 0, 1, &cz_gate(), 2);
        let norm = mps_norm_sq(&nodes);
        assert!((norm - 1.0).abs() < 1e-4, "CZ|++⟩ norm = {norm}");
    }

    // ── Test 12: chi_max=1 truncates Bell to product (lossy) ─────────────────

    #[test]
    fn test_2site_chi_max_1_truncates() {
        let (s, d) = ket0_tensor();
        let mut nodes = build_mps(&[(&s, &d), (&s, &d)]);
        mps_apply_gate(&mut nodes, 0, &hadamard());
        let chi_new = mps_apply_gate_2site(&mut nodes, 0, 1, &cnot_gate(), 1);
        assert_eq!(chi_new, 1, "truncated to chi_max=1");
        let norm = mps_norm_sq(&nodes);
        assert!(norm <= 1.0 + 1e-5, "truncated norm = {norm} must be ≤ 1");
    }

    // ── Test 7: edge weights encode bond dims for whole chain ─────────────────

    #[test]
    fn test_edge_weights_encode_bond_dimensions() {
        // Chain: bond dims 1, 2, 3 between 4 sites
        // shapes: [1,2,1], [1,2,2], [2,2,3], [3,2,1]
        let shapes: &[&[usize]] = &[&[1, 2, 1], &[1, 2, 2], &[2, 2, 3], &[3, 2, 1]];
        let datas: Vec<Vec<f32>> = shapes
            .iter()
            .map(|s| vec![0.0f32; s[0] * s[1] * s[2]])
            .collect();
        let pairs: Vec<(&[usize], &[f32])> = shapes
            .iter()
            .copied()
            .zip(datas.iter().map(|d| d.as_slice()))
            .collect();
        let nodes = build_mps(&pairs);

        // Each node i (except last) should have edge weight = χ_right of site i
        let expected_bonds = [1.0f32, 2.0, 3.0];
        for (i, &expected) in expected_bonds.iter().enumerate() {
            let w = nodes[i].successors.first().map(|e| e.weight).expect("edge");
            assert!(
                (w - expected).abs() < 1e-6,
                "site {i} bond = {w}, expected {expected}"
            );
        }
    }
}