// Cross-backend parity smoke test: run a representative fixture through
// `Device::Metal` (Apple GPU) and compare element-wise against the CPU
// reference.  Skipped automatically when Metal is unavailable.
//
// This is the first non-CPU validation in the test suite — every other
// test compiles to `Device::Cpu` (rlx-cpu).  We pick the end-to-end DP
// energy graph (se_e2_a + invar fitting + ΣE) since it covers the most
// op kinds in one graph: gather, matmul, activation, reduce, narrow,
// transpose, concat.

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

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 run_on(device: Device) -> Result<f32, String> {
    if !rlx_runtime::is_available(device) {
        return Err(format!("{device:?} backend not registered/available"));
    }

    let path = Path::new(env!("CARGO_MANIFEST_DIR"))
        .parent().unwrap().parent().unwrap()
        .join("tests/fixtures/dp_energy_small.json");
    let fx: Fixture = serde_json::from_str(&std::fs::read_to_string(&path).unwrap()).unwrap();

    let 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,
    };
    let 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,
    };
    let cfg = DPModelConfig { descriptor, fitting_net, type_map: None };

    let (g, handle) = build_dp_energy_graph(&cfg, fx.config.nf, fx.config.nloc)
        .map_err(|e| e.to_string())?;
    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 outputs = compiled.run(&[
        ("env_mat", em_raw.data.as_slice()),
        ("atype", atype_f32.as_slice()),
    ]);
    let exp = find_named(&fx.outputs, "atom_energy");
    let mut max_abs = 0f32;
    for (a, b) in outputs[0].iter().zip(exp.data.iter()) {
        max_abs = max_abs.max((a - b).abs());
    }
    Ok(max_abs)
}

#[test]
fn dp_energy_matches_cpu_reference_on_metal() {
    match run_on(Device::Metal) {
        Ok(drift) => {
            eprintln!("Metal drift (vs Python ref): {drift:.3e}");
            // Cross-backend f32 op ordering can vary the last bit;
            // accept anything within Python's own f32 round-trip
            // tolerance for the same graph.
            assert!(drift <= 1e-4, "Metal drift {drift:.3e} > 1e-4");
        }
        Err(msg) => {
            eprintln!("SKIP: {msg}");
        }
    }
}

#[test]
fn dp_energy_matches_cpu_reference_on_cpu_sanity() {
    // Sanity check: same harness on Cpu must still pass.  This guards
    // the cross-backend test from divergent reference loading bugs.
    let drift = run_on(Device::Cpu).expect("Cpu run");
    eprintln!("CPU drift: {drift:.3e}");
    assert!(drift <= 1e-5, "CPU drift {drift:.3e}");
}

#[test]
fn dp_energy_matches_cpu_reference_on_mlx() {
    match run_on(Device::Mlx) {
        Ok(drift) => {
            eprintln!("MLX drift (vs Python ref): {drift:.3e}");
            assert!(drift <= 1e-3, "MLX drift {drift:.3e} > 1e-3");
        }
        Err(msg) => eprintln!("SKIP: {msg}"),
    }
}

#[test]
fn dp_energy_matches_cpu_reference_on_cuda() {
    match run_on(Device::Cuda) {
        Ok(drift) => {
            eprintln!("CUDA drift (vs Python ref): {drift:.3e}");
            assert!(drift <= 1e-3, "CUDA drift {drift:.3e} > 1e-3");
        }
        Err(msg) => eprintln!("SKIP: {msg}"),
    }
}
