// Parity test for DPA-2 with use_three_body=True (se_t-style three-body block).

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

use deepmd::descriptor_dpa2::{
    build_dpa2_descriptor, DPA2Config, DPA2Inputs, RepinitArgs,
};
use deepmd::descriptor_repformers::RepformersConfig;
use deepmd::descriptor_se_t_tebd::TebdInputMode;
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_repinit: f64, rcut_smth_repinit: f64,
    rcut_repformer: f64, rcut_smth_repformer: f64,
    three_body_rcut: f64, three_body_rcut_smth: f64,
    nsel: usize, ntypes: usize, tebd_dim: usize,
    repinit_neuron: Vec<usize>, repinit_axis_neuron: usize,
    three_body_neuron: Vec<usize>, three_body_sel: usize,
    rep_nlayers: usize, rep_g1_dim: usize, rep_g2_dim: usize, rep_axis_neuron: usize,
    concat_output_tebd: bool, use_three_body: 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 dpa2_three_body_matches_python_reference() {
    let path = Path::new(env!("CARGO_MANIFEST_DIR"))
        .parent().unwrap().parent().unwrap()
        .join("tests/fixtures/dpa2_three_body.json");
    let fx: Fixture = serde_json::from_str(&std::fs::read_to_string(&path).unwrap()).unwrap();

    assert!(fx.config.use_three_body);

    let cfg = DPA2Config {
        ntypes: fx.config.ntypes,
        tebd_dim: fx.config.tebd_dim,
        concat_output_tebd: fx.config.concat_output_tebd,
        add_tebd_to_repinit_out: false,
        repinit_three_body: Some(RepinitArgs {
            rcut: fx.config.three_body_rcut,
            rcut_smth: fx.config.three_body_rcut_smth,
            nsel: fx.config.three_body_sel,
            neuron: fx.config.three_body_neuron.clone(),
            axis_neuron: 1,
            tebd_input_mode: TebdInputMode::Concat,
            type_one_side: true,
            resnet_dt: false,
            use_three_body: false,
        }),
        activation_function: "tanh".into(),
        repinit: RepinitArgs {
            rcut: fx.config.rcut_repinit,
            rcut_smth: fx.config.rcut_smth_repinit,
            nsel: fx.config.nsel,
            neuron: fx.config.repinit_neuron.clone(),
            axis_neuron: fx.config.repinit_axis_neuron,
            tebd_input_mode: TebdInputMode::Concat,
            type_one_side: true,
            resnet_dt: false,
            use_three_body: true,
        },
        repformers: RepformersConfig {
            rcut: fx.config.rcut_repformer,
            rcut_smth: fx.config.rcut_smth_repformer,
            sel: 2,
            ntypes: fx.config.ntypes,
            nlayers: fx.config.rep_nlayers,
            g1_dim: fx.config.rep_g1_dim,
            g2_dim: fx.config.rep_g2_dim,
            axis_neuron: fx.config.rep_axis_neuron,
            update_g1_has_conv: false, update_g1_has_drrd: true,
            update_g1_has_grrg: true, update_g1_has_attn: false,
            update_g2_has_g1g1: false, update_g2_has_attn: false,
            update_h2: false,
            attn1_hidden: 8, attn1_nhead: 1,
            attn2_hidden: 4, attn2_nhead: 1, attn2_has_gate: false,
            activation_function: "tanh".into(), update_style: "res_avg".into(),
            smooth: true, use_sqrt_nnei: true,
            g1_out_conv: true, g1_out_mlp: true,
            ln_eps: 1e-5, epsilon: 1e-4,
        },
    };

    let nf = fx.config.nf;
    let nloc = fx.config.nloc;
    let nall = fx.config.nall;
    let nsel_repinit = fx.config.nsel;
    let nsel_repformer = 2;
    let tb_sel = fx.config.three_body_sel;

    let mut g = Graph::new("dpa2_three_body_parity");
    let env_mat_repinit = g.input("env_mat_repinit", Shape::new(&[nf, nloc, nsel_repinit, 4], DType::F32));
    let env_mat_repformer = g.input("env_mat_repformer", Shape::new(&[nf, nloc, nsel_repformer, 4], DType::F32));
    let env_mat_tb = g.input("env_mat_three_body", Shape::new(&[nf, nloc, tb_sel, 4], DType::F32));
    let atype = g.input("atype", Shape::new(&[nf, nloc], DType::I32));
    let nei_atype = g.input("nei_atype", Shape::new(&[nf, nloc, nsel_repinit], DType::I32));
    let nei_atype_tb = g.input("nei_atype_tb", Shape::new(&[nf, nloc, tb_sel], DType::I32));
    let nlist_repinit = g.input("nlist_repinit", Shape::new(&[nf, nloc, nsel_repinit], DType::I32));
    let nlist_mask_repinit = g.input("nlist_mask_repinit", Shape::new(&[nf, nloc, nsel_repinit], DType::F32));
    let nlist_repformer = g.input("nlist_repformer", Shape::new(&[nf, nloc, nsel_repformer], DType::I32));
    let nlist_mask_repformer = g.input("nlist_mask_repformer", Shape::new(&[nf, nloc, nsel_repformer], DType::F32));
    let sw_repinit = g.input("sw_repinit", Shape::new(&[nf, nloc, nsel_repinit], DType::F32));
    let sw_repformer = g.input("sw_repformer", Shape::new(&[nf, nloc, nsel_repformer], DType::F32));
    let sw_tb = g.input("sw_three_body", Shape::new(&[nf, nloc, tb_sel], DType::F32));
    let type_embed = g.input("type_embed", Shape::new(&[cfg.ntypes + 1, cfg.tebd_dim], DType::F32));
    let g1_ext_tebd = g.input("g1_ext_tebd", Shape::new(&[nf, nall, cfg.tebd_dim], DType::F32));
    let mapping = g.input("mapping", Shape::new(&[nf, nall], DType::I32));

    let out = build_dpa2_descriptor(
        &mut g, &cfg,
        DPA2Inputs {
            env_mat_raw_repinit: env_mat_repinit,
            env_mat_raw_repformer: env_mat_repformer,
            atype_loc: atype, nei_atype,
            nlist_repinit, nlist_mask_repinit, sw_repinit,
            nlist_repformer, nlist_mask_repformer, sw_repformer,
            type_embedding: type_embed,
            g1_ext_tebd,
            mapping,
            env_mat_raw_three_body: Some(env_mat_tb),
            nei_atype_three_body: Some(nei_atype_tb),
            sw_three_body: Some(sw_tb),
        },
        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_ri = find_named(&fx.inputs, "env_mat_raw_repinit");
    let em_rf = find_named(&fx.inputs, "env_mat_raw_repformer");
    let em_tb = find_named(&fx.inputs, "env_mat_raw_three_body");
    let atype_ext = find_named(&fx.inputs, "atype_ext");
    let nei_atype_in = find_named(&fx.inputs, "nei_atype");
    let nei_atype_tb_in = find_named(&fx.inputs, "nei_atype_three_body");
    let nlist_ri = find_named(&fx.inputs, "nlist_repinit");
    let nlist_rf = find_named(&fx.inputs, "nlist_repformer");
    let nm_ri = find_named(&fx.inputs, "nlist_mask_repinit");
    let nm_rf = find_named(&fx.inputs, "nlist_mask_repformer");
    let sw_ri = find_named(&fx.inputs, "sw_repinit");
    let sw_rf = find_named(&fx.inputs, "sw_repformer");
    let sw_tb_in = find_named(&fx.inputs, "sw_three_body");
    let tab_in = find_named(&fx.inputs, "type_embedding");
    let g1_ext_in = find_named(&fx.inputs, "g1_ext_tebd");

    let atype_loc_f32: Vec<f32> = atype_ext.data.iter().take(nf * nloc).copied().collect();
    let mapping_f32: Vec<f32> = (0..nall).map(|i| (i % nloc) as f32).collect();
    let mapping_f32 = mapping_f32.repeat(nf);

    let outputs = compiled.run(&[
        ("env_mat_repinit", em_ri.data.as_slice()),
        ("env_mat_repformer", em_rf.data.as_slice()),
        ("env_mat_three_body", em_tb.data.as_slice()),
        ("atype", atype_loc_f32.as_slice()),
        ("nei_atype", nei_atype_in.data.as_slice()),
        ("nei_atype_tb", nei_atype_tb_in.data.as_slice()),
        ("nlist_repinit", nlist_ri.data.as_slice()),
        ("nlist_mask_repinit", nm_ri.data.as_slice()),
        ("nlist_repformer", nlist_rf.data.as_slice()),
        ("nlist_mask_repformer", nm_rf.data.as_slice()),
        ("sw_repinit", sw_ri.data.as_slice()),
        ("sw_repformer", sw_rf.data.as_slice()),
        ("sw_three_body", sw_tb_in.data.as_slice()),
        ("type_embed", tab_in.data.as_slice()),
        ("g1_ext_tebd", g1_ext_in.data.as_slice()),
        ("mapping", mapping_f32.as_slice()),
    ]);
    let exp_node = find_named(&fx.outputs, "node_ebd");
    let mut max_abs = 0f32;
    let mut max_at = 0usize;
    for (i, (a, b)) in outputs[0].iter().zip(exp_node.data.iter()).enumerate() {
        let d = (a - b).abs();
        if d > max_abs { max_abs = d; max_at = i; }
    }
    eprintln!("node_ebd: max_abs={max_abs:.3e} at [{max_at}] (actual={:.6e} expected={:.6e})",
              outputs[0][max_at], exp_node.data[max_at]);
    assert!(max_abs <= 1e-5, "drift {max_abs:.3e}");
}
