//! Builders and utilities for benchmarks.

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

use hugr::builder::{
    BuildError, CFGBuilder, DFGBuilder, Dataflow, DataflowHugr, DataflowSubContainer,
    FunctionBuilder, HugrBuilder, ModuleBuilder,
};
use hugr::extension::ExtensionRegistry;
use hugr::extension::prelude::{bool_t, qb_t, usize_t};
use hugr::hugr::hugrmut::HugrMut;
use hugr::ops::{OpName, Value, handle::FuncID};
use hugr::std_extensions::STD_REG;
use hugr::std_extensions::arithmetic::float_types::{ConstF64, float64_type};
use hugr::types::Signature;
use hugr::{CircuitUnit, Extension, Hugr, HugrView, Node, type_row};

pub fn simple_dfg_hugr() -> Hugr {
    let dfg_builder = DFGBuilder::new(Signature::new(vec![bool_t()], vec![bool_t()])).unwrap();
    let [i1] = dfg_builder.input_wires_arr();
    dfg_builder.finish_hugr_with_outputs([i1]).unwrap()
}

pub fn simple_cfg_builder<T: AsMut<Hugr> + AsRef<Hugr>>(
    cfg_builder: &mut CFGBuilder<T>,
) -> Result<(), BuildError> {
    let sum2_variants = vec![vec![usize_t()].into(), vec![usize_t()].into()];
    let mut entry_b = cfg_builder.entry_builder(sum2_variants.clone(), type_row![])?;
    let entry = {
        let [inw] = entry_b.input_wires_arr();

        let sum = entry_b.make_sum(1, sum2_variants, [inw])?;
        entry_b.finish_with_outputs(sum, [])?
    };
    let mut middle_b =
        cfg_builder.simple_block_builder(Signature::new(vec![usize_t()], vec![usize_t()]), 1)?;
    let middle = {
        let c = middle_b.add_load_const(hugr::ops::Value::unary_unit_sum());
        let [inw] = middle_b.input_wires_arr();
        middle_b.finish_with_outputs(c, [inw])?
    };
    let exit = cfg_builder.exit_block();
    cfg_builder.branch(&entry, 0, &middle)?;
    cfg_builder.branch(&middle, 0, &exit)?;
    cfg_builder.branch(&entry, 1, &exit)?;
    Ok(())
}

pub fn simple_cfg_hugr() -> Hugr {
    let mut cfg_builder =
        CFGBuilder::new(Signature::new(vec![usize_t()], vec![usize_t()])).unwrap();
    simple_cfg_builder(&mut cfg_builder).unwrap();
    cfg_builder.finish_hugr().unwrap()
}

pub fn dfg_calling_defn_decl() -> (Hugr, FuncID<true>, FuncID<false>) {
    let mut dfb = DFGBuilder::new(Signature::new([], [bool_t()])).unwrap();
    let new_defn = {
        let mut mb = dfb.module_root_builder();
        let fb = mb
            .define_function("helper_id", Signature::new_endo([bool_t()]))
            .unwrap();
        let [f_inp] = fb.input_wires_arr();
        fb.finish_with_outputs([f_inp]).unwrap()
    };
    let new_decl = dfb
        .module_root_builder()
        .declare("helper2", Signature::new_endo([bool_t()]).into())
        .unwrap();
    let cst = dfb.add_load_value(Value::true_val());
    let [c1] = dfb
        .call(new_defn.handle(), &[], [cst])
        .unwrap()
        .outputs_arr();
    let [c2] = dfb.call(&new_decl, &[], [c1]).unwrap().outputs_arr();
    (
        dfb.finish_hugr_with_outputs([c2]).unwrap(),
        *new_defn.handle(),
        new_decl,
    )
}

pub static QUANTUM_EXT: LazyLock<Arc<Extension>> = LazyLock::new(|| {
    Extension::new_arc(
        "bench.quantum".try_into().unwrap(),
        hugr::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([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([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
}
