//! The operator Loschmidt echo `f = 2⁻ⁿ Tr(O W† O W)` of a circuit written
//! `W = T · M · T⁻¹` (tail, middle, head), as a tensor network on the coupling
//! graph with an ancilla on every qubit: the purification of `W` as an
//! operator, `|W⟩⟩ = (W ⊗ I)|Φ⟩`, `|Φ⟩` a Bell pair on each qubit and its
//! ancilla. After each ancilla's `Z` is folded onto its qubit by a `CX`, `f` is
//! `⟨Π_{q ∈ O} Z_q⟩` in that state (see `common::purified_echo`).
//!
//! The network is evolved in the circuit's order (head, middle, tail) by
//! simple update with the dense blocked linear algebra
//! (`Tns::with_blocked_linalg`), and read three ways: along the arc of the
//! loop through the observable's sites with the simple-update environment
//! (the bond weights), the same with belief propagation's environment, and
//! around the whole loop with belief propagation's environment on every bond
//! off it. Without the middle the exact echo is 1, which gives the divisor of
//! the published "rescaled" values.

#![allow(dead_code)]

use wai_quantum::quantum_bptn::{BpEnv, Tns};
use wai_quantum::quantum_spd::{Op, Pauli, RotCircuit};

use crate::common;

/// The purified echo cut into its parts, on `2n` qubits (ancilla `n + q` of
/// qubit `q`).
pub struct Purified {
    pub n: u32,
    pub edges: Vec<(u32, u32)>,
    /// The Bell pairs.
    pub prep: RotCircuit,
    /// `T⁻¹`, on the qubits.
    pub head: RotCircuit,
    /// `M`, on the qubits.
    pub middle: RotCircuit,
    /// `T`, then the folds of the ancillas' `Z`s.
    pub tail: RotCircuit,
    /// The observable's qubits.
    pub sites: Vec<u32>,
    /// The loop through them.
    pub ring: Vec<u32>,
    /// The loop's arc from the first site through the others to the last.
    pub path: Vec<u32>,
}

/// Cut `circuit` (head of `head_len` ops, middle of `middle_len`, the rest the
/// tail) and purify it for the `Z`s on `sites`.
pub fn purified(circuit: &RotCircuit, head_len: usize, middle_len: usize, sites: &[u32]) -> Purified {
    let n = circuit.n;
    let mut edges = common::coupling(circuit);
    edges.extend((0..n).map(|q| (q, n + q)));
    let mut prep = RotCircuit::new(2 * n);
    for q in 0..n {
        prep.push(Op::H(q)).push(Op::CX(q, n + q));
    }
    let part = |ops: &[Op]| RotCircuit { n: 2 * n, ops: ops.to_vec() };
    let head = part(&circuit.ops[..head_len]);
    let middle = part(&circuit.ops[head_len..head_len + middle_len]);
    let mut tail = part(&circuit.ops[head_len + middle_len..]);
    for &q in sites {
        tail.push(Op::CX(n + q, q));
    }
    let ring = common::loop_through(&common::coupling(circuit), sites).expect("the observable's qubits lie on one loop");
    // The ring starts at the first site and meets the others in order.
    let last = ring.iter().position(|&q| q == *sites.last().expect("a site")).expect("on the loop");
    let path = ring[..=last].to_vec();
    Purified { n, edges, prep, head, middle, tail, sites: sites.to_vec(), ring, path }
}

/// What a network reads, and what it took.
#[derive(Clone, Debug)]
pub struct Reading {
    /// Along the path, the simple-update environment.
    pub vidal_path: f64,
    /// Along the path, belief propagation's environment.
    pub bp_path: Option<f64>,
    /// Around the loop, belief propagation's environment.
    pub bp_loop: Option<f64>,
    pub bp_sweeps: usize,
    pub bp_residual: f64,
    pub max_bond: usize,
    /// The loop's bond dimensions, in its order from its first vertex.
    pub loop_bonds: Vec<usize>,
    pub discarded: f64,
    pub seconds: f64,
}

/// Read `Π_{q ∈ sites} Z_q` from `t`: along the path with the simple-update
/// environment; with `bp_sweeps > 0` also with belief propagation's, and
/// around the loop when the largest bond is at most `loop_max`.
pub fn read(t: &Tns, p: &Purified, bp_sweeps: usize, bp_tol: f64, loop_max: usize) -> Reading {
    let started = std::time::Instant::now();
    let ops: Vec<(u32, Pauli)> = p.sites.iter().map(|&q| (q, Pauli::Z)).collect();
    let vidal_path = t.expect_path_bp(&t.vidal_env(), &p.path, &ops).expect("the path");
    let (mut bp_path, mut bp_loop, mut sweeps, mut residual) = (None, None, 0, f64::NAN);
    if bp_sweeps > 0 {
        let env: BpEnv = t.belief_propagation(bp_sweeps, bp_tol);
        sweeps = env.sweeps;
        residual = env.residual;
        bp_path = Some(t.expect_path_bp(&env, &p.path, &ops).expect("the path"));
        if t.max_bond() <= loop_max {
            bp_loop = Some(t.expect_loop_bp(&env, &p.ring, &ops).expect("the loop"));
        }
    }
    let weights = t.bond_weights();
    let loop_bonds = (0..p.ring.len())
        .map(|i| {
            let (a, b) = (p.ring[i], p.ring[(i + 1) % p.ring.len()]);
            let e = p.edges.iter().position(|&(x, y)| (x, y) == (a.min(b), a.max(b)) || (x, y) == (a, b) || (x, y) == (b, a)).expect("a loop edge");
            weights[e].len()
        })
        .collect();
    Reading {
        vidal_path,
        bp_path,
        bp_loop,
        bp_sweeps: sweeps,
        bp_residual: residual,
        max_bond: t.max_bond(),
        loop_bonds,
        discarded: t.discarded(),
        seconds: started.elapsed().as_secs_f64(),
    }
}

/// One region of the loop cluster expansion, read.
#[derive(Clone, Debug)]
pub struct RegionRead {
    /// The region's qubits (network vertices), ascending.
    pub region: Vec<u32>,
    /// How many of the graph's rings it joins (the observable's counted).
    pub cells: usize,
    pub cuts: Vec<(u32, u32, usize)>,
    pub combinations: u64,
    /// The contraction's estimated work (complex multiply-adds).
    pub work: f64,
    /// The most bytes its tables hold at once (estimated).
    pub bytes: f64,
    /// `None` where its work or memory was past the budget.
    pub value: Option<f64>,
    pub norm_ratio: Option<f64>,
    pub seconds: f64,
}

/// One order of the expansion: every union of the observable's ring with at
/// most `cells − 1` other rings as the largest regions.
#[derive(Clone, Debug)]
pub struct Order {
    pub cells: usize,
    /// (region, counting number), as `Tns::loop_cluster_terms` gives them.
    pub terms: Vec<(Vec<u32>, i64)>,
    /// The product and sum formulas, where every term's region was read.
    pub product: Option<f64>,
    pub sum: Option<f64>,
    /// Where some region was past the budget: the expansion over the
    /// largest regions that were read.
    pub partial: Option<Partial>,
}

/// An order's expansion with some of its largest regions left out (each one
/// past the budget): the regions left are closed under intersection and
/// counted afresh, so this is the loop cluster expansion over them.
#[derive(Clone, Debug)]
pub struct Partial {
    /// The largest regions left out, as unions of rings (not pruned).
    pub left_out: Vec<Vec<u32>>,
    pub terms: Vec<(Vec<u32>, i64)>,
    pub product: Option<f64>,
    pub sum: Option<f64>,
}

/// What [`clusters`] gives: the graph's rings, the orders, every region read
/// (in the order first met), and each ring's norm over its belief-propagation
/// value where it was within the budget.
#[derive(Clone, Debug)]
pub struct Expansion {
    pub cells: Vec<Vec<u32>>,
    pub orders: Vec<Order>,
    pub reads: Vec<RegionRead>,
    pub weights: Vec<Option<f64>>,
}

/// The loop cluster expansion of `Π_{q ∈ sites} Z_q` around the observable's
/// ring: the graph's rings (cycles of at most `ring_len` vertices), the
/// orders up to `max_cells`, each region contracted exactly with `env` on its
/// boundary when its estimated work is at most `max_work` multiply-adds and
/// its tables at most `max_bytes`. An order missing a region has no value;
/// the higher orders' regions within the budget are still read. Also each
/// ring's own correction to belief propagation (its norm over its
/// belief-propagation value), where within the budget.
pub fn clusters(t: &Tns, env: &BpEnv, p: &Purified, ring_len: usize, max_cells: usize, max_work: f64, max_bytes: f64) -> Expansion {
    use wai_quantum::quantum_bptn::{cell_unions, cluster_combine, short_cycles};
    let ops: Vec<(u32, Pauli)> = p.sites.iter().map(|&q| (q, Pauli::Z)).collect();
    let mut keep = p.sites.clone();
    keep.sort_unstable();
    keep.dedup();
    let cells = short_cycles(2 * p.n as usize, &p.edges, ring_len);
    let mut reads: Vec<RegionRead> = Vec::new();
    let mut orders = Vec::new();
    let count_cells = |r: &[u32]| cells.iter().filter(|c| c.iter().all(|v| r.binary_search(v).is_ok())).count();
    // A region's value, read once (within the budget) and kept.
    let value_of = |r: &[u32], reads: &mut Vec<RegionRead>| -> Option<f64> {
        if let Some(x) = reads.iter().find(|x| x.region == r) {
            return x.value;
        }
        let plan = t.region_plan(r, &ops).expect("a region");
        let started = std::time::Instant::now();
        let (value, norm_ratio) = if plan.work <= max_work && plan.bytes <= max_bytes {
            let got = t.expect_region_bp(env, r, &ops).expect("a region");
            (Some(got.value), Some(got.norm_ratio))
        } else {
            (None, None)
        };
        let read = RegionRead { region: r.to_vec(), cells: count_cells(r), cuts: plan.cuts, combinations: plan.combinations, work: plan.work, bytes: plan.bytes, value, norm_ratio, seconds: started.elapsed().as_secs_f64() };
        println!(
            "                region of {} rings, {} qubits: cuts {:?}, {} combinations, work {:.2e}, tables {:.2e} B: {}  ({:.1} s)",
            read.cells,
            r.len(),
            read.cuts.iter().map(|c| c.2).collect::<Vec<_>>(),
            read.combinations,
            read.work,
            read.bytes,
            read.value.map_or("past the budget".to_string(), |v| format!("{v:.6}, norm ratio {:.6}", read.norm_ratio.unwrap_or(f64::NAN))),
            read.seconds
        );
        reads.push(read);
        value
    };
    let combine = |terms: &[(Vec<u32>, i64)], values: &[Option<f64>]| -> (Option<f64>, Option<f64>) {
        match values.iter().copied().collect::<Option<Vec<f64>>>() {
            Some(v) => {
                let (a, b) = cluster_combine(terms, &v);
                (Some(a), Some(b))
            }
            None => (None, None),
        }
    };
    for k in 1..=max_cells {
        let maxima = cell_unions(&p.ring, &cells, k);
        let terms = t.loop_cluster_terms(&maxima, &ops).expect("regions holding the observable");
        let values: Vec<Option<f64>> = terms.iter().map(|(r, _)| value_of(r, &mut reads)).collect();
        let (product, sum) = combine(&terms, &values);
        println!("                order {k}: {} terms, product {:?}, sum {:?}", terms.len(), product, sum);
        // Where a region is past the budget: the expansion over the maxima
        // that were read, its intersections read too where within the budget.
        let mut partial = None;
        if product.is_none() {
            let (read, left_out): (Vec<Vec<u32>>, Vec<Vec<u32>>) = maxima.iter().cloned().partition(|m| {
                let r = t.prune_region(m, &keep).expect("a region");
                reads.iter().any(|x| x.region == r && x.value.is_some())
            });
            if !read.is_empty() {
                let terms = t.loop_cluster_terms(&read, &ops).expect("regions holding the observable");
                let values: Vec<Option<f64>> = terms.iter().map(|(r, _)| value_of(r, &mut reads)).collect();
                let (product, sum) = combine(&terms, &values);
                println!("                order {k} without {} of its {} largest regions: {} terms, product {:?}, sum {:?}", left_out.len(), maxima.len(), terms.len(), product, sum);
                partial = Some(Partial { left_out, terms, product, sum });
            }
        }
        orders.push(Order { cells: k, terms, product, sum, partial });
    }
    let weights = cells
        .iter()
        .map(|c| {
            let plan = t.region_plan(c, &[]).expect("a ring");
            (plan.work <= max_work && plan.bytes <= max_bytes).then(|| t.loop_weight(env, c).expect("a ring"))
        })
        .collect();
    Expansion { cells, orders, reads, weights }
}

/// The whole operator `|W⟩⟩` (or, without the middle, `|T T⁻¹⟩⟩`, whose exact
/// echo is 1) at bond cap `chi`. Returns the network and the seconds taken.
pub fn whole(p: &Purified, chi: usize, threads: usize, with_middle: bool) -> (Tns, f64) {
    let started = std::time::Instant::now();
    let mut t = Tns::new(2 * p.n, &p.edges, chi).expect("the graph").with_blocked_linalg(threads);
    t.run(&p.prep).expect("gates on the graph");
    t.run(&p.head).expect("gates on the graph");
    if with_middle {
        t.run(&p.middle).expect("gates on the graph");
    }
    t.run(&p.tail).expect("gates on the graph");
    (t, started.elapsed().as_secs_f64())
}