// Throughput benchmark for the end-to-end DP energy graph
// (se_e2_a + invar fitting) across CPU / Metal / MLX backends.
//
//   cargo bench -p deepmd --bench dp_energy_bench

use std::collections::HashMap;
use std::path::Path;

use criterion::{black_box, criterion_group, criterion_main, Criterion};
use deepmd::config::{DPModelConfig, EnerFittingConfig, SeAConfig};
use deepmd::model::build_dp_energy_graph;
use rlx_driver::Device;
use rlx_ir::DType;
use rlx_runtime::Session;
use serde::Deserialize;

#[derive(Deserialize)]
struct Tensor { name: String, #[allow(dead_code)] shape: Vec<usize>, data: Vec<f32> }

#[derive(Deserialize)]
struct Config {
    rcut: f64, rcut_smth: f64, sel: Vec<usize>, ntypes: usize,
    neuron: Vec<usize>, axis_neuron: usize, dim_descrpt: usize,
    fit_neuron: Vec<usize>, fit_dim_out: usize,
    activation_function: String,
    nf: usize, nloc: usize, #[allow(dead_code)] nall: usize, #[allow(dead_code)] nnei: usize,
}

#[derive(Deserialize)]
struct Fixture {
    config: Config, params: HashMap<String, Tensor>,
    inputs: Vec<Tensor>, outputs: Vec<Tensor>,
}

fn find_named<'a>(list: &'a [Tensor], name: &str) -> &'a Tensor {
    list.iter().find(|t| t.name == name).expect("missing tensor")
}

fn bench_on(c: &mut Criterion, device: Device, label: &str) {
    bench_fixture_on(c, device, label, "dp_energy_small.json");
}

fn bench_fixture_on(c: &mut Criterion, device: Device, label: &str, fixture: &str) {
    if !rlx_runtime::is_available(device) {
        eprintln!("SKIP {label}: backend not available");
        return;
    }
    let path = Path::new(env!("CARGO_MANIFEST_DIR"))
        .parent().unwrap().parent().unwrap()
        .join("tests/fixtures").join(fixture);
    let fx: Fixture = serde_json::from_str(&std::fs::read_to_string(&path).unwrap()).unwrap();

    let cfg = DPModelConfig {
        descriptor: SeAConfig {
            rcut: fx.config.rcut, rcut_smth: fx.config.rcut_smth,
            sel: fx.config.sel.clone(),
            neuron: fx.config.neuron.clone(),
            axis_neuron: fx.config.axis_neuron,
            resnet_dt: false, type_one_side: true,
            activation_function: fx.config.activation_function.clone(),
            type_map: None,
        },
        fitting_net: EnerFittingConfig {
            ntypes: fx.config.ntypes,
            dim_descrpt: fx.config.dim_descrpt,
            dim_out: fx.config.fit_dim_out,
            neuron: fx.config.fit_neuron.clone(),
            resnet_dt: false,
            numb_fparam: 0, numb_aparam: 0, dim_case_embd: 0,
            activation_function: fx.config.activation_function.clone(),
            mixed_types: true,
        },
        type_map: None,
    };

    let (g, handle) = build_dp_energy_graph(&cfg, fx.config.nf, fx.config.nloc).expect("build");
    let mut g = g;
    g.set_outputs(vec![handle.atom_energy, handle.energy]);

    let session = Session::new(device);
    let mut compiled = session.compile(g);
    for (name, t) in &fx.params {
        let bytes: Vec<u8> = t.data.iter().flat_map(|v| v.to_le_bytes()).collect();
        compiled.set_param_typed(name, &bytes, DType::F32);
    }
    let em_raw = find_named(&fx.inputs, "env_mat_raw");
    let atype_in = find_named(&fx.inputs, "atype");
    let atype_f32: Vec<f32> = atype_in.data.iter().copied().collect();

    let em_data = em_raw.data.clone();
    let atype_data = atype_f32.clone();

    c.bench_function(label, |b| {
        b.iter(|| {
            let outputs = compiled.run(&[
                ("env_mat", black_box(em_data.as_slice())),
                ("atype", black_box(atype_data.as_slice())),
            ]);
            black_box(outputs);
        })
    });
}

fn dp_energy_benches(c: &mut Criterion) {
    // nf=1 (tiny)
    bench_on(c, Device::Cpu, "dp_energy/cpu");
    bench_on(c, Device::Metal, "dp_energy/metal");
    bench_on(c, Device::Mlx, "dp_energy/mlx");
    bench_on(c, Device::Cuda, "dp_energy/cuda");
    bench_on(c, Device::Gpu, "dp_energy/wgpu");

    // Realistic batch sizes — this is where GPU latency amortises.
    for nf in [8, 32, 128] {
        let fixture = format!("dp_energy_nf{nf}.json");
        bench_fixture_on(c, Device::Cpu,   &format!("dp_energy_nf{nf}/cpu"),   &fixture);
        bench_fixture_on(c, Device::Metal, &format!("dp_energy_nf{nf}/metal"), &fixture);
        bench_fixture_on(c, Device::Mlx,   &format!("dp_energy_nf{nf}/mlx"),   &fixture);
        bench_fixture_on(c, Device::Cuda,  &format!("dp_energy_nf{nf}/cuda"),  &fixture);
        bench_fixture_on(c, Device::Gpu,   &format!("dp_energy_nf{nf}/wgpu"),  &fixture);
    }
}

criterion_group!(benches, dp_energy_benches);
criterion_main!(benches);
