// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, version 3.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program. If not, see <https://www.gnu.org/licenses/>.

//! Layer-0 last-token intermediates vs HF dump (`hf_l0_intermediates.json`).

use rlx_qwen3_tts::Qwen3TtsConfig;
use rlx_qwen3_tts::load::Qwen3TtsWeightStore;
use rlx_qwen3_tts::prompt::{build_custom_voice_prompt, load_text_tokenizer};
use rlx_qwen3_tts::talker::eager::TalkerEagerModel;
use rlx_qwen3_tts::text_embed::TextEmbedder;
use std::path::PathBuf;

fn max_abs(a: &[f32], b: &[f32]) -> f32 {
    a.iter()
        .zip(b.iter())
        .map(|(x, y)| (x - y).abs())
        .fold(0f32, f32::max)
}

#[test]
fn layer0_attn_and_block_match_hf() {
    let Some(model_dir) = std::env::var("RLX_QWEN3_TTS_DIR").ok().map(PathBuf::from) else {
        eprintln!("skip: RLX_QWEN3_TTS_DIR");
        return;
    };
    let hf_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
        .join("../../.cache/qwen3-tts/hf_l0_intermediates.json");
    if !hf_path.is_file() {
        eprintln!("skip: run dump_talker_l0_intermediates.py first");
        return;
    }
    let hf: serde_json::Value = serde_json::from_slice(&std::fs::read(&hf_path).unwrap()).unwrap();
    let hf_attn: Vec<f32> = hf["after_attn"]
        .as_array()
        .unwrap()
        .iter()
        .map(|v| v.as_f64().unwrap() as f32)
        .collect();
    let hf_block: Vec<f32> = hf["after_block"]
        .as_array()
        .unwrap()
        .iter()
        .map(|v| v.as_f64().unwrap() as f32)
        .collect();

    let cfg = Qwen3TtsConfig::from_model_dir(&model_dir).unwrap();
    let store = Qwen3TtsWeightStore::open(&model_dir).unwrap();
    let tokenizer = load_text_tokenizer(&model_dir).unwrap();
    let text_embedder = TextEmbedder::open(&store).unwrap();
    let prompt = build_custom_voice_prompt(
        &cfg,
        &store,
        &text_embedder,
        &tokenizer,
        "Hi.",
        "vivian",
        "english",
    )
    .unwrap();

    let mut eager = TalkerEagerModel::open(&store, cfg.talker()).unwrap();
    let native_attn = eager.layer0_after_attn_last(prompt.embeds.view()).unwrap();

    let through0 = eager
        .prefill_through_layer(prompt.embeds.view(), 0)
        .unwrap();
    let native_block: Vec<f32> = through0.row(through0.nrows() - 1).iter().copied().collect();

    let da = max_abs(&native_attn, &hf_attn);
    let db = max_abs(&native_block, &hf_block);
    eprintln!("layer0 after_attn max_abs={da}");
    eprintln!("layer0 after_block max_abs={db}");
    assert!(da < 0.012, "after_attn diverged (max_abs={da})");
    assert!(db < 0.012, "after_block diverged (max_abs={db})");
}
