prism-q 0.35.0

use super::*;
use crate::circuits;

#[test]
fn gf2_kernel_identity() {
    // Identity matrix: kernel is trivial (empty)
    let mut m = F2DenseMatrix::new(3, 3);
    m.set(0, 0);
    m.set(1, 1);
    m.set(2, 2);
    let k = gf2_kernel(&m);
    assert!(k.is_empty(), "Identity matrix should have trivial kernel");
}

#[test]
fn gf2_kernel_zero_matrix() {
    // Zero matrix: kernel is the full space
    let m = F2DenseMatrix::new(3, 4);
    let k = gf2_kernel(&m);
    assert_eq!(k.len(), 4, "Zero 3×4 matrix should have 4-dim kernel");
}

#[test]
fn gf2_kernel_rank_deficient() {
    // [1 1 0]
    // [0 1 1]
    // Row 0 + Row 1 = [1 0 1], so rank = 2, kernel dim = 3 - 2 = 1
    let mut m = F2DenseMatrix::new(2, 3);
    m.set(0, 0);
    m.set(0, 1);
    m.set(1, 1);
    m.set(1, 2);
    let k = gf2_kernel(&m);
    assert_eq!(k.len(), 1, "rank-2 2×3 matrix should have 1-dim kernel");
    // Kernel vector should be [1, 1, 1] (x₀ = x₁ = x₂)
    // Row 0: x₀ + x₁ = 0 → x₀ = x₁
    // Row 1: x₁ + x₂ = 0 → x₁ = x₂
    let kv = &k[0];
    assert_eq!(kv[0] & 0b111, 0b111, "kernel vector should be [1,1,1]");
}

#[test]
fn gf2_kernel_verifies() {
    // Verify Mx = 0 for all kernel vectors
    let mut m = F2DenseMatrix::new(3, 5);
    // Some arbitrary matrix
    m.set(0, 0);
    m.set(0, 2);
    m.set(0, 4);
    m.set(1, 1);
    m.set(1, 3);
    m.set(2, 0);
    m.set(2, 1);
    m.set(2, 2);

    let k = gf2_kernel(&m);
    for kv in &k {
        // Check M · kv = 0
        for r in 0..3 {
            let mut dot = 0u32;
            for c in 0..5 {
                if m.get(r, c) && (kv[c / 64] >> (c % 64)) & 1 != 0 {
                    dot ^= 1;
                }
            }
            assert_eq!(dot, 0, "kernel vector should satisfy Mx = 0");
        }
    }
}

#[test]
fn boundary_trivial_circuit_has_zero_homology() {
    let n = 4;
    let mut circuit = crate::circuit::Circuit::new(n, n);
    for i in 0..n {
        circuit.add_measure(i, i);
    }
    let noise = NoiseModel::uniform_depolarizing(&circuit, 0.001);
    let ecc = ErrorChainComplex::build(&circuit, &noise, 42).unwrap();
    assert_eq!(ecc.boundary_dim(), n);
    assert_eq!(ecc.homology_dim(), 0);
}

#[test]
fn boundary_ghz_has_one_logical_qubit() {
    for n in [3, 5, 8] {
        let mut circuit = circuits::ghz_circuit(n);
        circuit.measure_all();
        let noise = NoiseModel::uniform_depolarizing(&circuit, 0.001);
        let ecc = ErrorChainComplex::build(&circuit, &noise, 42).unwrap();
        assert_eq!(
            ecc.homology_dim(),
            1,
            "GHZ-{n} should have 1 logical error class"
        );
        assert_eq!(ecc.boundary_dim(), n - 1);
    }
}

#[test]
fn boundary_bell_pair_has_one_logical() {
    let mut circuit = crate::circuit::Circuit::new(2, 2);
    circuit.add_gate(crate::gates::Gate::H, &[0]);
    circuit.add_gate(crate::gates::Gate::Cx, &[0, 1]);
    circuit.add_measure(0, 0);
    circuit.add_measure(1, 1);
    let noise = NoiseModel::uniform_depolarizing(&circuit, 0.001);
    let ecc = ErrorChainComplex::build(&circuit, &noise, 42).unwrap();
    assert_eq!(ecc.homology_dim(), 1);
    assert_eq!(ecc.boundary_dim(), 1);
}

#[test]
fn boundary_independent_bell_pairs() {
    let n_pairs = 3;
    let mut circuit = circuits::independent_bell_pairs(n_pairs);
    circuit.measure_all();
    let noise = NoiseModel::uniform_depolarizing(&circuit, 0.001);
    let ecc = ErrorChainComplex::build(&circuit, &noise, 42).unwrap();
    assert_eq!(
        ecc.homology_dim(),
        n_pairs,
        "{n_pairs} bell pairs should have {n_pairs} logical error classes"
    );
}

#[test]
fn boundary_partial_measurement() {
    let mut circuit = crate::circuit::Circuit::new(3, 1);
    circuit.add_gate(crate::gates::Gate::H, &[0]);
    circuit.add_gate(crate::gates::Gate::Cx, &[0, 1]);
    circuit.add_measure(0, 0);
    let noise = NoiseModel::uniform_depolarizing(&circuit, 0.001);
    let ecc = ErrorChainComplex::build(&circuit, &noise, 42).unwrap();
    // 3 qubits, 1 measured: ker(σ) has dim 2*3-1=5
    // Stabilizers: X₀X₁, Z₀Z₁, Z₂ (3 generators)
    // X-projection on qubit 0: X₀X₁ has X on q0 → rank(A) = 1
    // boundary_dim = 3-1 = 2, homology_dim = 3-1+1 = 3
    assert_eq!(ecc.boundary_dim(), 2);
    assert_eq!(ecc.homology_dim(), 3);
}

#[test]
fn analytical_marginals_fold_in_readout_error() {
    // Deterministic bits under a readout-only model, so the closed form is the
    // whole answer: bit 0 is measured 1 and reads 1 with probability 1 - p10,
    // bit 1 is measured 0 and reads 1 with probability p01.
    let mut circuit = Circuit::new(2, 2);
    circuit.add_gate(crate::gates::Gate::X, &[0]);
    circuit.measure_all();
    let mut noise = NoiseModel::uniform_depolarizing(&circuit, 0.0);
    noise.set_bit_readout_error(0, 0.02, 0.07);
    noise.set_bit_readout_error(1, 0.03, 0.11);

    let marginals = noisy_marginals_analytical(&circuit, &noise, 42).unwrap();
    assert!((marginals[0] - 0.93).abs() < 1e-12, "{marginals:?}");
    assert!((marginals[1] - 0.03).abs() < 1e-12, "{marginals:?}");
}

#[test]
fn analytical_marginals_with_readout_match_sampled() {
    let n = 6;
    let mut circuit = circuits::ghz_circuit(n);
    circuit.measure_all();
    let mut noise = NoiseModel::uniform_depolarizing(&circuit, 0.01);
    for bit in 0..n {
        noise.set_bit_readout_error(bit, 0.04, 0.09);
    }

    let analytical = noisy_marginals_analytical(&circuit, &noise, 42).unwrap();

    let shots = 20_000;
    let sampled = crate::sim::noise::run_shots_noisy(&circuit, &noise, shots, 42).unwrap();
    let counts: Vec<f64> = (0..n)
        .map(|bit| sampled.shots.iter().filter(|record| record[bit]).count() as f64 / shots as f64)
        .collect();

    for bit in 0..n {
        // Five sigma on a binomial proportion at this shot count.
        let sigma = (analytical[bit] * (1.0 - analytical[bit]) / shots as f64).sqrt();
        assert!(
            (analytical[bit] - counts[bit]).abs() < 5.0 * sigma,
            "bit {bit}: analytical={:.6}, sampled={:.6}, sigma={sigma:.6}",
            analytical[bit],
            counts[bit],
        );
    }
}

#[test]
fn analytical_marginals_ghz_50q() {
    let n = 50;
    let mut circuit = circuits::ghz_circuit(n);
    circuit.measure_all();
    let noise = NoiseModel::uniform_depolarizing(&circuit, 0.001);

    let marginals = noisy_marginals_analytical(&circuit, &noise, 42).unwrap();
    assert_eq!(marginals.len(), n);
    for (i, &p) in marginals.iter().enumerate() {
        assert!(p > 0.0 && p < 1.0, "bit {i}: marginal {p} out of range");
        assert!(
            (p - 0.5).abs() < 0.05,
            "bit {i}: GHZ marginal should be near 0.5, got {p}"
        );
    }
}

#[test]
fn analytical_marginals_ghz_100q() {
    let n = 100;
    let mut circuit = circuits::ghz_circuit(n);
    circuit.measure_all();
    let noise = NoiseModel::uniform_depolarizing(&circuit, 0.001);

    let marginals = noisy_marginals_analytical(&circuit, &noise, 42).unwrap();
    assert_eq!(marginals.len(), n);
    for (i, &p) in marginals.iter().enumerate() {
        assert!(p > 0.0 && p < 1.0, "bit {i}: marginal {p} out of range");
    }
}

#[test]
fn analytical_marginals_bell_pairs_100q() {
    let n_pairs = 50;
    let n = n_pairs * 2;
    let mut circuit = circuits::independent_bell_pairs(n_pairs);
    circuit.measure_all();
    let noise = NoiseModel::uniform_depolarizing(&circuit, 0.001);

    let marginals = noisy_marginals_analytical(&circuit, &noise, 42).unwrap();
    assert_eq!(marginals.len(), n);
    for (i, &p) in marginals.iter().enumerate() {
        assert!(
            (p - 0.5).abs() < 0.05,
            "bit {i}: bell pair marginal should be near 0.5, got {p}"
        );
    }
}

#[test]
fn analytical_marginals_clifford_1000q() {
    let n = 1000;
    let mut circuit = circuits::clifford_heavy_circuit(n, 2, 42);
    circuit.measure_all();
    let noise = NoiseModel::uniform_depolarizing(&circuit, 0.001);

    let marginals = noisy_marginals_analytical(&circuit, &noise, 42).unwrap();
    assert_eq!(marginals.len(), n);
    for (i, &p) in marginals.iter().enumerate() {
        assert!(
            (0.0..=1.0).contains(&p),
            "bit {i}: marginal {p} out of range"
        );
    }
}

#[test]
fn analytical_marginals_deterministic_bits() {
    let mut circuit = crate::circuit::Circuit::new(4, 4);
    for i in 0..4 {
        circuit.add_gate(crate::gates::Gate::X, &[i]);
    }
    for i in 0..4 {
        circuit.add_measure(i, i);
    }
    let noise = NoiseModel::uniform_depolarizing(&circuit, 0.01);

    let marginals = noisy_marginals_analytical(&circuit, &noise, 42).unwrap();
    for (i, &p) in marginals.iter().enumerate() {
        assert!(
            p > 0.95,
            "bit {i}: X-then-measure should give p(1) near 1.0, got {p}"
        );
    }
}

#[test]
fn analytical_marginals_no_noise() {
    let n = 6;
    let mut circuit = circuits::ghz_circuit(n);
    circuit.measure_all();
    let noise = NoiseModel::uniform_depolarizing(&circuit, 0.0);

    let marginals = noisy_marginals_analytical(&circuit, &noise, 42).unwrap();
    for (i, &p) in marginals.iter().enumerate() {
        assert!(
            (p - 0.5).abs() < 1e-10,
            "bit {i}: GHZ with no noise should have marginal 0.5, got {p}"
        );
    }
}