prism-q 0.19.0

use super::DistributedStatevectorBackend;
use crate::backend::statevector::StatevectorBackend;
use crate::backend::Backend;
use crate::circuit::builder::CircuitBuilder;
use crate::circuit::Circuit;
use crate::distributed::{DistributedContext, RankComm};
use crate::sim::run_on;
use num_complex::Complex64;
use std::collections::VecDeque;
use std::sync::{Arc, Condvar, Mutex};

const SEED: u64 = 42;
const TOL: f64 = 1e-10;

/// Test transport that simulates ranks with threads.
///
/// This covers global gate exchange without an MPI launcher. Each rank runs the
/// same circuit and reaches collectives in the same order.
struct LoopbackShared {
    size: usize,
    state: Mutex<LoopbackState>,
    cv: Condvar,
}

struct LoopbackState {
    generation: u64,
    arrived: usize,
    cslots: Vec<Vec<Complex64>>,
    fslots: Vec<f64>,
    // Mailboxes indexed by `sender * size + receiver`. FIFO order matches MPI
    // sendrecv order and stays separate from collective barriers.
    mailbox: Vec<VecDeque<Vec<Complex64>>>,
}

impl LoopbackShared {
    fn new(size: usize) -> Arc<Self> {
        Arc::new(Self {
            size,
            state: Mutex::new(LoopbackState {
                generation: 0,
                arrived: 0,
                cslots: vec![Vec::new(); size],
                fslots: vec![0.0; size],
                mailbox: (0..size * size).map(|_| VecDeque::new()).collect(),
            }),
            cv: Condvar::new(),
        })
    }

    fn barrier(&self) {
        let mut st = self.state.lock().unwrap();
        let gen = st.generation;
        st.arrived += 1;
        if st.arrived == self.size {
            st.arrived = 0;
            st.generation = st.generation.wrapping_add(1);
            self.cv.notify_all();
        } else {
            while st.generation == gen {
                st = self.cv.wait(st).unwrap();
            }
        }
    }
}

#[derive(Clone)]
struct LoopbackComm {
    shared: Arc<LoopbackShared>,
    rank: usize,
}

impl std::fmt::Debug for LoopbackComm {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("LoopbackComm")
            .field("rank", &self.rank)
            .field("size", &self.shared.size)
            .finish()
    }
}

impl RankComm for LoopbackComm {
    fn rank(&self) -> usize {
        self.rank
    }

    fn size(&self) -> usize {
        self.shared.size
    }

    fn allgather_c64(&self, local: &[Complex64]) -> Vec<Complex64> {
        {
            let mut st = self.shared.state.lock().unwrap();
            st.cslots[self.rank] = local.to_vec();
        }
        self.shared.barrier();
        let out = {
            let st = self.shared.state.lock().unwrap();
            st.cslots.iter().flat_map(|s| s.iter().copied()).collect()
        };
        self.shared.barrier();
        out
    }

    fn allgather_f64(&self, local: &[f64]) -> Vec<f64> {
        let as_c: Vec<Complex64> = local.iter().map(|&v| Complex64::new(v, 0.0)).collect();
        self.allgather_c64(&as_c).iter().map(|c| c.re).collect()
    }

    fn allreduce_sum_f64(&self, value: f64) -> f64 {
        {
            let mut st = self.shared.state.lock().unwrap();
            st.fslots[self.rank] = value;
        }
        self.shared.barrier();
        let sum = {
            let st = self.shared.state.lock().unwrap();
            st.fslots.iter().sum()
        };
        self.shared.barrier();
        sum
    }

    fn sendrecv_c64(&self, partner: usize, send: &[Complex64], recv: &mut [Complex64]) {
        debug_assert_eq!(send.len(), recv.len());
        let size = self.shared.size;
        let mut st = self.shared.state.lock().unwrap();
        // Send to partner, then wait for partner to send back. Ranks that skip
        // an exchange do not block because their partner skips it too.
        st.mailbox[self.rank * size + partner].push_back(send.to_vec());
        self.shared.cv.notify_all();
        let inbox = partner * size + self.rank;
        loop {
            if let Some(msg) = st.mailbox[inbox].pop_front() {
                recv.copy_from_slice(&msg);
                return;
            }
            st = self.shared.cv.wait(st).unwrap();
        }
    }

    fn barrier(&self) {
        self.shared.barrier();
    }
}

/// Run `circuit` across simulated ranks with the given backend configuration
/// and return rank 0's probabilities.
fn loopback_probs_with(circuit: &Circuit, size: usize, chunk: usize, relabel: bool) -> Vec<f64> {
    let shared = LoopbackShared::new(size);
    let handles: Vec<_> = (0..size)
        .map(|rank| {
            let comm = LoopbackComm {
                shared: shared.clone(),
                rank,
            };
            let circuit = circuit.clone();
            std::thread::Builder::new()
                .stack_size(64 * 1024 * 1024)
                .spawn(move || {
                    let ctx = DistributedContext::from_comm(Arc::new(comm));
                    let mut backend = DistributedStatevectorBackend::new(ctx, SEED);
                    backend.set_exchange_chunk(chunk);
                    backend.set_relabel(relabel);
                    run_on(&mut backend, &circuit)
                        .expect("distributed run")
                        .probabilities
                        .expect("probabilities")
                        .to_vec()
                })
                .expect("spawn rank thread")
        })
        .collect();
    let mut results: Vec<Vec<f64>> = handles.into_iter().map(|h| h.join().unwrap()).collect();
    // Every rank holds the same gathered vector.
    let first = results.swap_remove(0);
    results.clear();
    first
}

fn assert_loopback_matches(circuit: &Circuit, sizes: &[usize]) {
    let expected = reference_probs(circuit);
    for &relabel in &[true, false] {
        for &size in sizes {
            let actual = loopback_probs_with(circuit, size, usize::MAX, relabel);
            assert_eq!(
                expected.len(),
                actual.len(),
                "length mismatch at size {size} relabel {relabel}"
            );
            for (i, (e, a)) in expected.iter().zip(actual.iter()).enumerate() {
                assert!(
                    (e - a).abs() < TOL,
                    "size {size} relabel {relabel}: prob[{i}] expected {e}, got {a}"
                );
            }
        }
    }
}

/// Run `circuit` across simulated ranks and return rank 0's probabilities and
/// classical bits.
fn loopback_run_with(circuit: &Circuit, size: usize, relabel: bool) -> (Vec<f64>, Vec<bool>) {
    let shared = LoopbackShared::new(size);
    let handles: Vec<_> = (0..size)
        .map(|rank| {
            let comm = LoopbackComm {
                shared: shared.clone(),
                rank,
            };
            let circuit = circuit.clone();
            std::thread::Builder::new()
                .stack_size(64 * 1024 * 1024)
                .spawn(move || {
                    let ctx = DistributedContext::from_comm(Arc::new(comm));
                    let mut backend = DistributedStatevectorBackend::new(ctx, SEED);
                    backend.set_relabel(relabel);
                    let out = run_on(&mut backend, &circuit).expect("distributed run");
                    let probs = out.probabilities.expect("probabilities").to_vec();
                    (probs, out.classical_bits)
                })
                .expect("spawn rank thread")
        })
        .collect();
    let mut results: Vec<(Vec<f64>, Vec<bool>)> =
        handles.into_iter().map(|h| h.join().unwrap()).collect();
    let first = results.swap_remove(0);
    results.clear();
    first
}

fn loopback_run(circuit: &Circuit, size: usize) -> (Vec<f64>, Vec<bool>) {
    loopback_run_with(circuit, size, true)
}

/// Assert that probabilities and classical bits are identical across all rank
/// counts. This checks measurement determinism across rank counts.
fn assert_loopback_deterministic(circuit: &Circuit, sizes: &[usize]) {
    for &relabel in &[true, false] {
        let (ref_probs, ref_bits) = loopback_run_with(circuit, sizes[0], relabel);
        for &size in &sizes[1..] {
            let (probs, bits) = loopback_run_with(circuit, size, relabel);
            assert_eq!(
                ref_bits, bits,
                "classical bits differ at size {size} relabel {relabel}"
            );
            assert_eq!(
                ref_probs.len(),
                probs.len(),
                "length differs at size {size} relabel {relabel}"
            );
            for (i, (e, a)) in ref_probs.iter().zip(probs.iter()).enumerate() {
                assert!(
                    (e - a).abs() < TOL,
                    "size {size} relabel {relabel}: prob[{i}] {e} vs {a} differ across ranks"
                );
            }
        }
    }
}

fn reference_probs(circuit: &Circuit) -> Vec<f64> {
    let mut backend = StatevectorBackend::new(SEED);
    run_on(&mut backend, circuit)
        .expect("statevector run")
        .probabilities
        .expect("probabilities")
        .to_vec()
}

fn distributed_serial_probs(circuit: &Circuit) -> Vec<f64> {
    let ctx = DistributedContext::serial();
    let mut backend = DistributedStatevectorBackend::new(ctx, SEED);
    run_on(&mut backend, circuit)
        .expect("distributed run")
        .probabilities
        .expect("probabilities")
        .to_vec()
}

fn assert_probs_match(circuit: &Circuit) {
    let expected = reference_probs(circuit);
    let actual = distributed_serial_probs(circuit);
    assert_eq!(expected.len(), actual.len(), "length mismatch");
    for (i, (e, a)) in expected.iter().zip(actual.iter()).enumerate() {
        assert!(
            (e - a).abs() < TOL,
            "prob[{i}] mismatch: expected {e}, got {a}"
        );
    }
}

#[test]
fn serial_matches_statevector_bell() {
    let mut b = CircuitBuilder::new(2);
    b.h(0).cx(0, 1);
    assert_probs_match(&b.build());
}

#[test]
fn serial_matches_statevector_rotations_and_entanglers() {
    let mut b = CircuitBuilder::new(4);
    b.h(0)
        .rx(0.37, 1)
        .ry(1.1, 2)
        .rz(-0.6, 3)
        .cx(0, 1)
        .cz(1, 2)
        .swap(2, 3)
        .t(0)
        .s(1);
    assert_probs_match(&b.build());
}

#[test]
fn serial_matches_statevector_ghz() {
    let n = 6;
    let mut b = CircuitBuilder::new(n);
    b.h(0);
    for q in 0..n - 1 {
        b.cx(q, q + 1);
    }
    assert_probs_match(&b.build());
}

#[test]
fn serial_export_statevector_matches() {
    let mut b = CircuitBuilder::new(3);
    b.h(0).cx(0, 1).ry(0.9, 2).cz(0, 2);
    let circuit = b.build();

    let mut sv = StatevectorBackend::new(SEED);
    sv.init(circuit.num_qubits, circuit.num_classical_bits)
        .unwrap();
    sv.apply_instructions(&circuit.instructions).unwrap();
    let expected = sv.export_statevector().unwrap();

    let ctx = DistributedContext::serial();
    let mut dist = DistributedStatevectorBackend::new(ctx, SEED);
    dist.init(circuit.num_qubits, circuit.num_classical_bits)
        .unwrap();
    dist.apply_instructions(&circuit.instructions).unwrap();
    let actual = dist.export_statevector().unwrap();

    assert_eq!(expected.len(), actual.len());
    for (e, a) in expected.iter().zip(actual.iter()) {
        assert!((e - a).norm() < TOL);
    }
}

/// Relax the local qubit floor before constructing any distributed backend.
fn relax_min_local_qubits() {
    std::env::set_var("PRISM_DIST_MIN_LOCAL_QUBITS", "1");
    assert_eq!(crate::distributed::min_local_qubits(), 1);
}

#[test]
fn loopback_global_hadamard_wall() {
    relax_min_local_qubits();
    let mut b = CircuitBuilder::new(4);
    for q in 0..4 {
        b.h(q);
    }
    assert_loopback_matches(&b.build(), &[1, 2, 4]);
}

#[test]
fn loopback_global_rotations_and_diagonals() {
    relax_min_local_qubits();
    let mut b = CircuitBuilder::new(4);
    b.h(0)
        .rx(0.7, 3)
        .ry(1.3, 2)
        .rz(-0.4, 3)
        .t(2)
        .s(3)
        .x(1)
        .h(2)
        .h(3);
    assert_loopback_matches(&b.build(), &[1, 2, 4]);
}

#[test]
fn loopback_local_only_matches_across_ranks() {
    relax_min_local_qubits();
    let mut b = CircuitBuilder::new(5);
    b.h(0).cx(0, 1).cx(1, 2);
    assert_loopback_matches(&b.build(), &[1, 2, 4]);
}

// With 4 qubits across 4 ranks, qubits 0,1 are local and 2,3 are global, so a
// two qubit gate can place operands in every local and global combination.

fn spread_4q() -> CircuitBuilder {
    let mut b = CircuitBuilder::new(4);
    b.h(0).h(1).h(2).h(3);
    b
}

#[test]
fn loopback_cx_all_qubit_splits() {
    relax_min_local_qubits();
    for &(c, t) in &[(0usize, 1usize), (1, 2), (2, 0), (2, 3)] {
        let mut b = spread_4q();
        b.cx(c, t);
        assert_loopback_matches(&b.build(), &[1, 2, 4]);
    }
}

#[test]
fn loopback_cz_all_qubit_splits() {
    relax_min_local_qubits();
    for &(a, t) in &[(0usize, 1usize), (1, 3), (3, 0), (2, 3)] {
        let mut b = spread_4q();
        b.cz(a, t);
        assert_loopback_matches(&b.build(), &[1, 2, 4]);
    }
}

#[test]
fn loopback_swap_all_qubit_splits() {
    relax_min_local_qubits();
    for &(a, t) in &[(0usize, 1usize), (1, 2), (3, 0), (2, 3)] {
        let mut b = spread_4q();
        b.swap(a, t);
        assert_loopback_matches(&b.build(), &[1, 2, 4]);
    }
}

#[test]
fn loopback_swap_asymmetric_state() {
    relax_min_local_qubits();
    // Distinct rotation per qubit, so every marginal differs and a wrong
    // readout permutation cannot hide behind symmetric probabilities. A
    // uniform H wall would mask map bugs.
    let mut b = CircuitBuilder::new(4);
    b.rx(0.3, 0).rx(0.7, 1).rx(1.1, 2).rx(1.5, 3);
    b.swap(0, 3).swap(1, 2).swap(0, 1);
    b.ry(0.4, 3).cx(3, 0);
    assert_loopback_matches(&b.build(), &[1, 2, 4]);
}

#[test]
fn loopback_gates_after_relabel_use_moved_qubits() {
    relax_min_local_qubits();
    // Repeated non-diagonal gates on global qubits force relabels and
    // evictions; later gates must follow the moved qubits through the map.
    let n = 5;
    let mut b = CircuitBuilder::new(n);
    b.rx(0.2, 0).rx(0.5, 1).rx(0.9, 2);
    b.h(3).h(4);
    b.cx(3, 0).cz(4, 1).rzz(0.6, 2, 3);
    b.h(0).h(1);
    b.swap(2, 4).ry(0.8, 4).cx(4, 2);
    assert_loopback_matches(&b.build(), &[1, 2, 4]);
}

#[test]
fn loopback_export_statevector_after_swap() {
    relax_min_local_qubits();
    // Export must reorder the gathered amplitudes back to circuit qubit order
    // after SWAPs leave the map permuted.
    let n = 4;
    let mut b = CircuitBuilder::new(n);
    b.rx(0.3, 0).ry(0.8, 1).rx(1.2, 2).t(3).h(3);
    b.swap(0, 3).swap(1, 2);
    b.rz(0.5, 3).ry(0.2, 0);
    let circuit = b.build();

    let mut sv = StatevectorBackend::new(SEED);
    sv.init(circuit.num_qubits, circuit.num_classical_bits)
        .unwrap();
    sv.apply_instructions(&circuit.instructions).unwrap();
    let expected = sv.export_statevector().unwrap();

    let size = 4;
    let shared = LoopbackShared::new(size);
    let handles: Vec<_> = (0..size)
        .map(|rank| {
            let comm = LoopbackComm {
                shared: shared.clone(),
                rank,
            };
            let circuit = circuit.clone();
            std::thread::Builder::new()
                .stack_size(64 * 1024 * 1024)
                .spawn(move || {
                    let ctx = DistributedContext::from_comm(Arc::new(comm));
                    let mut backend = DistributedStatevectorBackend::new(ctx, SEED);
                    backend
                        .init(circuit.num_qubits, circuit.num_classical_bits)
                        .unwrap();
                    backend.apply_instructions(&circuit.instructions).unwrap();
                    backend.export_statevector().unwrap()
                })
                .expect("spawn rank thread")
        })
        .collect();
    let results: Vec<Vec<Complex64>> = handles.into_iter().map(|h| h.join().unwrap()).collect();
    for actual in &results {
        assert_eq!(expected.len(), actual.len());
        for (i, (e, a)) in expected.iter().zip(actual.iter()).enumerate() {
            assert!((e - a).norm() < TOL, "amp[{i}] expected {e}, got {a}");
        }
    }
}

#[test]
fn loopback_measure_after_swap_reads_moved_qubit() {
    relax_min_local_qubits();
    // |0001> swapped across the boundary becomes |1000>: the top qubit must
    // read 1 and the bottom 0 through the permuted map.
    let n = 4;
    let mut b = CircuitBuilder::new_with_classical(n, 2);
    b.x(0).swap(0, n - 1);
    b.measure(n - 1, 0).measure(0, 1);
    let circuit = b.build();
    for &size in &[1usize, 2, 4] {
        let (probs, bits) = loopback_run(&circuit, size);
        assert!(bits[0], "size {size}: moved excitation must read 1");
        assert!(!bits[1], "size {size}: vacated qubit must read 0");
        let expected = 1usize << (n - 1);
        assert!(
            (probs[expected] - 1.0).abs() < TOL,
            "size {size}: state must be |1000>, got p[{expected}]={}",
            probs[expected]
        );
    }
}

#[test]
fn loopback_reset_after_swap_clears_moved_qubit() {
    relax_min_local_qubits();
    let n = 4;
    let mut b = CircuitBuilder::new(n);
    b.rx(0.4, 0).rx(0.9, 1).h(3);
    b.swap(0, 3);
    let mut circuit = b.build();
    circuit.add_reset(3);
    assert_loopback_matches(&circuit, &[1, 2, 4]);
}

#[test]
fn loopback_relabel_victim_starvation_falls_back() {
    relax_min_local_qubits();
    // At 4 ranks of a 4 qubit register only two positions are local. An Mcu
    // referencing every qubit leaves no eviction victim, so the direct global
    // exchange path must apply the gate.
    let x = [
        [Complex64::new(0.0, 0.0), Complex64::new(1.0, 0.0)],
        [Complex64::new(1.0, 0.0), Complex64::new(0.0, 0.0)],
    ];
    let mut b = spread_4q();
    b.mcu(x, &[0, 1, 2], 3);
    assert_loopback_matches(&b.build(), &[1, 2, 4]);
}

#[test]
fn loopback_rzz_all_qubit_splits() {
    relax_min_local_qubits();
    for &(a, t) in &[(0usize, 1usize), (1, 2), (2, 0), (2, 3)] {
        let mut b = spread_4q();
        b.rzz(0.85, a, t);
        assert_loopback_matches(&b.build(), &[1, 2, 4]);
    }
}

#[test]
fn loopback_cphase_all_qubit_splits() {
    relax_min_local_qubits();
    for &(c, t) in &[(0usize, 1usize), (1, 2), (2, 0), (2, 3)] {
        let mut b = spread_4q();
        b.cphase(0.6, c, t);
        assert_loopback_matches(&b.build(), &[1, 2, 4]);
    }
}

#[test]
fn loopback_controlled_unitary_global_target() {
    relax_min_local_qubits();
    let ry = |theta: f64| {
        let (s, c) = (theta / 2.0).sin_cos();
        [
            [Complex64::new(c, 0.0), Complex64::new(-s, 0.0)],
            [Complex64::new(s, 0.0), Complex64::new(c, 0.0)],
        ]
    };
    for &(c, t) in &[(1usize, 2usize), (2, 0), (2, 3)] {
        let mut b = spread_4q();
        b.cu(ry(0.9), c, t);
        assert_loopback_matches(&b.build(), &[1, 2, 4]);
    }
}

#[test]
fn loopback_toffoli_mixed_splits() {
    relax_min_local_qubits();
    let x = [
        [Complex64::new(0.0, 0.0), Complex64::new(1.0, 0.0)],
        [Complex64::new(1.0, 0.0), Complex64::new(0.0, 0.0)],
    ];
    for &(c0, c1, t) in &[(0usize, 1usize, 2usize), (0, 2, 3), (2, 3, 0), (1, 2, 3)] {
        let mut b = spread_4q();
        b.mcu(x, &[c0, c1], t);
        assert_loopback_matches(&b.build(), &[1, 2, 4]);
    }
}

#[test]
fn loopback_mixed_circuit_qft_like() {
    relax_min_local_qubits();
    // H walls interleaved with controlled phases spanning all splits, plus an
    // entangling tail.
    let mut b = CircuitBuilder::new(4);
    b.h(0).h(1).h(2).h(3);
    b.cphase(0.5, 0, 2)
        .cphase(0.25, 1, 3)
        .cx(2, 1)
        .rzz(0.4, 0, 3)
        .swap(1, 2)
        .cz(0, 3);
    assert_loopback_matches(&b.build(), &[1, 2, 4]);
}

// These circuits are large enough to trigger the fusion pipeline (>= 10 qubits
// for 1q fusion, >= 12 for 2q, >= 14 for multi 1q, >= 16 for diagonal batch).
// The reference statevector fuses identically, so a match confirms the
// distributed backend decomposes each fused or batched variant correctly.

#[test]
fn loopback_fused_multifused_and_2q() {
    relax_min_local_qubits();
    // HEA pattern: rotation layers fuse into MultiFused, CX ladders into
    // Fused2q and Multi2q, reaching the global qubits at the top.
    let n = 14;
    let mut b = CircuitBuilder::new(n);
    for layer in 0..3 {
        for q in 0..n {
            b.ry(0.3 + 0.01 * (layer * n + q) as f64, q);
            b.rz(-0.2 + 0.02 * q as f64, q);
        }
        for q in 0..n - 1 {
            b.cx(q, q + 1);
        }
    }
    assert_loopback_matches(&b.build(), &[1, 2, 4]);
}

#[test]
fn loopback_fused_batchphase_qft() {
    relax_min_local_qubits();
    // Textbook QFT produces H walls plus controlled phase batches (BatchPhase)
    // and trailing swaps, spanning the full register including global qubits.
    let n = 12;
    let mut b = CircuitBuilder::new(n);
    for q in 0..n {
        b.h(q);
        for (j, target) in (q + 1..n).enumerate() {
            let angle = std::f64::consts::PI / (1u64 << (j + 1)) as f64;
            b.cphase(angle, target, q);
        }
    }
    for q in 0..n / 2 {
        b.swap(q, n - 1 - q);
    }
    assert_loopback_matches(&b.build(), &[1, 2, 4]);
}

#[test]
fn loopback_fused_batchrzz_qaoa() {
    relax_min_local_qubits();
    // QAOA pattern: Rzz on every edge (fuse into BatchRzz at >= 16q) plus Rx
    // mixers (MultiFused), spanning global qubits.
    let n = 16;
    let mut b = CircuitBuilder::new(n);
    for q in 0..n {
        b.h(q);
    }
    for round in 0..2 {
        for q in 0..n - 1 {
            b.rzz(0.7 + 0.01 * round as f64, q, q + 1);
        }
        for q in 0..n {
            b.rx(0.4, q);
        }
    }
    assert_loopback_matches(&b.build(), &[1, 2, 4]);
}

#[test]
fn loopback_fused_diagonal_batch() {
    relax_min_local_qubits();
    // Mixed diagonal run (cphase + rzz + diagonal 1q) at >= 16 qubits triggers
    // DiagonalBatch fusion; ensure the entries decompose across the boundary.
    let n = 16;
    let mut b = CircuitBuilder::new(n);
    for q in 0..n {
        b.h(q);
    }
    for q in 0..n - 1 {
        b.cphase(0.3, q, q + 1);
        b.rzz(0.5, q, q + 1);
    }
    for q in 0..n {
        b.t(q);
        b.rz(0.15, q);
    }
    assert_loopback_matches(&b.build(), &[1, 2, 4]);
}

#[test]
fn loopback_fused_2q_both_global() {
    relax_min_local_qubits();
    // Adjacent 1q gates and a CX on the top pair fuse into a Fused2q spanning
    // both global qubits at 4 ranks, exercising the four rank gather path.
    let n = 12;
    let mut b = CircuitBuilder::new(n);
    for q in 0..n {
        b.h(q);
    }
    b.ry(0.6, n - 2)
        .rz(0.3, n - 1)
        .cx(n - 2, n - 1)
        .ry(-0.4, n - 1);
    assert_loopback_matches(&b.build(), &[1, 2, 4]);
}

#[test]
fn loopback_measure_deterministic_basis_state() {
    relax_min_local_qubits();
    // X on every qubit yields |1...1>, so every measurement must read 1.
    let n = 5;
    let mut b = CircuitBuilder::new_with_classical(n, n);
    for q in 0..n {
        b.x(q);
    }
    for q in 0..n {
        b.measure(q, q);
    }
    let circuit = b.build();
    for &size in &[1usize, 2, 4] {
        let (_probs, bits) = loopback_run(&circuit, size);
        assert_eq!(bits, vec![true; n], "size {size}: expected all one readout");
    }
}

#[test]
fn loopback_measure_determinism_across_ranks() {
    relax_min_local_qubits();
    // With a fixed seed the outcomes and post measurement probabilities must be
    // identical for every rank count (lockstep meas_rng plus Allreduce).
    let n = 6;
    let mut b = CircuitBuilder::new_with_classical(n, n);
    b.h(0).h(3).h(5);
    for q in 0..n - 1 {
        b.cx(q, q + 1);
    }
    b.measure(0, 0).measure(3, 3).measure(5, 5);
    assert_loopback_deterministic(&b.build(), &[1, 2, 4]);
}

#[test]
fn loopback_ghz_measure_correlated() {
    relax_min_local_qubits();
    // A GHZ state collapses to all zeros or all ones, so the measured qubits
    // must agree at every rank count.
    let n = 5;
    let mut b = CircuitBuilder::new_with_classical(n, 2);
    b.h(0);
    for q in 0..n - 1 {
        b.cx(q, q + 1);
    }
    b.measure(0, 0).measure(n - 1, 1);
    let circuit = b.build();
    for &size in &[1usize, 2, 4] {
        let (_probs, bits) = loopback_run(&circuit, size);
        assert_eq!(bits[0], bits[1], "size {size}: GHZ qubits must correlate");
    }
}

#[test]
fn loopback_reset_clears_global_qubit() {
    relax_min_local_qubits();
    // Resetting every qubit of a full superposition must yield |0...0>.
    let n = 5;
    let mut b = CircuitBuilder::new(n);
    for q in 0..n {
        b.h(q);
    }
    let mut circuit = b.build();
    for q in 0..n {
        circuit.add_reset(q);
    }
    for &size in &[1usize, 2, 4] {
        let (probs, _bits) = loopback_run(&circuit, size);
        assert!(
            (probs[0] - 1.0).abs() < TOL,
            "size {size}: reset should yield |0...0>, got p[0]={}",
            probs[0]
        );
    }
}

#[test]
fn loopback_reset_empty_zero_branch_matches_statevector() {
    relax_min_local_qubits();
    // Statevector reset projects onto the |0> branch. If that branch is empty,
    // it reinitializes to |0...0>; distributed reset must not preserve other
    // qubits by flipping the |1> branch.
    let n = 5;
    let mut b = CircuitBuilder::new(n);
    b.x(0).x(n - 1);
    let mut circuit = b.build();
    circuit.add_reset(n - 1);

    let expected = reference_probs(&circuit);
    assert!(
        (expected[0] - 1.0).abs() < TOL,
        "statevector reset should reinitialize empty zero branch"
    );
    assert_loopback_matches(&circuit, &[1, 2, 4]);
}

#[test]
fn loopback_conditional_on_global_measurement() {
    relax_min_local_qubits();
    // Measure a |1> qubit into a bit, then conditionally X another qubit. The
    // conditional must fire identically on every rank.
    let n = 4;
    let mut b = CircuitBuilder::new_with_classical(n, 1);
    b.x(n - 1);
    b.measure(n - 1, 0);
    b.conditional(
        crate::circuit::ClassicalCondition::BitIsOne(0),
        crate::gates::Gate::X,
        &[0],
    );
    let circuit = b.build();
    for &size in &[1usize, 2, 4] {
        let (probs, bits) = loopback_run(&circuit, size);
        assert!(bits[0], "size {size}: measured bit should be 1");
        let expected = 1usize | (1usize << (n - 1));
        assert!(
            (probs[expected] - 1.0).abs() < TOL,
            "size {size}: conditional X should set qubit 0"
        );
    }
}

#[test]
fn loopback_tiled_exchange_matches_full() {
    relax_min_local_qubits();
    // Every chunk size must match the statevector reference, so tiling the
    // exchange preserves correctness.
    let n = 6;
    let mut b = CircuitBuilder::new(n);
    for q in 0..n {
        b.h(q);
    }
    b.rx(0.5, n - 1).ry(0.8, n - 2).h(n - 1).h(n - 2);
    let circuit = b.build();
    let expected = reference_probs(&circuit);

    // The local slice at 4 ranks is 16 amplitudes; these chunks span the tiling
    // boundaries from one element up to a single whole slice message. Both the
    // direct exchange and the relabel exchange honor the chunk size.
    for &relabel in &[true, false] {
        for &size in &[1usize, 2, 4] {
            for &chunk in &[1usize, 3, 16, 1 << 20] {
                let actual = loopback_probs_with(&circuit, size, chunk, relabel);
                assert_eq!(expected.len(), actual.len());
                for (i, (e, a)) in expected.iter().zip(actual.iter()).enumerate() {
                    assert!(
                        (e - a).abs() < TOL,
                        "size {size} chunk {chunk} relabel {relabel}: prob[{i}] expected {e}, got {a}"
                    );
                }
            }
        }
    }
}

#[test]
fn exchange_counters_track_communication() {
    relax_min_local_qubits();
    // A single rank keeps every qubit local, so no gate exchanges.
    let ctx = DistributedContext::from_comm(Arc::new(LoopbackComm {
        shared: LoopbackShared::new(1),
        rank: 0,
    }));
    let mut dist = DistributedStatevectorBackend::new(ctx, SEED);
    dist.init(4, 0).unwrap();
    dist.apply(&inst_h(3)).unwrap();
    assert_eq!(dist.exchange_messages(), 0, "single rank never exchanges");

    // At 2 ranks qubit 3 is global and the local slice is 8 amplitudes. The
    // diagonal Z and Rz are free. Direct mode exchanges the full slice per H.
    // Relabel mode pays one half-slice exchange for the first H, after which
    // qubit 3 is local and the second H is free.
    let circuit = {
        let mut b = CircuitBuilder::new(4);
        b.z(3).rz(0.3, 3).h(3).h(3);
        b.build()
    };
    let direct = loopback_exchange_stats(&circuit, 2, usize::MAX, false);
    assert_eq!(
        direct,
        (2, 16),
        "two global H gates exchange the slice twice"
    );
    let relabeled = loopback_exchange_stats(&circuit, 2, usize::MAX, true);
    assert_eq!(relabeled, (1, 4), "one relabel moves half the slice once");
}

#[test]
fn relabel_makes_global_swap_free() {
    relax_min_local_qubits();
    let circuit = {
        let mut b = CircuitBuilder::new(5);
        b.x(0).swap(0, 4);
        b.build()
    };
    let direct = loopback_exchange_stats(&circuit, 2, usize::MAX, false);
    let relabeled = loopback_exchange_stats(&circuit, 2, usize::MAX, true);
    assert_eq!(direct.0, 1, "direct boundary swap exchanges once");
    assert_eq!(relabeled, (0, 0), "relabel swap is a map update");
}

#[test]
fn relabel_reduces_phased_exchange_volume() {
    relax_min_local_qubits();
    // Activity shifts from the bottom qubits to the top qubits, as after
    // fusion passes that concentrate work. Direct exchange pays for every
    // layer touching the global qubits; relabeling pays two half-slice moves
    // when the working set shifts and the remaining layers run locally.
    let n = 6;
    let mut b = CircuitBuilder::new(n);
    for q in 0..4 {
        b.ry(0.3 + 0.01 * q as f64, q);
    }
    b.cx(0, 1).cx(2, 3);
    for layer in 0..3 {
        for q in 2..n {
            b.ry(0.2 + 0.01 * (layer * n + q) as f64, q);
        }
        b.cx(2, 3).cx(3, 4).cx(4, 5);
    }
    let circuit = b.build();
    let direct = loopback_exchange_stats(&circuit, 4, usize::MAX, false);
    let relabeled = loopback_exchange_stats(&circuit, 4, usize::MAX, true);
    assert_eq!(
        relabeled,
        (2, 16),
        "two half-slice relabels cover every layer"
    );
    assert!(
        relabeled.1 < direct.1 / 4,
        "relabel volume {} should be far below direct volume {}",
        relabeled.1,
        direct.1
    );
}

#[test]
fn tiled_exchange_splits_messages_not_volume() {
    relax_min_local_qubits();
    let circuit = {
        let mut b = CircuitBuilder::new(5);
        b.h(4);
        b.build()
    };
    let full = loopback_exchange_stats(&circuit, 2, 1 << 20, false);
    let tiled = loopback_exchange_stats(&circuit, 2, 4, false);
    assert_eq!(full.0, 1, "full slice is one message");
    assert_eq!(tiled.0, 4, "16 amplitudes in chunks of 4 is four messages");
    assert_eq!(full.1, tiled.1, "total amplitudes exchanged is unchanged");

    let relabel_full = loopback_exchange_stats(&circuit, 2, 1 << 20, true);
    let relabel_tiled = loopback_exchange_stats(&circuit, 2, 4, true);
    assert_eq!(
        relabel_full,
        (1, 8),
        "relabel moves the half slice in one message"
    );
    assert_eq!(
        relabel_tiled,
        (2, 8),
        "8 amplitudes in chunks of 4 is two messages"
    );
}

fn inst_h(q: usize) -> crate::circuit::Instruction {
    crate::circuit::Instruction::Gate {
        gate: crate::gates::Gate::H,
        targets: crate::circuit::smallvec![q],
    }
}

/// Run `circuit` across `size` ranks with the given exchange chunk and relabel
/// setting; return rank 0's `(message_count, amplitude_count)`.
fn loopback_exchange_stats(
    circuit: &Circuit,
    size: usize,
    chunk: usize,
    relabel: bool,
) -> (u64, u64) {
    let shared = LoopbackShared::new(size);
    let handles: Vec<_> = (0..size)
        .map(|rank| {
            let comm = LoopbackComm {
                shared: shared.clone(),
                rank,
            };
            let circuit = circuit.clone();
            std::thread::Builder::new()
                .stack_size(64 * 1024 * 1024)
                .spawn(move || {
                    let ctx = DistributedContext::from_comm(Arc::new(comm));
                    let mut backend = DistributedStatevectorBackend::new(ctx, SEED);
                    backend.set_exchange_chunk(chunk);
                    backend.set_relabel(relabel);
                    backend
                        .init(circuit.num_qubits, circuit.num_classical_bits)
                        .unwrap();
                    backend.apply_instructions(&circuit.instructions).unwrap();
                    (backend.exchange_messages(), backend.exchange_amplitudes())
                })
                .expect("spawn rank thread")
        })
        .collect();
    let mut results: Vec<(u64, u64)> = handles.into_iter().map(|h| h.join().unwrap()).collect();
    let first = results.swap_remove(0);
    results.clear();
    first
}

#[test]
fn reports_global_qubit_count_for_single_rank() {
    let ctx = DistributedContext::serial();
    let mut dist = DistributedStatevectorBackend::new(ctx, SEED);
    dist.init(5, 0).unwrap();
    assert_eq!(dist.num_qubits(), 5);
    assert!(dist.supports_fused_gates());
}