//! The cost of one amplitude of a circuit by tensor-network contraction, as
//! the hyper-optimised path search finds it.
//!
//! ```text
//! cargo run --release --example tn_cost --features quantum_tn -- CIRCUIT.qasm TRIALS SEED [LOG2_MAX_SIZE|-] [RESTARTS]
//! ```
//!
//! `CIRCUIT.qasm`: OpenQASM 3 with `rz`, `rx`, `ry`, `h`, `x`, `sx`, `cz`,
//! `cx` (barriers and measurements skipped). Prints log10 of the complex
//! multiply-adds for one amplitude (every slice), log2 of the largest
//! intermediate, and the number of slices. The network splits each `cz` and
//! `cx` at the qubits it is block-diagonal on ([`Network::amplitude_split`]);
//! `TN_DENSE=1` keeps every gate whole.

use std::time::Instant;
use wai_quantum::gates::Gate;
use wai_quantum::quantum_tn::{Network, PathOptions, find_path_restarts};

fn main() {
    let args: Vec<String> = std::env::args().collect();
    let src = std::fs::read_to_string(&args[1]).expect("the circuit reads");
    let trials: usize = args[2].parse().unwrap();
    let seed: u64 = args[3].parse().unwrap();
    let max_size = args.get(4).filter(|s| *s != "-").map(|s| s.parse::<f64>().unwrap().exp2());
    let restarts: usize = args.get(5).map(|s| s.parse().unwrap()).unwrap_or(1);
    let mut n = 0u32;
    let mut gates = Vec::new();
    for line in src.lines().map(str::trim) {
        let line = line.trim_end_matches(';');
        if let Some(rest) = line.strip_prefix("qubit[") {
            n = rest.split(']').next().unwrap().parse().unwrap();
            continue;
        }
        if line.is_empty() || ["OPENQASM", "include", "bit", "barrier", "c["].iter().any(|p| line.starts_with(p)) || line.contains("measure") {
            continue;
        }
        let (head, args) = line.split_once(' ').expect("a gate and its qubits");
        let (name, theta) = match head.split_once('(') {
            Some((n, p)) => (n, p.trim_end_matches(')').parse::<f64>().expect("one angle")),
            None => (head, 0.0),
        };
        let qs: Vec<u32> = args.split(',').map(|a| a.trim().trim_start_matches("q[").trim_end_matches(']').parse().unwrap()).collect();
        gates.push(match (name, qs.as_slice()) {
            ("rz", [a]) => Gate::rz(*a, theta),
            ("rx", [a]) => Gate::rx(*a, theta),
            ("ry", [a]) => Gate::ry(*a, theta),
            ("h", [a]) => Gate::h(*a),
            ("x", [a]) => Gate::x(*a),
            ("sx", [a]) => Gate::sqrt_x(*a),
            ("cz", [a, b]) => Gate::cz(*a, *b),
            ("cx", [a, b]) => Gate::cx(*a, *b),
            _ => panic!("gate {name} on {qs:?}"),
        });
    }
    let bits = vec![0u8; n as usize];
    // TN_DENSE=1: every gate one tensor, as before splitting.
    let dense = std::env::var("TN_DENSE").as_deref() == Ok("1");
    let mut net = if dense { Network::amplitude(n, &gates, &bits, &[]) } else { Network::amplitude_split(n, &gates, &bits, &[]) }.expect("the network builds");
    // Split networks are reduced (projections fixed through diagonal gates);
    // TN_DENSE=1 is the pipeline before both.
    if dense {
        net.simplify();
    } else {
        net.reduce();
    }
    let t = Instant::now();
    let plan = find_path_restarts(&net.shape(), &net.dims, &net.open, &PathOptions { trials, seed, max_size, ..PathOptions::default() }, restarts);
    println!(
        "{n} qubits, {} gates: log10 flops {:.3}, log2 largest {:.1}, slices {:.3e}, searched in {:?}",
        gates.len(),
        plan.total_flops().log10(),
        plan.per_slice.max_size.log2(),
        plan.slices,
        t.elapsed()
    );
}