wai-quantum 0.3.20

//! Deterministic neutral-atom REARRANGEMENT — `wai.quantum.rearrange`
//! (extensions/quantum-ops § Register preparation).
//!
//! A neutral-atom quantum processor loads atoms **stochastically** into a large
//! array of optical traps — each site fills with probability ~½ — then, before the
//! computation, mobile optical tweezers **rearrange** the loaded atoms into a
//! defect-free target register. Which atom moves to which site is a combinatorial
//! optimization: the **linear sum assignment problem** (assign every target site a
//! distinct atom, minimizing total move distance), solved exactly by the
//! **Hungarian algorithm** in `O(n³)`. Getting it right — every target filled, from
//! the atoms actually present, at minimum cost — is the difference between a
//! defect-free register and a failed shot, and today no planner signs its plan.
//!
//! This module is the WAI take: a deterministic Hungarian solver over integer
//! Manhattan costs (byte-identical on every machine), turning a low-yield load into
//! a defect-free array, sealing the plan into a signed
//! [`crate::quantum_ops::RearrangeReceipt`]. The `success` verdict — every target
//! site covered by a valid bijection of present atoms — is re-checkable from the
//! bound configs alone; the same check rejects a tampered plan.
//!
//! Honest boundary: a geometric/combinatorial model of the rearrangement (the
//! solver and its optimality are exact and reproducible); the physics of tweezer
//! transport, loss, and AOD parallelization is out of scope. The move distance is
//! the plan's cost, not a wall-clock or fidelity claim.

use crate::quantum_ops::{content_hash, GrantRef, RearrangeReceipt};
use ed25519_dalek::SigningKey;

fn splitmix64(state: &mut u64) -> u64 {
    *state = state.wrapping_add(0x9E37_79B9_7F4A_7C15);
    let mut z = *state;
    z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
    z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
    z ^ (z >> 31)
}

// ===========================================================================
// Trap array
// ===========================================================================

/// A `w × h` array of optical traps; site `(x, y)` has index `y*w + x`.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Array {
    pub w: u32,
    pub h: u32,
    /// Occupancy per site (row-major).
    pub occ: Vec<bool>,
}

impl Array {
    pub fn empty(w: u32, h: u32) -> Array {
        Array { w, h, occ: vec![false; (w * h) as usize] }
    }
    /// Stochastically load each site with fill probability `fill_permille`/1000.
    pub fn load(w: u32, h: u32, fill_permille: u32, seed: u64) -> Array {
        let mut st = seed.wrapping_mul(0xA24B_AED4).wrapping_add(1);
        let occ = (0..w * h)
            .map(|_| (splitmix64(&mut st) % 1000) < fill_permille as u64)
            .collect();
        Array { w, h, occ }
    }
    pub fn atoms(&self) -> Vec<usize> {
        (0..self.occ.len()).filter(|&i| self.occ[i]).collect()
    }
    pub fn count(&self) -> u32 {
        self.occ.iter().filter(|&&o| o).count() as u32
    }
    fn pos(&self, idx: usize) -> (i64, i64) {
        ((idx as u32 % self.w) as i64, (idx as u32 / self.w) as i64)
    }
    /// Canonical bytes of the occupancy (for content-hashing).
    pub fn bytes(&self) -> Vec<u8> {
        let mut b = Vec::with_capacity(self.occ.len() / 8 + 12);
        b.extend_from_slice(b"wai:qc-array\x01");
        b.extend_from_slice(&self.w.to_le_bytes());
        b.extend_from_slice(&self.h.to_le_bytes());
        for chunk in self.occ.chunks(8) {
            let mut byte = 0u8;
            for (i, &o) in chunk.iter().enumerate() {
                if o {
                    byte |= 1 << i;
                }
            }
            b.push(byte);
        }
        b
    }
}

/// A centered `tw × th` block of target sites (the defect-free register).
pub fn target_block(w: u32, h: u32, tw: u32, th: u32) -> Vec<usize> {
    let x0 = (w - tw) / 2;
    let y0 = (h - th) / 2;
    let mut t = Vec::with_capacity((tw * th) as usize);
    for y in y0..y0 + th {
        for x in x0..x0 + tw {
            t.push((y * w + x) as usize);
        }
    }
    t
}

/// Canonical bytes of a target site set.
pub fn target_bytes(target: &[usize], w: u32, h: u32) -> Vec<u8> {
    let mut b = Vec::with_capacity(target.len() * 4 + 16);
    b.extend_from_slice(b"wai:qc-target\x01");
    b.extend_from_slice(&w.to_le_bytes());
    b.extend_from_slice(&h.to_le_bytes());
    let mut sorted = target.to_vec();
    sorted.sort_unstable();
    for &s in &sorted {
        b.extend_from_slice(&(s as u32).to_le_bytes());
    }
    b
}

// ===========================================================================
// Hungarian algorithm (min-cost linear sum assignment)
// ===========================================================================

/// Optimal assignment of `n` rows to distinct columns of an `n × m` cost matrix
/// (`n ≤ m`), minimizing total cost. `O(n²m)`. Returns `assignment[row] = column`.
/// The classic Kuhn–Munkres with potentials (deterministic).
fn hungarian(cost: &[Vec<i64>], n: usize, m: usize) -> Vec<usize> {
    const INF: i64 = i64::MAX / 4;
    let mut u = vec![0i64; n + 1];
    let mut v = vec![0i64; m + 1];
    let mut p = vec![0usize; m + 1]; // p[col] = row matched to col (1-indexed; 0 = none)
    let mut way = vec![0usize; m + 1];
    for i in 1..=n {
        p[0] = i;
        let mut j0 = 0usize;
        let mut minv = vec![INF; m + 1];
        let mut used = vec![false; m + 1];
        loop {
            used[j0] = true;
            let i0 = p[j0];
            let mut delta = INF;
            let mut j1 = 0usize;
            for j in 1..=m {
                if !used[j] {
                    let cur = cost[i0 - 1][j - 1] - u[i0] - v[j];
                    if cur < minv[j] {
                        minv[j] = cur;
                        way[j] = j0;
                    }
                    if minv[j] < delta {
                        delta = minv[j];
                        j1 = j;
                    }
                }
            }
            for j in 0..=m {
                if used[j] {
                    u[p[j]] += delta;
                    v[j] -= delta;
                } else {
                    minv[j] -= delta;
                }
            }
            j0 = j1;
            if p[j0] == 0 {
                break;
            }
        }
        loop {
            let j1 = way[j0];
            p[j0] = p[j1];
            j0 = j1;
            if j0 == 0 {
                break;
            }
        }
    }
    let mut assignment = vec![usize::MAX; n];
    for j in 1..=m {
        if p[j] != 0 {
            assignment[p[j] - 1] = j - 1;
        }
    }
    assignment
}

// ===========================================================================
// Planning + verification
// ===========================================================================

/// A rearrangement plan: `moves[k] = (from_site, to_site)`.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Plan {
    pub moves: Vec<(usize, usize)>,
    pub total_distance: u64,
    pub atoms_moved: u32,
    /// Every target site is covered by a distinct present atom.
    pub feasible: bool,
}

impl Plan {
    pub fn bytes(&self) -> Vec<u8> {
        let mut b = Vec::with_capacity(self.moves.len() * 8 + 8);
        b.extend_from_slice(b"wai:qc-plan\x01");
        for &(f, t) in &self.moves {
            b.extend_from_slice(&(f as u32).to_le_bytes());
            b.extend_from_slice(&(t as u32).to_le_bytes());
        }
        b
    }
    pub fn plan_hash(&self) -> [u8; 32] {
        content_hash(&self.bytes())
    }
}

/// Plan the rearrangement of `array` into `target` by minimum-distance assignment.
/// If fewer atoms are loaded than target sites, the plan is infeasible.
pub fn solve(array: &Array, target: &[usize]) -> Plan {
    let atoms = array.atoms();
    let nt = target.len();
    let na = atoms.len();
    if na < nt {
        return Plan { moves: Vec::new(), total_distance: 0, atoms_moved: 0, feasible: false };
    }
    // cost[target_i][atom_j] = Manhattan distance
    let cost: Vec<Vec<i64>> = target
        .iter()
        .map(|&t| {
            let (tx, ty) = array.pos(t);
            atoms
                .iter()
                .map(|&a| {
                    let (ax, ay) = array.pos(a);
                    (tx - ax).abs() + (ty - ay).abs()
                })
                .collect()
        })
        .collect();
    let assign = hungarian(&cost, nt, na);
    let mut moves = Vec::new();
    let mut total = 0u64;
    let mut moved = 0u32;
    for (i, &t) in target.iter().enumerate() {
        let a = atoms[assign[i]];
        let d = cost[i][assign[i]] as u64;
        total += d;
        if a != t {
            moves.push((a, t));
            moved += 1;
        }
    }
    Plan { moves, total_distance: total, atoms_moved: moved, feasible: true }
}

/// Does the plan turn `array` into a configuration covering every `target` site?
/// Order-independent set check: sources must be present atoms, targets a bijection,
/// and the resulting occupancy must include all target sites.
pub fn verify(array: &Array, target: &[usize], plan: &Plan) -> bool {
    if !plan.feasible {
        return false;
    }
    let mut occ = array.occ.clone();
    let mut froms = std::collections::HashSet::new();
    let mut tos = std::collections::HashSet::new();
    for &(f, t) in &plan.moves {
        if f >= occ.len() || t >= occ.len() || !array.occ[f] {
            return false; // must move an atom that is actually present
        }
        if !froms.insert(f) || !tos.insert(t) {
            return false; // no atom or site used twice
        }
    }
    // apply as a set: vacate sources, fill targets
    for &(f, _) in &plan.moves {
        occ[f] = false;
    }
    for &(_, t) in &plan.moves {
        occ[t] = true;
    }
    target.iter().all(|&t| occ[t])
}

/// Solve + verify + seal a `RearrangeReceipt` binding the initial array, target,
/// and plan to the outcome.
#[allow(clippy::too_many_arguments)]
pub fn rearrange_and_seal(
    signer: &SigningKey,
    signer_id: &str,
    array: &Array,
    target: &[usize],
    joules_micro: u64,
    grant: GrantRef,
) -> (Plan, bool, RearrangeReceipt) {
    let plan = solve(array, target);
    let success = verify(array, target, &plan);
    let receipt = RearrangeReceipt::seal(
        signer,
        signer_id,
        "hungarian-lsap",
        content_hash(&array.bytes()),
        content_hash(&target_bytes(target, array.w, array.h)),
        plan.plan_hash(),
        success,
        plan.atoms_moved,
        plan.total_distance,
        joules_micro,
        grant,
        None,
    );
    (plan, success, receipt)
}

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

    fn key(s: u8) -> SigningKey {
        SigningKey::from_bytes(&[s; 32])
    }

    // brute-force min assignment cost for small n, to check Hungarian optimality
    fn brute(cost: &[Vec<i64>], n: usize, m: usize) -> i64 {
        let mut cols: Vec<usize> = (0..m).collect();
        let mut best = i64::MAX;
        // permutations of choosing n distinct columns in order
        fn rec(cost: &[Vec<i64>], n: usize, i: usize, used: &mut Vec<bool>, acc: i64, best: &mut i64) {
            if i == n {
                *best = (*best).min(acc);
                return;
            }
            for j in 0..cost[i].len() {
                if !used[j] {
                    used[j] = true;
                    rec(cost, n, i + 1, used, acc + cost[i][j], best);
                    used[j] = false;
                }
            }
        }
        let mut used = vec![false; m];
        rec(cost, n, 0, &mut used, 0, &mut best);
        let _ = &mut cols;
        best
    }

    #[test]
    fn hungarian_matches_brute_force() {
        let mut st = 12345u64;
        for _ in 0..30 {
            let n = 3 + (splitmix64(&mut st) % 3) as usize;
            let m = n + (splitmix64(&mut st) % 3) as usize;
            let cost: Vec<Vec<i64>> = (0..n)
                .map(|_| (0..m).map(|_| (splitmix64(&mut st) % 20) as i64).collect())
                .collect();
            let assign = hungarian(&cost, n, m);
            let got: i64 = (0..n).map(|i| cost[i][assign[i]]).sum();
            assert_eq!(got, brute(&cost, n, m), "Hungarian must be optimal");
            // valid: distinct columns
            let mut s = std::collections::HashSet::new();
            assert!(assign.iter().all(|&c| s.insert(c)));
        }
    }

    #[test]
    fn rearrangement_makes_defect_free() {
        // load a half-full array, rearrange into a centered block, verify it fills.
        let arr = Array::load(12, 12, 550, 7);
        let target = target_block(12, 12, 6, 6);
        assert!(arr.count() as usize >= target.len(), "enough atoms for the test");
        let plan = solve(&arr, &target);
        assert!(plan.feasible);
        assert!(verify(&arr, &target, &plan), "target must be defect-free after the plan");
    }

    #[test]
    fn already_placed_atoms_do_not_move() {
        // an array already defect-free over the target needs zero moves.
        let mut arr = Array::empty(8, 8);
        let target = target_block(8, 8, 4, 4);
        for &t in &target {
            arr.occ[t] = true;
        }
        let plan = solve(&arr, &target);
        assert!(plan.feasible);
        assert_eq!(plan.atoms_moved, 0);
        assert_eq!(plan.total_distance, 0);
        assert!(verify(&arr, &target, &plan));
    }

    #[test]
    fn too_few_atoms_is_infeasible() {
        let mut arr = Array::empty(8, 8);
        arr.occ[0] = true; // one atom
        let target = target_block(8, 8, 3, 3); // 9 sites
        let plan = solve(&arr, &target);
        assert!(!plan.feasible);
        assert!(!verify(&arr, &target, &plan));
    }

    #[test]
    fn tamper_breaks_verification() {
        let arr = Array::load(10, 10, 600, 3);
        let target = target_block(10, 10, 5, 5);
        let mut plan = solve(&arr, &target);
        assert!(verify(&arr, &target, &plan));
        // corrupt a move's destination to a non-target site → target no longer covered
        if let Some(mv) = plan.moves.first_mut() {
            mv.1 = 0; // divert to corner (not in the centered target)
        }
        // (only asserts when the diverted site actually left a target uncovered)
        if !target.contains(&0) {
            assert!(!verify(&arr, &target, &plan), "a diverted move must fail verification");
        }
    }

    #[test]
    fn solve_is_deterministic() {
        let arr = Array::load(14, 14, 520, 99);
        let target = target_block(14, 14, 7, 7);
        assert_eq!(solve(&arr, &target), solve(&arr, &target));
    }

    #[test]
    fn seals_verifying_receipt() {
        let arr = Array::load(12, 12, 560, 42);
        let target = target_block(12, 12, 6, 6);
        let (plan, success, rec) = rearrange_and_seal(
            &key(1), "did:key:lab", &arr, &target, 400_000, GrantRef::unbounded("quantum.rearrange"),
        );
        assert!(success);
        assert!(rec.verify());
        assert!(rec.success);
        assert!(rec.plan_matches(&plan.bytes()));
    }
}