//! A circuit as an operator, grown from its middle outward: the method the
//! peaked circuits are read by (`tracker_peaked`, `mpo`).

#![allow(dead_code)]

use wai_quantum::gates::{C, Gate};
use wai_quantum::quantum_mpsim::{Chain, MpsConfig};
use wai_quantum::quantum_spd::{Op, Pauli, RotCircuit, sin_cos};

/// A gate's matrix on the qubits it names (the first the most significant bit).
pub fn matrix(op: &Op) -> (Vec<u32>, Vec<C>) {
    let c = |re: f64, im: f64| C { re, im };
    let pauli = |p: Pauli| -> [C; 4] {
        match p {
            Pauli::X => [c(0.0, 0.0), c(1.0, 0.0), c(1.0, 0.0), c(0.0, 0.0)],
            Pauli::Y => [c(0.0, 0.0), c(0.0, -1.0), c(0.0, 1.0), c(0.0, 0.0)],
            Pauli::Z => [c(1.0, 0.0), c(0.0, 0.0), c(0.0, 0.0), c(-1.0, 0.0)],
        }
    };
    let kron = |a: &[C], b: &[C], da: usize, db: usize| -> Vec<C> {
        let d = da * db;
        let mut out = vec![c(0.0, 0.0); d * d];
        for i in 0..d {
            for j in 0..d {
                let (x, y) = (a[(i / db) * da + j / db], b[(i % db) * db + j % db]);
                out[i * d + j] = c(x.re * y.re - x.im * y.im, x.re * y.im + x.im * y.re);
            }
        }
        out
    };
    // exp(−iθP/2) = cos(θ/2)·I − i·sin(θ/2)·P, for the string P on `qs`.
    let rotation = |axis: &[(u32, Pauli)], cos: f64, sin: f64| -> (Vec<u32>, Vec<C>) {
        let qs: Vec<u32> = axis.iter().map(|(q, _)| *q).collect();
        let mut p = pauli(axis[0].1).to_vec();
        let mut d = 2;
        for &(_, f) in &axis[1..] {
            p = kron(&p, &pauli(f), d, 2);
            d *= 2;
        }
        let d = 1usize << qs.len();
        let m = (0..d * d).map(|k| {
            let id = if k / d == k % d { cos } else { 0.0 };
            // −i·sin·P: (re, im) of P times (0, −sin).
            c(id + sin * p[k].im, -sin * p[k].re)
        });
        (qs, m.collect())
    };
    let h = std::f64::consts::FRAC_1_SQRT_2;
    match op {
        Op::Rot { axis, theta } => {
            let (s, co) = sin_cos(0.5 * theta);
            rotation(axis, co, s)
        }
        Op::QuarterRot { axis, k } => {
            let (co, s) = match k {
                1 => (h, h),
                2 => (0.0, 1.0),
                _ => (-h, h),
            };
            rotation(axis, co, s)
        }
        Op::CZ(a, b) => (vec![*a, *b], (0..16).map(|k| if k % 5 == 0 { c(if k == 15 { -1.0 } else { 1.0 }, 0.0) } else { c(0.0, 0.0) }).collect()),
        Op::CX(a, b) => {
            let perm = [0usize, 1, 3, 2];
            (vec![*a, *b], (0..16).map(|k| c(if perm[k / 4] == k % 4 { 1.0 } else { 0.0 }, 0.0)).collect())
        }
        Op::H(q) => (vec![*q], vec![c(h, 0.0), c(h, 0.0), c(h, 0.0), c(-h, 0.0)]),
        Op::X(q) => (vec![*q], pauli(Pauli::X).to_vec()),
        Op::Y(q) => (vec![*q], pauli(Pauli::Y).to_vec()),
        Op::Z(q) => (vec![*q], pauli(Pauli::Z).to_vec()),
        Op::S(q) => (vec![*q], vec![c(1.0, 0.0), c(0.0, 0.0), c(0.0, 0.0), c(0.0, 1.0)]),
        Op::Sdg(q) => (vec![*q], vec![c(1.0, 0.0), c(0.0, 0.0), c(0.0, 0.0), c(0.0, -1.0)]),
        Op::SX(q) => (vec![*q], vec![c(0.5, 0.5), c(0.5, -0.5), c(0.5, -0.5), c(0.5, 0.5)]),
        Op::SXdg(q) => (vec![*q], vec![c(0.5, -0.5), c(0.5, 0.5), c(0.5, 0.5), c(0.5, -0.5)]),
        Op::Swap(a, b) => (vec![*a, *b], Gate::swap(0, 1).matrix),
    }
}

/// The transpose of a square row-major matrix.
pub fn transposed(m: &[C]) -> Vec<C> {
    let d = (m.len() as f64).sqrt() as usize;
    (0..d * d).map(|k| m[(k % d) * d + k / d]).collect()
}

/// What growing the operator from the middle gives.
pub struct Grown {
    /// Each qubit's `⟨Z⟩` in the state the circuit prepares from `|0…0⟩`.
    pub z: Vec<f64>,
    /// The string the signs of `z` give, bit `q` qubit `q`.
    pub string: Vec<u8>,
    /// The string's probability.
    pub probability: f64,
    /// The chain's qubit order.
    pub order: Vec<u32>,
    pub max_bond: usize,
    pub discarded: f64,
    pub fidelity_estimate: f64,
}

/// The circuit as an operator, a matrix-product state on `2n` sites (each
/// qubit's row index beside its column index, the qubits in breadth-first
/// order of `edges`), started as the identity (a Bell pair a qubit) and grown
/// from the middle of the circuit outward: the gates after the midpoint on the
/// rows, in order, and each gate before it transposed on the columns, in
/// reverse, one two-qubit gate a side at a time. That is
/// `(A ⊗ Bᵀ)|Φ⟩ = (A·B ⊗ I)|Φ⟩` for the circuit `A·B`. The columns are then
/// projected on `|0⟩`, which leaves `2^{−n/2}·A·B|0…0⟩` on the rows.
pub fn grow_from_middle(circuit: &RotCircuit, edges: &[(u32, u32)], chi: usize) -> Grown {
    grow_from(circuit, edges, chi, circuit.ops.len() / 2, usize::MAX).expect("the whole circuit")
}

/// A circuit's operator on a chain, as [`grow_chain`] leaves it.
pub struct Chained {
    pub chain: Chain,
    /// The qubit in each slot of the chain at the end, and each qubit's slot.
    pub order: Vec<u32>,
    pub pos: Vec<u32>,
    /// Whether every gate was applied.
    pub whole: bool,
}

impl Chained {
    /// The site of qubit `q`'s row index, and of its column index.
    pub fn row(&self, q: u32) -> u32 {
        2 * self.pos[q as usize]
    }
    pub fn col(&self, q: u32) -> u32 {
        2 * self.pos[q as usize] + 1
    }
}

/// The circuit's operator `(A ⊗ Bᵀ)|Φ⟩`, `A` the gates from `mid` on and `B`
/// those before it, grown outward from `mid` by at most `window` two-qubit
/// gates a side, at bond dimension `chi`.
///
/// Each qubit's row and column sit side by side in a slot of the chain, the
/// slots first in breadth-first order of `edges`. A two-qubit gate on qubits in
/// slots apart first moves the second qubit's slot, both its sites together,
/// next to the first's, and the slots stay where they are moved: a qubit's row
/// never leaves its column, so routing stretches no pair across the chain.
pub fn grow_chain(circuit: &RotCircuit, edges: &[(u32, u32)], chi: usize, mid: usize, window: usize) -> Chained {
    let n = circuit.n;
    let order = super::common::bfs_chain(n, edges);
    let mut pos = vec![0u32; n as usize];
    for (i, &q) in order.iter().enumerate() {
        pos[q as usize] = i as u32;
    }
    let mut at = order.clone();
    let swap = Gate::swap(0, 1).matrix;
    let ops = &circuit.ops;
    let two = |op: &Op| matrix(op).0.len() == 2;
    let mut chain = Chain::new(2 * n, MpsConfig { max_bond: chi, cutoff: 1e-14 }).expect("a chain");
    for q in 0..n {
        let s = 2 * pos[q as usize];
        chain.apply(&Gate { qubits: vec![s], matrix: matrix(&Op::H(0)).1 }).expect("a gate");
        chain.apply(&Gate { qubits: vec![s, s + 1], matrix: matrix(&Op::CX(0, 1)).1 }).expect("a gate");
    }
    // Exchange the pairs in slots `s` and `s + 1`: sites (r, c, r', c') become
    // (r', c', r, c) by four swaps of neighbours.
    let swap_slots = |chain: &mut Chain, pos: &mut [u32], at: &mut [u32], s: u32| {
        let b = 2 * s;
        for i in [b + 1, b, b + 2, b + 1] {
            chain.apply(&Gate { qubits: vec![i, i + 1], matrix: swap.clone() }).expect("a swap");
        }
        let (x, y) = (at[s as usize], at[s as usize + 1]);
        at.swap(s as usize, s as usize + 1);
        pos[x as usize] = s + 1;
        pos[y as usize] = s;
    };
    let apply = |chain: &mut Chain, pos: &mut [u32], at: &mut [u32], op: &Op, columns: bool| {
        let (qs, m) = matrix(op);
        if let [a, b] = qs[..] {
            while pos[b as usize] > pos[a as usize] + 1 {
                swap_slots(chain, pos, at, pos[b as usize] - 1);
            }
            while pos[b as usize] + 1 < pos[a as usize] {
                swap_slots(chain, pos, at, pos[b as usize]);
            }
        }
        let site = |q: u32| 2 * pos[q as usize] + u32::from(columns);
        let m = if columns { transposed(&m) } else { m };
        chain.apply(&Gate { qubits: qs.iter().map(|&q| site(q)).collect(), matrix: m }).expect("a gate");
    };
    let (mut f, mut b) = (mid, mid);
    let mut sides = 0usize;
    while (f < ops.len() || b > 0) && sides < window {
        sides += 1;
        while f < ops.len() {
            apply(&mut chain, &mut pos, &mut at, &ops[f], false);
            f += 1;
            if two(&ops[f - 1]) {
                break;
            }
        }
        while b > 0 {
            b -= 1;
            apply(&mut chain, &mut pos, &mut at, &ops[b], true);
            if two(&ops[b]) {
                break;
            }
        }
    }
    let whole = f == ops.len() && b == 0;
    Chained { chain, order: at, pos, whole }
}

/// The echo `2⁻ⁿ Tr(O · W† O W)` of a circuit `W` for `O` the `Z`s on `sites`,
/// from its operator grown outward from gate `mid`: with `|W⟩⟩ = (W ⊗ I)|Φ⟩`,
/// it is `⟨⟨W|(O ⊗ O)|W⟩⟩`, since `O` is real and diagonal. For a circuit
/// written `U · V · U†`, `mid` at `V`'s centre starts the operator at `V`, a
/// product near the identity, and grows it by `V`'s dressing alone. Returns the
/// echo and the chain's report.
pub fn echo_by_operator(circuit: &RotCircuit, edges: &[(u32, u32)], chi: usize, mid: usize, sites: &[u32]) -> (f64, wai_quantum::quantum_mpsim::MpsReport) {
    let g = grow_chain(circuit, edges, chi, mid, usize::MAX);
    let mut z: Vec<(u32, char)> = Vec::new();
    for &q in sites {
        z.push((g.row(q), 'Z'));
        z.push((g.col(q), 'Z'));
    }
    let f = g.chain.expectation(&z).expect("sites on the chain") / g.chain.norm2();
    (f, g.chain.report())
}

/// The operator grown outward from gate `mid` by at most `window` two-qubit
/// gates a side: with the whole circuit, what [`grow_from_middle`] reads;
/// with less, the chain's report alone (no readout), to find where the
/// circuit's halves undo each other.
pub fn grow_from(circuit: &RotCircuit, edges: &[(u32, u32)], chi: usize, mid: usize, window: usize) -> Option<Grown> {
    let n = circuit.n;
    let Chained { mut chain, order, pos, whole } = grow_chain(circuit, edges, chi, mid, window);
    let (row, col) = (|q: u32| 2 * pos[q as usize], |q: u32| 2 * pos[q as usize] + 1);
    let report = chain.report();
    if !whole {
        return Some(Grown {
            z: Vec::new(),
            string: Vec::new(),
            probability: f64::NAN,
            order,
            max_bond: report.max_bond_used,
            discarded: report.discarded,
            fidelity_estimate: report.fidelity_estimate,
        });
    }
    let zero = C { re: 0.0, im: 0.0 };
    let project = vec![C { re: 1.0, im: 0.0 }, zero, zero, zero];
    for q in 0..n {
        chain.apply(&Gate { qubits: vec![col(q)], matrix: project.clone() }).expect("a projection");
    }
    let norm = chain.norm2();
    let z: Vec<f64> = (0..n).map(|q| chain.expectation(&[(row(q), 'Z')]).expect("a qubit") / norm).collect();
    let string: Vec<u8> = z.iter().map(|&v| u8::from(v < 0.0)).collect();
    let mut bits = vec![0u8; 2 * n as usize];
    for q in 0..n {
        bits[row(q) as usize] = string[q as usize];
    }
    let a = chain.amplitude(&bits).expect("a string");
    Some(Grown {
        z,
        string,
        probability: (a.re * a.re + a.im * a.im) / norm,
        order,
        max_bond: report.max_bond_used,
        discarded: report.discarded,
        fidelity_estimate: report.fidelity_estimate,
    })
}