// Parity test for DPA-3 with add_chg_spin_ebd=True.

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

use deepmd::descriptor_dpa3::{build_dpa3_descriptor, DPA3Config, DPA3Inputs};
use deepmd::descriptor_repflows::RepflowsConfig;
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,
    a_rcut: f64, a_rcut_smth: f64,
    nsel: usize, a_sel: usize, ntypes: usize,
    nlayers: usize, n_dim: usize, e_dim: usize, a_dim: usize,
    tebd_dim: usize,
    axis_neuron: usize,
    update_angle: bool,
    activation_function: String, smooth: bool,
    concat_output_tebd: bool, add_chg_spin_ebd: bool,
    nf: usize, nloc: usize, 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")
}

#[test]
fn dpa3_chgspin_matches_python_reference() {
    let path = Path::new(env!("CARGO_MANIFEST_DIR"))
        .parent().unwrap().parent().unwrap()
        .join("tests/fixtures/dpa3_chgspin.json");
    let fx: Fixture = serde_json::from_str(&std::fs::read_to_string(&path).unwrap()).unwrap();

    let cfg = DPA3Config {
        repflows: RepflowsConfig {
            rcut: fx.config.rcut, rcut_smth: fx.config.rcut_smth,
            nsel: fx.config.nsel,
            a_rcut: fx.config.a_rcut, a_rcut_smth: fx.config.a_rcut_smth,
            a_sel: fx.config.a_sel, ntypes: fx.config.ntypes,
            nlayers: fx.config.nlayers,
            n_dim: fx.config.n_dim, e_dim: fx.config.e_dim, a_dim: fx.config.a_dim,
            activation_function: fx.config.activation_function.clone(),
            smooth: fx.config.smooth, ln_eps: 1e-5,
            axis_neuron: fx.config.axis_neuron,
            update_angle: fx.config.update_angle,
        },
        ntypes: fx.config.ntypes,
        tebd_dim: fx.config.tebd_dim,
        concat_output_tebd: fx.config.concat_output_tebd,
        use_loc_mapping: false,
        add_chg_spin_ebd: fx.config.add_chg_spin_ebd,
        activation_function: fx.config.activation_function.clone(),
    };

    let nf = fx.config.nf;
    let nloc = fx.config.nloc;
    let nall = fx.config.nall;
    let nnei = cfg.repflows.nsel;
    let a_sel = cfg.repflows.a_sel;

    let mut g = Graph::new("dpa3_chgspin_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 nlist = g.input("nlist", Shape::new(&[nf, nloc, nnei], DType::I32));
    let nlist_mask = g.input("nlist_mask", Shape::new(&[nf, nloc, nnei], DType::F32));
    let sw = g.input("sw", Shape::new(&[nf, nloc, nnei], DType::F32));
    let angle_input = g.input("angle_input", Shape::new(&[nf, nloc, a_sel, a_sel, 1], DType::F32));
    let angle_mask = g.input("angle_mask", Shape::new(&[nf, nloc, a_sel, a_sel], DType::F32));
    let type_embed = g.input("type_embed", Shape::new(&[cfg.ntypes + 1, cfg.tebd_dim], DType::F32));
    let node_ext = g.input("node_ext", Shape::new(&[nf, nall, cfg.repflows.n_dim], DType::F32));
    let charge_spin = g.input("charge_spin", Shape::new(&[nf, 2], DType::F32));
    let chg_table = g.input("chg_table", Shape::new(&[201, cfg.tebd_dim], DType::F32));
    let spin_table = g.input("spin_table", Shape::new(&[101, cfg.tebd_dim], DType::F32));

    let out = build_dpa3_descriptor(
        &mut g, &cfg,
        DPA3Inputs {
            env_mat_raw: env_mat,
            atype_loc: atype,
            nlist, nlist_mask, sw,
            angle_input, angle_mask,
            type_embedding: type_embed,
            node_ebd_ext: node_ext,
            charge_spin: Some(charge_spin),
            chg_table: Some(chg_table),
            spin_table: Some(spin_table),
        },
        nf, nloc, nall,
    ).expect("build");
    g.set_outputs(vec![out.descriptor]);

    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 nlist_in = find_named(&fx.inputs, "nlist");
    let nlist_mask_in = find_named(&fx.inputs, "nlist_mask");
    let sw_in = find_named(&fx.inputs, "sw");
    let node_ext_in = find_named(&fx.inputs, "node_ebd_ext");
    let tab_in = find_named(&fx.inputs, "type_embedding");
    let cs_in = find_named(&fx.inputs, "charge_spin");
    let chg_in = find_named(&fx.inputs, "chg_table");
    let spin_in = find_named(&fx.inputs, "spin_table");
    let atype_loc_f32: Vec<f32> = atype_ext.data.iter().take(nf * nloc).copied().collect();
    let zero_ai = vec![0f32; nf * nloc * a_sel * a_sel];
    let zero_am = vec![0f32; nf * nloc * a_sel * a_sel];

    let outputs = compiled.run(&[
        ("env_mat", em_raw.data.as_slice()),
        ("atype", atype_loc_f32.as_slice()),
        ("nlist", nlist_in.data.as_slice()),
        ("nlist_mask", nlist_mask_in.data.as_slice()),
        ("sw", sw_in.data.as_slice()),
        ("angle_input", zero_ai.as_slice()),
        ("angle_mask", zero_am.as_slice()),
        ("type_embed", tab_in.data.as_slice()),
        ("node_ext", node_ext_in.data.as_slice()),
        ("charge_spin", cs_in.data.as_slice()),
        ("chg_table", chg_in.data.as_slice()),
        ("spin_table", spin_in.data.as_slice()),
    ]);
    let exp = find_named(&fx.outputs, "node_ebd");
    let mut max_abs = 0f32;
    for (a, b) in outputs[0].iter().zip(exp.data.iter()) {
        max_abs = max_abs.max((a - b).abs());
    }
    eprintln!("node_ebd: max_abs={max_abs:.3e}");
    assert!(max_abs <= 1e-5, "drift {max_abs:.3e}");
}
