geographdb-core 0.5.4

//! 4D Tensor Network backed by the spatiotemporal graph.
//!
//! Sites occupy a (nx × ny × nz) spatial grid. Each site exists at `depth`
//! time layers. The full structure has `nx * ny * nz * depth` nodes.
//!
//! Two edge types share the same `TemporalEdge` struct — distinguished by the
//! destination node's layer:
//!
//!   Spatial edge   — dst.begin_ts == src.begin_ts  (same time layer, adjacent space)
//!   Temporal edge  — dst.begin_ts >  src.begin_ts  (same space, next time layer)
//!
//! Tensor shape per node: `[χ_t_prev, d, χ_t_next]`
//!   - χ_t_prev = temporal bond to previous layer (1 for layer 0)
//!   - d        = physical dimension (2 for qubits)
//!   - χ_t_next = temporal bond to next layer (1 for final layer)
//!
//! Spatial bonds are recorded as edges with weight = χ_spatial and do not
//! change the tensor rank in this model. They carry entanglement information
//! but full spatial contraction requires a separate PEPS engine (future work).
//!
//! Node id scheme: id = tl * nz*ny*nx  +  sz * ny*nx  +  sy * nx  +  sx

use crate::algorithms::four_d::{GraphNode4D, GraphProperties, TemporalEdge};
use crate::algorithms::mps::{get_tensor, mps_apply_gate_2site, mps_norm_sq, set_tensor};

/// Spatial grid dimensions passed as a single argument to avoid clippy `too_many_arguments`.
#[derive(Clone, Copy, Debug)]
pub struct GridDims {
    pub nx: usize,
    pub ny: usize,
    pub nz: usize,
}

/// A single spatial coordinate `(sx, sy, sz)` within a [`GridDims`] grid.
#[derive(Clone, Copy, Debug)]
pub struct SiteCoord {
    pub sx: usize,
    pub sy: usize,
    pub sz: usize,
}

// ── Identity helpers ──────────────────────────────────────────────────────────

/// Compute the node id for site (sx, sy, sz) at time layer tl.
pub fn site_id(sx: usize, sy: usize, sz: usize, tl: usize, nx: usize, ny: usize, nz: usize) -> u64 {
    (tl * nz * ny * nx + sz * ny * nx + sy * nx + sx) as u64
}

// ── Construction ──────────────────────────────────────────────────────────────

/// Build a (nx × ny × nz) × depth tensor network, all sites initialised to |0⟩.
///
/// Physical dimension is 2 (qubit). All bond dimensions are 1 (product state).
/// Spatial edges connect nearest neighbours within each layer; temporal edges
/// connect the same site across adjacent layers.
pub fn build_tnet4d(nx: usize, ny: usize, nz: usize, depth: usize) -> Vec<GraphNode4D> {
    let total = nx * ny * nz * depth;
    let mut nodes: Vec<GraphNode4D> = Vec::with_capacity(total);

    // Phase 1: create all nodes with |0⟩ tensors.
    for tl in 0..depth {
        for sz in 0..nz {
            for sy in 0..ny {
                for sx in 0..nx {
                    let id = site_id(sx, sy, sz, tl, nx, ny, nz);
                    let mut node = GraphNode4D {
                        id,
                        x: sx as f32,
                        y: sy as f32,
                        z: sz as f32,
                        begin_ts: tl as u64,
                        end_ts: tl as u64 + 1,
                        properties: GraphProperties::new(),
                        successors: Vec::new(),
                    };
                    // |0⟩: shape [1, 2, 1], data [1.0, 0.0]
                    set_tensor(&mut node, [1, 2, 1], vec![1.0, 0.0]);
                    nodes.push(node);
                }
            }
        }
    }

    // Phase 2: wire spatial edges (same layer, adjacent positions).
    let directions: &[(i32, i32, i32)] = &[
        (1, 0, 0),
        (-1, 0, 0),
        (0, 1, 0),
        (0, -1, 0),
        (0, 0, 1),
        (0, 0, -1),
    ];
    for tl in 0..depth {
        for sz in 0..nz {
            for sy in 0..ny {
                for sx in 0..nx {
                    let src_id = site_id(sx, sy, sz, tl, nx, ny, nz);
                    for &(dx, dy, dz) in directions {
                        let nx2 = sx as i32 + dx;
                        let ny2 = sy as i32 + dy;
                        let nz2 = sz as i32 + dz;
                        if nx2 < 0 || nx2 >= nx as i32 {
                            continue;
                        }
                        if ny2 < 0 || ny2 >= ny as i32 {
                            continue;
                        }
                        if nz2 < 0 || nz2 >= nz as i32 {
                            continue;
                        }
                        let dst_id =
                            site_id(nx2 as usize, ny2 as usize, nz2 as usize, tl, nx, ny, nz);
                        let src_idx = src_id as usize;
                        nodes[src_idx].successors.push(TemporalEdge {
                            dst: dst_id,
                            weight: 1.0, // spatial bond dim = 1
                            begin_ts: tl as u64,
                            end_ts: tl as u64, // begin_ts == end_ts marks spatial
                        });
                    }
                }
            }
        }
    }

    // Phase 3: wire temporal edges (same position, next layer).
    for tl in 0..depth.saturating_sub(1) {
        for sz in 0..nz {
            for sy in 0..ny {
                for sx in 0..nx {
                    let src_id = site_id(sx, sy, sz, tl, nx, ny, nz);
                    let dst_id = site_id(sx, sy, sz, tl + 1, nx, ny, nz);
                    let src_idx = src_id as usize;
                    nodes[src_idx].successors.push(TemporalEdge {
                        dst: dst_id,
                        weight: 1.0, // temporal bond dim = 1
                        begin_ts: tl as u64,
                        end_ts: tl as u64 + 1, // begin_ts < end_ts marks temporal
                    });
                }
            }
        }
    }

    nodes
}

// ── Lookup ────────────────────────────────────────────────────────────────────

/// Return the slice index of the node at (sx, sy, sz, tl), or `None` if out of bounds.
pub fn tnet4d_find(
    nodes: &[GraphNode4D],
    sx: usize,
    sy: usize,
    sz: usize,
    tl: usize,
    dims: GridDims,
) -> Option<usize> {
    if sx >= dims.nx || sy >= dims.ny || sz >= dims.nz {
        return None;
    }
    let target_id = site_id(sx, sy, sz, tl, dims.nx, dims.ny, dims.nz);
    nodes.iter().position(|n| n.id == target_id)
}

// ── Gates ─────────────────────────────────────────────────────────────────────

/// Apply a d×d single-site gate at (sx, sy, sz) in layer tl.
///
/// The gate is a d×d matrix in row-major order. Does nothing if the site does
/// not exist.
pub fn tnet4d_apply_gate_1site(
    nodes: &mut [GraphNode4D],
    sx: usize,
    sy: usize,
    sz: usize,
    tl: usize,
    dims: GridDims,
    gate: &[f32],
) {
    if let Some(idx) = tnet4d_find(nodes, sx, sy, sz, tl, dims) {
        let (shape, data) = match get_tensor(&nodes[idx]) {
            Some(t) => t,
            None => return,
        };
        let (chi_l, d, chi_r) = (shape[0], shape[1], shape[2]);
        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[idx], shape, new_data);
    }
}

/// Apply a two-site gate to two spatially adjacent sites at time layer `tl`.
///
/// `site0` and `site1` must be neighbors in the grid (differ by one step in one
/// axis). The gate is a `d²×d²` unitary; see [`mps_apply_gate_2site`] for the
/// index convention. Only works for χ_in = 1 (product state) and d = 2.
///
/// Returns the actual bond dimension kept after SVD truncation.
pub fn tnet4d_apply_gate_2site(
    nodes: &mut [GraphNode4D],
    site0: SiteCoord,
    site1: SiteCoord,
    tl: usize,
    dims: GridDims,
    gate: &[f32],
    chi_max: usize,
) -> usize {
    let idx_a = match tnet4d_find(nodes, site0.sx, site0.sy, site0.sz, tl, dims) {
        Some(i) => i,
        None => return 0,
    };
    let idx_b = match tnet4d_find(nodes, site1.sx, site1.sy, site1.sz, tl, dims) {
        Some(i) => i,
        None => return 0,
    };
    mps_apply_gate_2site(nodes, idx_a, idx_b, gate, chi_max)
}

// ── Norm ──────────────────────────────────────────────────────────────────────

/// Compute ⟨ψ|ψ⟩ for a product-state 4D network (all spatial bonds χ=1).
///
/// Contracts each spatial column (fixed x,y,z through all layers) as a 1D MPS
/// using `mps_norm_sq`, then multiplies all column norms together.
///
/// For entangled states (spatial bond dim > 1) a full PEPS contraction is
/// required — this function returns the correct answer only for χ_spatial = 1.
pub fn tnet4d_norm_sq(nodes: &[GraphNode4D], nx: usize, ny: usize, nz: usize, depth: usize) -> f64 {
    // For product states (all spatial bonds χ=1), the norm factors into
    // independent temporal columns. Contract each column as a 1D MPS.
    let mut total_norm = 1.0f64;
    for sz in 0..nz {
        for sy in 0..ny {
            for sx in 0..nx {
                let dims = GridDims { nx, ny, nz };
                let column: Vec<GraphNode4D> = (0..depth)
                    .filter_map(|tl| {
                        tnet4d_find(nodes, sx, sy, sz, tl, dims).map(|idx| nodes[idx].clone())
                    })
                    .collect();
                total_norm *= mps_norm_sq(&column);
            }
        }
    }
    total_norm
}

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

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

    // ── test 1: node count 2×2×1×1 ───────────────────────────────────────────

    #[test]
    fn test_build_node_count_single_layer() {
        let nodes = build_tnet4d(2, 2, 1, 1);
        assert_eq!(nodes.len(), 4, "2×2×1 grid × 1 layer = 4 nodes");
    }

    // ── test 2: node count 2×2×1×2 ───────────────────────────────────────────

    #[test]
    fn test_build_node_count_two_layers() {
        let nodes = build_tnet4d(2, 2, 1, 2);
        assert_eq!(nodes.len(), 8, "2×2×1 grid × 2 layers = 8 nodes");
    }

    // ── test 3: temporal edges connect layers ─────────────────────────────────

    #[test]
    fn test_temporal_edges_connect_layers() {
        let nodes = build_tnet4d(2, 2, 1, 2);
        // Node at (0,0,0,tl=0) should have a temporal edge to (0,0,0,tl=1).
        let src_id = site_id(0, 0, 0, 0, 2, 2, 1);
        let dst_id = site_id(0, 0, 0, 1, 2, 2, 1);
        let src = nodes.iter().find(|n| n.id == src_id).expect("src node");
        let has_temporal = src.successors.iter().any(|e| e.dst == dst_id);
        assert!(has_temporal, "no temporal edge from layer 0 to layer 1");
        // Temporal edge: dst node is at a later layer (begin_ts differs).
        let dst = nodes.iter().find(|n| n.id == dst_id).expect("dst node");
        assert!(
            dst.begin_ts > src.begin_ts,
            "temporal dst should be at a later time layer"
        );
    }

    // ── test 4: spatial edges connect grid neighbours ────────────────────────

    #[test]
    fn test_spatial_edges_connect_neighbours() {
        let nodes = build_tnet4d(2, 2, 1, 1);
        // Node at (0,0,0,tl=0) should have spatial edges to (1,0,0) and (0,1,0).
        let src_id = site_id(0, 0, 0, 0, 2, 2, 1);
        let right_id = site_id(1, 0, 0, 0, 2, 2, 1);
        let up_id = site_id(0, 1, 0, 0, 2, 2, 1);
        let src = nodes.iter().find(|n| n.id == src_id).expect("src");
        let has_right = src.successors.iter().any(|e| e.dst == right_id);
        let has_up = src.successors.iter().any(|e| e.dst == up_id);
        assert!(has_right, "missing spatial edge to (1,0,0)");
        assert!(has_up, "missing spatial edge to (0,1,0)");
        // Spatial edges: dst node is at the same layer (same begin_ts as src).
        let right = nodes.iter().find(|n| n.id == right_id).expect("right");
        assert_eq!(
            right.begin_ts, src.begin_ts,
            "spatial neighbour must be same layer"
        );
    }

    // ── test 5: tnet4d_find locates nodes correctly ───────────────────────────

    #[test]
    fn test_find_locates_correct_node() {
        let nodes = build_tnet4d(2, 2, 1, 2);
        let dims = GridDims {
            nx: 2,
            ny: 2,
            nz: 1,
        };
        // Site (1,1,0) at layer 0 should be index 3 in a 2×2×1×1 layout.
        let idx = tnet4d_find(&nodes, 1, 1, 0, 0, dims).expect("node must exist");
        let expected_id = site_id(1, 1, 0, 0, 2, 2, 1);
        assert_eq!(nodes[idx].id, expected_id);
        // Out of bounds returns None.
        assert!(tnet4d_find(&nodes, 5, 0, 0, 0, dims).is_none());
    }

    // ── test 6: all-|0⟩ product state norm = 1.0 ─────────────────────────────

    #[test]
    fn test_product_state_norm_is_one() {
        let nodes = build_tnet4d(2, 2, 1, 1);
        let norm = tnet4d_norm_sq(&nodes, 2, 2, 1, 1);
        assert!((norm - 1.0).abs() < 1e-5, "norm = {norm}");
    }

    fn cnot_gate() -> Vec<f32> {
        vec![
            1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 0.0,
        ]
    }

    // ── test 7: Hadamard on all sites preserves norm ──────────────────────────

    #[test]
    fn test_hadamard_all_sites_preserves_norm() {
        let mut nodes = build_tnet4d(2, 2, 1, 1);
        let h = hadamard();
        let dims = GridDims {
            nx: 2,
            ny: 2,
            nz: 1,
        };
        for sx in 0..2 {
            for sy in 0..2 {
                tnet4d_apply_gate_1site(&mut nodes, sx, sy, 0, 0, dims, &h);
            }
        }
        let norm = tnet4d_norm_sq(&nodes, 2, 2, 1, 1);
        assert!((norm - 1.0).abs() < 1e-5, "norm after H = {norm}");
    }

    // ── test 8: 2-site CNOT creates entanglement between spatial neighbors ────

    #[test]
    fn test_tnet4d_2site_gate_creates_entanglement() {
        let mut nodes = build_tnet4d(2, 2, 1, 1);
        let dims = GridDims {
            nx: 2,
            ny: 2,
            nz: 1,
        };
        // H on (0,0,0): |+⟩|0⟩
        tnet4d_apply_gate_1site(&mut nodes, 0, 0, 0, 0, dims, &hadamard());
        // CNOT (0,0,0)↔(1,0,0) → Bell pair
        let site0 = SiteCoord {
            sx: 0,
            sy: 0,
            sz: 0,
        };
        let site1 = SiteCoord {
            sx: 1,
            sy: 0,
            sz: 0,
        };
        let chi_new = tnet4d_apply_gate_2site(&mut nodes, site0, site1, 0, dims, &cnot_gate(), 2);
        assert_eq!(
            chi_new, 2,
            "Bell pair between spatial neighbours: bond must be 2"
        );
    }
}