// Ghost-atom parity test for se_e2_a (nf=2, nall=5, nloc=3).

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).expect("missing tensor")
}

#[test]
fn se_e2_a_ghost_matches_python_reference() {
    let path = Path::new(env!("CARGO_MANIFEST_DIR"))
        .parent().unwrap().parent().unwrap()
        .join("tests/fixtures/se_e2_a_nf2_ghost.json");
    let fx: Fixture = serde_json::from_str(&std::fs::read_to_string(&path).unwrap()).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_ghost_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);
    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_ext = find_named(&fx.inputs, "atype_ext");
    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 exp = find_named(&fx.outputs, "descriptor");
    let exp_gr = find_named(&fx.outputs, "gr");

    let mut max_d = 0f32;
    let mut max_at = 0usize;
    for (i, (a, b)) in outputs[0].iter().zip(exp.data.iter()).enumerate() {
        let d = (a - b).abs();
        if d > max_d { max_d = d; max_at = i; }
    }
    eprintln!("max descriptor drift at [{max_at}]: actual={:.6e} expected={:.6e}",
              outputs[0][max_at], exp.data[max_at]);
    // First 8 actual vs expected
    for i in 0..8 {
        eprintln!("  [{i}]: actual={:.6e} expected={:.6e} diff={:.3e}",
                  outputs[0][i], exp.data[i], (outputs[0][i] - exp.data[i]).abs());
    }
    let mut max_gr = 0f32;
    for (a, b) in outputs[1].iter().zip(exp_gr.data.iter()) {
        max_gr = max_gr.max((a - b).abs());
    }
    eprintln!("descriptor: max_abs={max_d:.3e}");
    eprintln!("gr:         max_abs={max_gr:.3e}");
    assert!(max_d <= 1e-5, "descriptor drift {max_d:.3e}");
    assert!(max_gr <= 1e-5, "gr drift {max_gr:.3e}");
}
