//! Builders and utilities for benchmarks.

use std::sync::{Arc, LazyLock};

use hugr_core::builder::{
    Dataflow, DataflowSubContainer, FunctionBuilder, HugrBuilder, ModuleBuilder,
};
use hugr_core::extension::ExtensionRegistry;
use hugr_core::extension::prelude::qb_t;
use hugr_core::hugr::hugrmut::HugrMut;
use hugr_core::ops::OpName;
use hugr_core::std_extensions::STD_REG;
use hugr_core::std_extensions::arithmetic::float_types::{ConstF64, float64_type};
use hugr_core::types::Signature;
use hugr_core::{CircuitUnit, Extension, Hugr, HugrView, Node};

pub static QUANTUM_EXT: LazyLock<Arc<Extension>> = LazyLock::new(|| {
    Extension::new_arc(
        "bench.quantum".try_into().unwrap(),
        hugr_core::extension::Version::new(0, 0, 0),
        |ext, extension_ref| {
            ext.add_op(
                OpName::new_inline("H"),
                String::new(),
                Signature::new_endo([qb_t()]),
                extension_ref,
            )
            .unwrap();
            ext.add_op(
                OpName::new_inline("Rz"),
                String::new(),
                Signature::new(vec![qb_t(), float64_type()], vec![qb_t()]),
                extension_ref,
            )
            .unwrap();

            ext.add_op(
                OpName::new_inline("CX"),
                String::new(),
                Signature::new_endo(vec![qb_t(), qb_t()]),
                extension_ref,
            )
            .unwrap();
        },
    )
});

pub static BENCH_EXTENSIONS: LazyLock<ExtensionRegistry> = LazyLock::new(|| {
    let mut reg = STD_REG.clone();
    reg.register_updated(QUANTUM_EXT.clone());
    reg
});

/// The node ids for a layer generated by the [`circuit`] function.
pub struct CircuitLayer {
    pub h: Node,
    pub cx1: Node,
    pub cx2: Node,
}

/// Construct a quantum circuit with two qubits and `layers` layers applying `H q0; CX q0, q1; CX q1, q0`.
pub fn circuit(layers: usize) -> (Hugr, Vec<CircuitLayer>) {
    let h_gate = QUANTUM_EXT.instantiate_extension_op("H", []).unwrap();
    let cx_gate = QUANTUM_EXT.instantiate_extension_op("CX", []).unwrap();
    let rz = QUANTUM_EXT.instantiate_extension_op("Rz", []).unwrap();
    let signature = Signature::new_endo(vec![qb_t(), qb_t()]);
    let mut module_builder = ModuleBuilder::new();
    let mut f_build = module_builder.define_function("main", signature).unwrap();

    let wires: Vec<_> = f_build.input_wires().collect();

    let mut linear = f_build.as_circuit(wires);

    let mut layer_ids = Vec::with_capacity(layers);
    for _ in 0..layers {
        linear.append(h_gate.clone(), [0]).unwrap();
        let h = linear.tracked_wire(0).unwrap().node();
        linear.append(cx_gate.clone(), [0, 1]).unwrap();
        let cx1 = linear.tracked_wire(0).unwrap().node();
        linear.append(cx_gate.clone(), [1, 0]).unwrap();
        let cx2 = linear.tracked_wire(0).unwrap().node();

        let angle = linear.add_constant(ConstF64::new(0.5));
        linear
            .append_and_consume(
                rz.clone(),
                [CircuitUnit::Linear(0), CircuitUnit::Wire(angle)],
            )
            .unwrap();

        layer_ids.push(CircuitLayer { h, cx1, cx2 });
    }

    let outs = linear.finish();
    f_build.finish_with_outputs(outs).unwrap();

    (module_builder.finish_hugr().unwrap(), layer_ids)
}

/// Returns a Hugr that serializes to approximately `byte_size` bytes in binary
/// format by defining a large metadata payload.
pub fn big_hugr(byte_size: usize) -> Hugr {
    let big_payload: String = "a".repeat(byte_size);

    let mut hugr = FunctionBuilder::new("main", Signature::new_endo(vec![]))
        .unwrap()
        .finish_hugr()
        .unwrap();
    hugr.set_metadata_any(hugr.entrypoint(), "big", big_payload);
    hugr
}
