//! The reference modules (extensions/quantum-wasmq.md §8): the engine's qLDPC decoder and
//! its atom rearrangement planner, built as wasmQ modules of record, give exactly what the
//! same code gives natively, called directly and in the middle of a circuit, on every engine.

use wai_quantum::quantum_atom::{Array, target_block, target_bytes};
use wai_quantum::quantum_qec::CssCode;
use wai_wasmq::engine::{Engine, Interpreter};
use wai_wasmq::envelope::{Run, open, pack};
use wai_wasmq::interp::Metering;
use wai_wasmq::reference::{self, checks_prior, target_prior};
use wai_wasmq::run::{EngineKind, run, shot_seed};
use wai_wasmq::{abi, meter};
use wai_quantum::ed25519_dalek::SigningKey;

const DECODER: &[u8] = include_bytes!("../../quantum-conformance/wasmq/modules/qldpc_decoder.wasm");
const PLANNER: &[u8] = include_bytes!("../../quantum-conformance/wasmq/modules/atom_planner.wasm");

fn engines(module: &[u8], prepared: &'static [u8]) -> Vec<Box<dyn Engine>> {
    let module: &'static [u8] = Box::leak(module.to_vec().into_boxed_slice());
    #[allow(unused_mut)] // the native engine is added where this build has it
    let mut v: Vec<Box<dyn Engine>> = vec![Box::new(Interpreter::new(module, Metering::Own).unwrap()), Box::new(Interpreter::new(prepared, Metering::Module).unwrap())];
    #[cfg(all(feature = "wasmtime", not(target_arch = "wasm32")))]
    v.push(Box::new(wai_wasmq::engine::Native::new(prepared).unwrap()));
    v
}

fn bits(v: u64, n: usize) -> Vec<bool> {
    (0..n).map(|i| v >> i & 1 == 1).collect()
}

struct Rng(u64);
impl Rng {
    fn next(&mut self) -> u64 {
        self.0 = self.0.wrapping_add(0x9e37_79b9_7f4a_7c15);
        let mut z = self.0;
        z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
        z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
        z ^ (z >> 31)
    }
}

#[test]
fn the_priors_are_the_engines_own_bytes() {
    for code in [CssCode::surface(3), CssCode::surface(5), CssCode::gross()] {
        assert_eq!(checks_prior(code.n, &code.checks), code.check_matrix_bytes());
    }
    assert_eq!(target_prior(8, 8, &target_block(8, 8, 4, 4)), target_bytes(&target_block(8, 8, 4, 4), 8, 8));
    assert_eq!(reference::target_block(8, 8, 4, 4), target_block(8, 8, 4, 4));
}

fn decoder_instance<'a>(e: &'a dyn Engine, code: &CssCode, seed: u64) -> Box<dyn wai_wasmq::engine::Instance + 'a> {
    let mut i = e.instantiate().unwrap();
    i.set_fuel(1 << 40);
    abi::lifecycle(i.as_mut(), "wasmq-seed", &[seed]).unwrap();
    abi::prior(i.as_mut(), 0, &checks_prior(code.n, &code.checks)).unwrap();
    i
}

#[test]
fn the_decoder_module_decodes_as_the_engine_does_on_every_engine() {
    let prepared: &'static [u8] = Box::leak(meter::prepare(DECODER).unwrap().into_boxed_slice());
    let small = CssCode::surface(3);
    let decl = wai_quantum::qasm::ExternDecl { name: "decode64".into(), args: vec![wai_quantum::qasm::ExternArg::Bit(small.checks.len() as u32)], ret: Some(small.n as u32) };
    let seed = 11;
    // Every error of weight at most two on surface(3).
    let mut errors = vec![vec![]];
    for a in 0..small.n {
        errors.push(vec![a]);
        for b in a + 1..small.n {
            errors.push(vec![a, b]);
        }
    }
    let mut fuels: Vec<Vec<u64>> = Vec::new();
    for e in engines(DECODER, prepared) {
        let mut i = decoder_instance(e.as_ref(), &small, seed);
        let mut f = Vec::new();
        for err in &errors {
            let mut x = vec![false; small.n];
            for &q in err {
                x[q] = true;
            }
            let syn = small.syndrome(&x);
            i.set_fuel(1 << 30);
            let got = abi::call(i.as_mut(), &decl, "decode64", &syn, false).unwrap();
            f.push((1u64 << 30) - i.fuel() as u64);
            assert_eq!(got, wasmq_qldpc_decoder::decode_bits(&small, &syn, seed), "{} on {err:?}", e.name());
            if err.len() <= 1 {
                let residual: Vec<usize> = (0..small.n).filter(|&q| got[q] != x[q]).collect();
                assert!(small.syndrome(&(0..small.n).map(|q| got[q] != x[q]).collect::<Vec<_>>()).iter().all(|b| !b));
                assert!(small.is_trivial(&residual), "a weight-1 error decoded to a logical");
            }
        }
        fuels.push(f);
    }
    assert!(fuels.windows(2).all(|w| w[0] == w[1]), "each call's fuel is the same on every engine");

    // The gross code through the list form, with random errors of weight up to four.
    let gross = CssCode::gross();
    let decl = wai_quantum::qasm::ExternDecl { name: "decode".into(), args: vec![wai_quantum::qasm::ExternArg::Bit(gross.checks.len() as u32)], ret: Some(gross.n as u32) };
    let mut rng = Rng(5);
    let cases: Vec<Vec<bool>> = (0..12)
        .map(|_| {
            let mut x = vec![false; gross.n];
            for _ in 0..1 + rng.next() % 4 {
                x[(rng.next() % gross.n as u64) as usize] = true;
            }
            gross.syndrome(&x)
        })
        .collect();
    let mut fuels: Vec<Vec<u64>> = Vec::new();
    for e in engines(DECODER, prepared) {
        let mut i = decoder_instance(e.as_ref(), &gross, seed);
        let mut f = Vec::new();
        for syn in &cases {
            i.set_fuel(1 << 34);
            let got = abi::call(i.as_mut(), &decl, "decode", syn, false).unwrap();
            f.push((1u64 << 34) - i.fuel() as u64);
            assert_eq!(got, wasmq_qldpc_decoder::decode_bits(&gross, syn, seed), "{}", e.name());
        }
        fuels.push(f);
    }
    assert!(fuels.windows(2).all(|w| w[0] == w[1]));
}

#[test]
fn the_planner_module_plans_as_the_engine_does_on_every_engine() {
    let prepared: &'static [u8] = Box::leak(meter::prepare(PLANNER).unwrap().into_boxed_slice());
    for (w, h, tw, th, f) in [(4u32, 4u32, 2u32, 2u32, "plan64"), (5, 5, 3, 3, "plan64"), (9, 9, 5, 5, "plan")] {
        let n = (w * h) as usize;
        let target = target_block(w, h, tw, th);
        let decl = wai_quantum::qasm::ExternDecl { name: f.into(), args: vec![wai_quantum::qasm::ExternArg::Bit(n as u32)], ret: Some(2 * n as u32 + 1) };
        let mut rng = Rng(w as u64);
        let loads: Vec<Vec<bool>> = (0..16).map(|k| Array::load(w, h, 300 + 40 * k, rng.next()).occ).collect();
        let mut fuels: Vec<Vec<u64>> = Vec::new();
        for e in engines(PLANNER, prepared) {
            let mut i = e.instantiate().unwrap();
            i.set_fuel(1 << 40);
            abi::prior(i.as_mut(), 0, &target_prior(w, h, &target)).unwrap();
            let mut fu = Vec::new();
            for occ in &loads {
                i.set_fuel(1 << 34);
                let got = abi::call(i.as_mut(), &decl, f, occ, false).unwrap();
                fu.push((1u64 << 34) - i.fuel() as u64);
                assert_eq!(got, wasmq_atom_planner::plan_bits(w, h, &target, occ), "{} {w}x{h}", e.name());
                // The engine's own check of the plan, on what the module returned.
                let plan = wai_quantum::quantum_atom::solve(&Array { w, h, occ: occ.clone() }, &target);
                assert_eq!(got[2 * n], plan.feasible && wai_quantum::quantum_atom::verify(&Array { w, h, occ: occ.clone() }, &target, &plan));
            }
            fuels.push(fu);
        }
        assert!(fuels.windows(2).all(|w| w[0] == w[1]));
    }
}

fn kinds() -> Vec<EngineKind> {
    #[allow(unused_mut)] // the native engine is added where this build has it
    let mut k = vec![EngineKind::Interpreter, EngineKind::InterpreterPrepared];
    #[cfg(all(feature = "wasmtime", not(target_arch = "wasm32")))]
    k.insert(0, EngineKind::Native);
    k
}

#[test]
fn a_memory_round_decodes_in_the_middle_of_the_circuit_and_every_shot_ends_in_the_code() {
    let key = SigningKey::from_bytes(&[7; 32]);
    for (code, errors, shots, state) in [(CssCode::surface(3), 1, 48, "shot"), (CssCode::surface(5), 2, 24, "job"), (CssCode::gross(), 3, 6, "shot")] {
        let mut r = Run::new(shots, 1234);
        r.state = state.into();
        let env = pack(&reference::qec_bundle(DECODER, code.n, &code.checks, errors, r)).unwrap();
        let o = open(&env).unwrap();
        let mut bodies = Vec::new();
        for k in kinds() {
            let (body, log) = run(&o, k, &key, true).unwrap();
            let log = log.unwrap();
            assert_eq!(body["result"]["status"], "ok", "{}: {}", code.id, body["result"]);
            // Every call's result is the native decoder's on the same syndrome.
            for c in &log {
                let syn: Vec<bool> = (0..code.checks.len()).map(|i| c.input[i / 8] >> (i % 8) & 1 == 1).collect();
                let seed = if state == "shot" { shot_seed(1234, c.shot) } else { 1234 };
                let want = abi::pack(&wasmq_qldpc_decoder::decode_bits(&code, &syn, seed));
                assert_eq!(c.output, want, "shot {} of {}", c.shot, code.id);
            }
            assert_eq!(log.len() as u64, shots);
            // The record replays: every call's output and fuel again, from its input alone.
            assert!(wai_wasmq::run::replay(&o, &log, k).unwrap().is_empty());
            let mut bad = log.clone();
            bad[0].output[0] ^= 1;
            bad[1].fuel += 1;
            assert_eq!(wai_wasmq::run::replay(&o, &bad, k).unwrap().len(), 2, "this test can fail");
            bodies.push(wai_quantum::engine::receipt::canonical(&body));
        }
        assert!(bodies.windows(2).all(|w| w[0] == w[1]), "{}: the same job bytes on every engine", code.id);
        let body: serde_json::Value = serde_json::from_slice(&bodies[0]).unwrap();
        // Every shot: the second syndrome is clear and the residual is no logical.
        let (n, m) = (code.n, code.checks.len());
        for (key, count) in body["result"]["counts"].as_object().unwrap() {
            let bit = |i: usize| key.as_bytes()[key.len() - 1 - i] == b'1';
            let k = key.len();
            // Bits in declaration order: e*, s, corr, s2, fin; fin are the last n declared.
            let fin: Vec<bool> = (k - n..k).map(bit).collect();
            let s2: Vec<bool> = (k - n - m..k - n).map(bit).collect();
            assert!(s2.iter().all(|b| !b), "{}: the syndrome after correction ({count} shots)", code.id);
            let residual: Vec<usize> = (0..n).filter(|&q| fin[q]).collect();
            assert!(code.syndrome(&fin).iter().all(|b| !b));
            assert!(code.is_trivial(&residual), "{}: a shot ended in a logical error", code.id);
        }
        assert!(wai_wasmq::verify(&body, Some(&o), kinds()[0]).unwrap().ok());
    }
}

#[test]
fn atoms_are_rearranged_in_the_middle_of_the_circuit_and_fill_the_target() {
    let key = SigningKey::from_bytes(&[7; 32]);
    for (w, h, tw, th, shots) in [(4u32, 4u32, 2u32, 2u32, 64u64), (9, 9, 3, 3, 12)] {
        let env = pack(&reference::atom_bundle(PLANNER, w, h, tw, th, Run::new(shots, 99))).unwrap();
        let o = open(&env).unwrap();
        let target = target_block(w, h, tw, th);
        let n = (w * h) as usize;
        let mut bodies = Vec::new();
        for k in kinds() {
            let (body, log) = run(&o, k, &key, true).unwrap();
            for c in log.unwrap() {
                let occ: Vec<bool> = (0..n).map(|i| c.input[i / 8] >> (i % 8) & 1 == 1).collect();
                assert_eq!(c.output, abi::pack(&wasmq_atom_planner::plan_bits(w, h, &target, &occ)));
            }
            bodies.push(wai_quantum::engine::receipt::canonical(&body));
        }
        assert!(bodies.windows(2).all(|w| w[0] == w[1]));
        let body: serde_json::Value = serde_json::from_slice(&bodies[0]).unwrap();
        let mut feasible = 0;
        for (key, count) in body["result"]["counts"].as_object().unwrap() {
            let bit = |i: usize| key.as_bytes()[key.len() - 1 - i] == b'1';
            let k = key.len();
            // Declared: occ, plan, fin.
            let plan_ok = bit(n + 2 * n);
            let fin: Vec<bool> = (k - n..k).map(bit).collect();
            let occ: Vec<bool> = (0..n).map(bit).collect();
            if plan_ok {
                feasible += count.as_u64().unwrap();
                assert!(target.iter().all(|&t| fin[t]), "a feasible plan left a target site empty");
                assert_eq!(fin.iter().filter(|&&b| b).count(), occ.iter().filter(|&&b| b).count(), "atoms are moved, never made or lost");
            }
        }
        assert!(feasible > 0);
        let _ = bits(0, 0);
    }
}