// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.

//! End-to-end numerical-parity test for the se_e2_a descriptor.
//!
//! Loads `tests/fixtures/se_e2_a_small.json` (produced by
//! `scripts/gen_se_e2_a_fixture.py` running against the Python
//! reference at `/Users/Shared/deepmd-kit`), builds the equivalent
//! Rust graph, binds parameters, executes through the CPU backend,
//! and asserts element-wise close to the Python output.

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

use deepmd::config::SeAConfig;
use deepmd::descriptor::{build_se_a_descriptor, SeAExtraInputs};
use rlx_ir::{DType, Graph, Shape};
use rlx_runtime::{Device, 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>,
    neuron: Vec<usize>,
    axis_neuron: usize,
    resnet_dt: bool,
    type_one_side: bool,
    activation_function: String,
    nf: usize,
    nloc: usize,
    #[allow(dead_code)]
    nall: usize,
    #[allow(dead_code)]
    ntypes: 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)
        .unwrap_or_else(|| panic!("missing tensor: {name}"))
}

fn assert_allclose(actual: &[f32], expected: &[f32], atol: f32, rtol: f32, label: &str) {
    assert_eq!(
        actual.len(),
        expected.len(),
        "{label}: length mismatch ({} vs {})",
        actual.len(),
        expected.len()
    );
    let mut max_abs = 0f32;
    let mut max_rel = 0f32;
    for (i, (a, b)) in actual.iter().zip(expected.iter()).enumerate() {
        let abs = (a - b).abs();
        let tol = atol + rtol * b.abs();
        max_abs = max_abs.max(abs);
        if b.abs() > 1e-8 {
            max_rel = max_rel.max(abs / b.abs());
        }
        assert!(
            abs <= tol && a.is_nan() == b.is_nan(),
            "{label}[{i}]: actual={a:.6e} expected={b:.6e} |Δ|={abs:.3e} > tol={tol:.3e}",
        );
    }
    eprintln!("{label}: max_abs={max_abs:.3e} max_rel={max_rel:.3e}");
}

#[test]
fn se_e2_a_matches_python_reference() {
    let path = Path::new(env!("CARGO_MANIFEST_DIR"))
        .parent()
        .unwrap()
        .parent()
        .unwrap()
        .join("tests/fixtures/se_e2_a_small.json");
    let text = std::fs::read_to_string(&path)
        .unwrap_or_else(|e| panic!("read fixture {path:?}: {e}"));
    let fx: Fixture = serde_json::from_str(&text).unwrap();

    let cfg = 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: fx.config.resnet_dt,
        type_one_side: fx.config.type_one_side,
        activation_function: fx.config.activation_function.clone(),
        type_map: None,
    };
    let nf = fx.config.nf;
    let nloc = fx.config.nloc;
    let nnei: usize = cfg.sel.iter().sum();

    let mut g = Graph::new("se_e2_a_parity");
    let env_mat = g.input("env_mat", Shape::new(&[nf, nloc, nnei, 4], DType::F32));
    let atype = g.input("atype", Shape::new(&[nf, nloc], DType::I32));
    let out = build_se_a_descriptor(
        &mut g,
        &cfg,
        env_mat,
        atype,
        nf,
        nloc,
        SeAExtraInputs::default(),
    )
    .expect("build");
    g.set_outputs(vec![out.descriptor, out.gr]);

    let session = Session::new(Device::Cpu);
    let mut compiled = session.compile(g);

    // Bind every params/* tensor.
    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);
    }

    // Inputs: the raw env-matrix that the Python EnvMat produces
    // *before* davg/dstd normalization (the Rust graph applies those).
    // And atype_loc (the first nloc entries of atype_ext).
    let em_raw = find_named(&fx.inputs, "env_mat_raw");
    let atype_ext = find_named(&fx.inputs, "atype_ext");
    // CPU backend uses an f32 arena — integer inputs are passed as
    // f32 values (the kernel truncates to int internally).
    let atype_loc_f32: Vec<f32> = atype_ext
        .data
        .iter()
        .take(nf * nloc)
        .copied()
        .collect();

    let outputs = compiled.run(&[
        ("env_mat", em_raw.data.as_slice()),
        ("atype", atype_loc_f32.as_slice()),
    ]);
    let descriptor_actual = &outputs[0];
    let gr_actual = &outputs[1];

    let expected_descriptor = find_named(&fx.outputs, "descriptor");
    let expected_gr = find_named(&fx.outputs, "gr");

    assert_allclose(
        descriptor_actual,
        &expected_descriptor.data,
        1e-4,
        1e-4,
        "descriptor",
    );
    assert_allclose(gr_actual, &expected_gr.data, 1e-4, 1e-4, "gr");
}

