// Parity test for DPA-1 (simplest case: attn_layer=0).

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

use deepmd::descriptor_dpa1::{build_dpa1_descriptor, DPA1Config, DPA1Inputs};
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: f64,
    rcut_smth: f64,
    sel: Vec<usize>,
    ntypes: usize,
    neuron: Vec<usize>,
    axis_neuron: usize,
    tebd_dim: usize,
    tebd_input_mode: String,
    resnet_dt: bool,
    type_one_side: bool,
    attn: usize,
    attn_layer: usize,
    attn_dotr: bool,
    num_heads: usize,
    normalize: bool,
    scaling_factor: f32,
    ln_eps: f32,
    smooth: bool,
    activation_function: String,
    concat_output_tebd: bool,
    nf: usize,
    nloc: 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 assert_allclose(actual: &[f32], expected: &[f32], atol: f32, rtol: f32, label: &str) {
    assert_eq!(actual.len(), expected.len());
    let mut max_abs = 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);
        assert!(abs <= tol, "{label}[{i}]: actual={a:.6e} expected={b:.6e} |Δ|={abs:.3e}");
    }
    eprintln!("{label}: max_abs={max_abs:.3e}");
}

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

    let mode = match fx.config.tebd_input_mode.as_str() {
        "concat" => TebdInputMode::Concat,
        "strip" => TebdInputMode::Strip,
        _ => panic!(),
    };
    let cfg = DPA1Config {
        rcut: fx.config.rcut,
        rcut_smth: fx.config.rcut_smth,
        sel: fx.config.sel.clone(),
        ntypes: fx.config.ntypes,
        neuron: fx.config.neuron.clone(),
        axis_neuron: fx.config.axis_neuron,
        tebd_dim: fx.config.tebd_dim,
        tebd_input_mode: mode,
        resnet_dt: fx.config.resnet_dt,
        attn: fx.config.attn,
        attn_layer: fx.config.attn_layer,
        attn_dotr: fx.config.attn_dotr,
        num_heads: fx.config.num_heads,
        normalize: fx.config.normalize,
        scaling_factor: fx.config.scaling_factor,
        ln_eps: fx.config.ln_eps,
        smooth: fx.config.smooth,
        type_one_side: fx.config.type_one_side,
        activation_function: fx.config.activation_function.clone(),
        concat_output_tebd: fx.config.concat_output_tebd,
    };
    let nf = fx.config.nf;
    let nloc = fx.config.nloc;
    let nnei: usize = cfg.sel.iter().sum();

    let mut g = Graph::new("dpa1_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 nei_atype = g.input("nei_atype", 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 t_table = g.input(
        "type_embed_in",
        Shape::new(&[cfg.ntypes + 1, cfg.tebd_dim], DType::F32),
    );
    let inputs = DPA1Inputs {
        env_mat_raw: env_mat,
        atype_loc: atype,
        nei_atype,
        nlist_mask,
        sw,
        type_embedding: t_table,
    };
    let out = build_dpa1_descriptor(&mut g, &cfg, inputs, nf, nloc).expect("build");
    let mut outs = vec![out.descriptor, out.gg];
    if let Some(q) = out.debug_attn_query {
        outs.push(q);
    }
    if let Some(awp) = out.debug_attn_aw_post {
        outs.push(awp);
    }
    if let Some(op) = out.debug_attn_o_pre_proj {
        outs.push(op);
    }
    if let Some(ir) = out.debug_input_r {
        outs.push(ir);
    }
    g.set_outputs(outs);

    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 nei_atype_in = find_named(&fx.inputs, "nei_atype");
    let nlist_mask_in = find_named(&fx.inputs, "nlist_mask");
    let sw_in = find_named(&fx.inputs, "sw");
    let tab_in = find_named(&fx.inputs, "type_embedding");

    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()),
        ("nei_atype", nei_atype_in.data.as_slice()),
        ("nlist_mask", nlist_mask_in.data.as_slice()),
        ("sw", sw_in.data.as_slice()),
        ("type_embed_in", tab_in.data.as_slice()),
    ]);
    let tol = if fx.config.smooth { 1e-2 } else { 1e-4 };
    if let Some(t) = fx.outputs.iter().find(|t| t.name == "attn_query") {
        eprintln!("checking attn_query");
        let _ = assert_allclose_relaxed(&outputs[2], &t.data, "attn_query");
    }
    if let Some(t) = fx.outputs.iter().find(|t| t.name == "attn_aw_post") {
        eprintln!("checking attn_aw_post");
        let _ = assert_allclose_relaxed(&outputs[3], &t.data, "attn_aw_post");
    }
    if let Some(t) = fx.outputs.iter().find(|t| t.name == "attn_o_pre_proj") {
        eprintln!("checking attn_o_pre_proj");
        let _ = assert_allclose_relaxed(&outputs[4], &t.data, "attn_o_pre_proj");
    }
    if let Some(t) = fx.outputs.iter().find(|t| t.name == "attn_input_r") {
        eprintln!("checking attn_input_r");
        let _ = assert_allclose_relaxed(&outputs[5], &t.data, "attn_input_r");
    }
    if let Some(gg_expected) = fx.outputs.iter().find(|t| t.name == "gg_after_attention") {
        assert_allclose(&outputs[1], &gg_expected.data, tol, tol, "gg_after_attention");
    }
    let expected = find_named(&fx.outputs, "descriptor");
    assert_allclose(&outputs[0], &expected.data, tol, tol, "descriptor");
}

#[test]
fn dpa1_matches_python_reference() {
    run_dpa1_fixture("dpa1_small.json");
}

#[test]
fn dpa1_nf2_matches_python_reference() {
    run_dpa1_fixture("dpa1_nf2.json");
}

fn assert_allclose_relaxed(actual: &[f32], expected: &[f32], label: &str) -> f32 {
    let max_abs = actual
        .iter()
        .zip(expected.iter())
        .map(|(a, b)| (a - b).abs())
        .fold(0f32, f32::max);
    eprintln!("{label}: max_abs={max_abs:.3e}");
    max_abs
}
