//! The tracker's guided sparse ground-state problem: its Hamiltonian, a list of
//! (Pauli string, coefficient) pairs, and its guiding basis state, each as a
//! `.npy` array file; read here, and written as the operator sum DMRG takes.

#![allow(dead_code)]

use wai_quantum::quantum_dmrg::{Mps, Op, OpSum};

/// One term: the factors as (qubit, letter), and the coefficient.
pub type Term = (Vec<(usize, char)>, f64);

/// The header of a version 1 or 2 `.npy` file, and the bytes after it.
fn npy(bytes: &[u8]) -> (String, &[u8]) {
    assert!(bytes.len() > 12 && &bytes[..6] == b"\x93NUMPY", "not a .npy file"); // allow-name: the .npy format's magic bytes
    let (len, at) = match bytes[6] {
        1 => (u16::from_le_bytes([bytes[8], bytes[9]]) as usize, 10),
        2 | 3 => (u32::from_le_bytes([bytes[8], bytes[9], bytes[10], bytes[11]]) as usize, 12),
        v => panic!(".npy version {v}"),
    };
    let header = String::from_utf8(bytes[at..at + len].to_vec()).expect("an ASCII header");
    (header, &bytes[at + len..])
}

/// The Hamiltonian file: a `(rows, 2)` array of `<U64` strings, each row a Pauli
/// string in the tracker's order (its last letter is qubit 0) and a complex
/// coefficient written `(re+imj)`. Refused unless every imaginary part is 0.
pub fn hamiltonian(bytes: &[u8]) -> (usize, Vec<Term>) {
    let (header, body) = npy(bytes);
    assert!(header.contains("'descr': '<U64'") && header.contains("'fortran_order': False"), "unexpected array: {header}");
    let shape = header.split("'shape': (").nth(1).and_then(|s| s.split(')').next()).expect("a shape");
    let dims: Vec<usize> = shape.split(',').filter(|s| !s.trim().is_empty()).map(|s| s.trim().parse().expect("a dimension")).collect();
    assert!(dims.len() == 2 && dims[1] == 2, "shape ({shape})");
    let cell = |i: usize| -> String {
        body[i * 256..(i + 1) * 256].chunks(4).map(|c| u32::from_le_bytes([c[0], c[1], c[2], c[3]])).take_while(|&u| u != 0).map(|u| char::from_u32(u).expect("a character")).collect()
    };
    let mut n = 0;
    let terms = (0..dims[0])
        .map(|r| {
            let (p, c) = (cell(2 * r), cell(2 * r + 1));
            n = p.len();
            let factors = p.chars().rev().enumerate().filter(|(_, l)| *l != 'I').collect::<Vec<_>>();
            assert!(factors.iter().all(|(_, l)| matches!(l, 'X' | 'Y' | 'Z')), "a letter in {p}");
            (factors, coefficient(&c))
        })
        .collect();
    (n, terms)
}

/// `(re+imj)` with `im` zero, as the double `re` reads as.
fn coefficient(s: &str) -> f64 {
    let s = s.trim().trim_start_matches('(').trim_end_matches(')');
    let body = s.strip_suffix('j').expect("a complex number");
    // The imaginary part starts at the last sign that does not follow an exponent.
    let cut = body.char_indices().rev().find(|&(i, ch)| (ch == '+' || ch == '-') && i > 0 && !matches!(body.as_bytes()[i - 1], b'e' | b'E')).map(|(i, _)| i).expect("a real and an imaginary part");
    let (re, im) = body.split_at(cut);
    let im: f64 = im.parse().expect("an imaginary part");
    assert_eq!(im, 0.0, "a coefficient with an imaginary part: {s}");
    re.parse().expect("a real part")
}

/// The guiding state file: a `<i8` scalar, the basis state's index, bit `q`
/// qubit `q`.
pub fn guide(bytes: &[u8]) -> u64 {
    let (header, body) = npy(bytes);
    assert!(header.contains("'descr': '<i8'") && header.contains("'shape': ()"), "unexpected array: {header}");
    let v = i64::from_le_bytes(body[..8].try_into().unwrap());
    assert!(v >= 0, "a negative index");
    v as u64
}

/// The Hamiltonian as DMRG takes it, its qubits placed on the chain by
/// `order` (`order[site]` is the qubit there). `Y` is `−i·J` with `J` real, so
/// a term with an even number of `Y`s is real: `(−1)^{k/2}` times the product
/// with each `Y` replaced by `J`. A term with an odd number is refused.
pub fn op_sum(n: usize, terms: &[Term], order: &[usize]) -> OpSum {
    let site = |q: usize| order.iter().position(|&s| s == q).expect("every qubit on the chain");
    let j = Op([0.0, 1.0, -1.0, 0.0]);
    let mut h = OpSum::new(n);
    for (factors, c) in terms {
        let ys = factors.iter().filter(|(_, l)| *l == 'Y').count();
        assert!(ys.is_multiple_of(2), "a term with an odd number of Y is not real");
        let sign = if ys.is_multiple_of(4) { 1.0 } else { -1.0 };
        let ops: Vec<(usize, Op)> = factors
            .iter()
            .map(|&(q, l)| (site(q), match l {
                'X' => Op::x(),
                'Y' => j,
                _ => Op::z(),
            }))
            .collect();
        h.add(sign * c, &ops);
    }
    h
}

/// The basis state `index` (bit `q` qubit `q`, `1` the `−1` eigenstate of `Z`)
/// as a product MPS on the chain `order`.
pub fn basis_mps(index: u64, order: &[usize]) -> Mps {
    let sites = order.iter().map(|&q| if (index >> q) & 1 == 1 { vec![0.0, 1.0] } else { vec![1.0, 0.0] }).collect();
    Mps { dims: vec![1; order.len() + 1], sites }
}

/// `⟨index| H |index⟩`: the diagonal terms' signs on the basis state.
pub fn diagonal_energy(terms: &[Term], index: u64) -> f64 {
    terms
        .iter()
        .filter(|(f, _)| f.iter().all(|(_, l)| *l == 'Z'))
        .map(|(f, c)| if f.iter().filter(|(q, _)| (index >> q) & 1 == 1).count().is_multiple_of(2) { *c } else { -*c })
        .sum()
}

/// A chain order for the qubits: breadth first through the coupling graph the
/// two-qubit terms draw, from qubit 0, each vertex's neighbours in increasing
/// order (a bandwidth-reducing order, so the MPO's bonds stay small).
pub fn bfs_order(n: usize, terms: &[Term]) -> Vec<usize> {
    let mut adj = vec![Vec::new(); n];
    for (f, _) in terms {
        if let [(a, _), (b, _)] = f.as_slice() {
            adj[*a].push(*b);
            adj[*b].push(*a);
        }
    }
    for a in &mut adj {
        a.sort_unstable();
        a.dedup();
    }
    let mut seen = vec![false; n];
    let mut order = Vec::with_capacity(n);
    for start in 0..n {
        if seen[start] {
            continue;
        }
        seen[start] = true;
        let mut queue = std::collections::VecDeque::from([start]);
        while let Some(v) = queue.pop_front() {
            order.push(v);
            for &w in &adj[v] {
                if !seen[w] {
                    seen[w] = true;
                    queue.push_back(w);
                }
            }
        }
    }
    order
}