candding 0.1.0

use std::collections::HashMap;

use candle_core::DType;
use candle_core::Device;
use candle_core::Tensor;
use candle_nn::VarBuilder;
use candle_nn::VarMap;

use super::*;
use crate::layers::layer_norm::LayerNorm;
use crate::layers::mask::AttentionMask;
use crate::models::Encoder;
use crate::models::EncoderInput;
use crate::models::gte_new::RopeScaling;
use crate::models::gte_new::attention::GteNewAttention;
use crate::models::gte_new::embeddings::GteNewEmbeddings;
use crate::models::gte_new::layer::GteNewLayer;
use crate::models::gte_new::mlp::GteNewMlp;

fn tiny_config() -> GteNewConfig {
  GteNewConfig {
    hidden_size: 16,
    num_hidden_layers: 2,
    num_attention_heads: 4,
    intermediate_size: 32,
    vocab_size: 32,
    type_vocab_size: 1,
    hidden_act: "gelu".to_string(),
    layer_norm_type: "layer_norm".to_string(),
    layer_norm_eps: 1e-12,
    position_embedding_type: "rope".to_string(),
    rope_theta: 20_000.0,
    rope_scaling: Some(RopeScaling {
      factor: 8.0,
      r#type: "ntk".to_string(),
    }),
    max_position_embeddings: 64,
    pack_qkv: true,
    unpad_inputs: false,
    use_memory_efficient_attention: false,
  }
}

fn tiny_model() -> (VarMap, GteNewModel) {
  let varmap = VarMap::new();
  let vb = VarBuilder::from_varmap(&varmap, DType::F32, &Device::Cpu);
  let model = GteNewModel::load(tiny_config(), vb).unwrap();
  (varmap, model)
}

fn run(model: &GteNewModel, rows: &[Vec<u32>]) -> Vec<Vec<Vec<f32>>> {
  let seq = rows.iter().map(Vec::len).max().unwrap();
  let mut ids = Vec::new();
  let mut mask = Vec::new();
  for row in rows {
    ids.extend(
      row
        .iter()
        .copied()
        .chain(std::iter::repeat_n(0, seq - row.len())),
    );
    mask
      .extend(std::iter::repeat_n(1u32, row.len()).chain(std::iter::repeat_n(0, seq - row.len())));
  }
  let input_ids = Tensor::from_vec(ids, (rows.len(), seq), &Device::Cpu).unwrap();
  let attention_mask = Tensor::from_vec(mask, (rows.len(), seq), &Device::Cpu).unwrap();
  let token_type_ids = input_ids.zeros_like().unwrap();
  let lengths: Vec<usize> = rows.iter().map(Vec::len).collect();
  let input = EncoderInput {
    input_ids: &input_ids,
    attention_mask: &attention_mask,
    token_type_ids: &token_type_ids,
    lengths: &lengths,
  };
  model.forward(&input).unwrap().to_vec3::<f32>().unwrap()
}

fn max_shift(a: &[Vec<Vec<f32>>], b: &[Vec<Vec<f32>>]) -> f32 {
  a.iter()
    .flatten()
    .flatten()
    .zip(b.iter().flatten().flatten())
    .map(|(x, y)| (x - y).abs())
    .fold(0.0, f32::max)
}

#[test]
fn output_has_hidden_size_per_token() {
  let (_varmap, model) = tiny_model();
  let out = run(&model, &[vec![3, 5, 7, 9]]);
  assert_eq!((out.len(), out[0].len(), out[0][0].len()), (1, 4, 16));
  assert_eq!(model.hidden_size(), 16);
}

#[test]
fn padding_does_not_change_real_tokens() {
  let (_varmap, model) = tiny_model();
  let alone = run(&model, &[vec![3, 9, 5]]);
  let batched = run(&model, &[vec![3, 5, 7, 8, 9], vec![3, 9, 5]]);
  for (pos, token) in alone[0].iter().enumerate() {
    for (a, b) in token.iter().zip(&batched[1][pos]) {
      assert!((a - b).abs() < 1e-4, "position {pos}: {a} vs {b}");
    }
  }
}

/// Without rotary the network is reversal-equivariant; reversing an unpadded row leaves its mask unchanged.
#[test]
fn reversing_the_row_is_not_a_reversal_of_the_output() {
  let (_varmap, model) = tiny_model();
  let tokens = vec![3u32, 5, 7, 9, 11, 13];
  let reversed: Vec<u32> = tokens.iter().rev().copied().collect();
  let straight = run(&model, &[tokens]);
  let reversed_out = run(&model, &[reversed]);
  let undone: Vec<Vec<f32>> = reversed_out[0].iter().rev().cloned().collect();
  let shift = max_shift(&[straight[0].clone()], &[undone]);
  assert!(
    shift > 1e-3,
    "rotary is the only positional signal, so reversing the row merely reversed its output rows: max shift {shift}"
  );
}

#[test]
fn the_token_type_row_reaches_the_sum() {
  let (mut varmap, model) = tiny_model();
  let before = run(&model, &[vec![3, 5, 7, 9]]);
  varmap
    .set_one(
      "embeddings.token_type_embeddings.weight",
      Tensor::zeros((1, 16), DType::F32, &Device::Cpu).unwrap(),
    )
    .unwrap();
  let after = run(&model, &[vec![3, 5, 7, 9]]);
  let shift = max_shift(&[before[0].clone()], &[after[0].clone()]);
  assert!(
    shift > 1e-6,
    "zeroing the token-type row did not move the output: {shift}"
  );

  let mut config = tiny_config();
  config.type_vocab_size = 0;
  let empty_varmap = VarMap::new();
  let vb = VarBuilder::from_varmap(&empty_varmap, DType::F32, &Device::Cpu);
  GteNewModel::load(config, vb).unwrap();
  let data = empty_varmap.data().lock().unwrap();
  assert!(!data.contains_key("embeddings.token_type_embeddings.weight"));
}

/// Harvested from the embeddings and layer loads directly, so it stays a fixed yardstick when the prefix probe is mutated.
fn all_tensors(prefix: &str) -> HashMap<String, Tensor> {
  let config = tiny_config();
  let varmap = VarMap::new();
  let vb = VarBuilder::from_varmap(&varmap, DType::F32, &Device::Cpu);
  GteNewEmbeddings::load(&config, vb.clone()).unwrap();
  for i in 0..config.num_hidden_layers {
    GteNewLayer::load(&config, vb.pp(format!("encoder.layer.{i}"))).unwrap();
  }
  varmap
    .data()
    .lock()
    .unwrap()
    .iter()
    .map(|(name, var)| (format!("{prefix}{name}"), var.as_tensor().clone()))
    .collect()
}

/// `VarMap::get` auto-creates missing names, so only an exact-match `VarBuilder::from_tensors` can fail a wrong prefix.
#[test]
fn the_prefix_probe_finds_both_shapes() {
  let prefixed = VarBuilder::from_tensors(all_tensors("new."), DType::F32, &Device::Cpu);
  assert!(GteNewModel::load(tiny_config(), prefixed).is_ok());
  let bare = VarBuilder::from_tensors(all_tensors(""), DType::F32, &Device::Cpu);
  assert!(GteNewModel::load(tiny_config(), bare).is_ok());
  let wrong_prefix = VarBuilder::from_tensors(all_tensors("other."), DType::F32, &Device::Cpu);
  assert!(GteNewModel::load(tiny_config(), wrong_prefix).is_err());
}

/// Rebuilds the one-layer forward from its parts to pin the sublayer order and the absence of a final norm.
/// The `mlp_ln` overrides matter: a default norm re-normalizing a standardized signal is a no-op that would hide an extra norm.
#[test]
fn the_forward_is_one_post_norm_block_with_no_final_norm() {
  let mut config = tiny_config();
  config.num_hidden_layers = 1;
  let mut varmap = VarMap::new();
  let vb = VarBuilder::from_varmap(&varmap, DType::F32, &Device::Cpu);
  let model = GteNewModel::load(config.clone(), vb.clone()).unwrap();
  varmap
    .set_one(
      "encoder.layer.0.mlp_ln.weight",
      Tensor::new(&[2f32; 16], &Device::Cpu).unwrap(),
    )
    .unwrap();
  varmap
    .set_one(
      "encoder.layer.0.mlp_ln.bias",
      Tensor::new(&[1f32; 16], &Device::Cpu).unwrap(),
    )
    .unwrap();

  let ids = vec![3u32, 5, 7, 9];
  let out = run(&model, std::slice::from_ref(&ids));

  let input_ids = Tensor::from_vec(ids, (1, 4), &Device::Cpu).unwrap();
  let hidden = GteNewEmbeddings::load(&config, vb.clone())
    .unwrap()
    .forward(&input_ids)
    .unwrap();
  let mask = AttentionMask::padding(
    &Tensor::ones((1, 4), DType::U32, &Device::Cpu).unwrap(),
    &[4],
    config.num_attention_heads,
    config.head_dim(),
    DType::F32,
  )
  .unwrap();
  let layer_vb = vb.pp("encoder.layer.0");
  let attended = GteNewAttention::load(&config, layer_vb.clone())
    .unwrap()
    .forward(&hidden, &mask, model.rope())
    .unwrap();
  let normed = LayerNorm::load(layer_vb.pp("attn_ln"), 16, config.layer_norm_eps)
    .unwrap()
    .forward(&(attended + &hidden).unwrap())
    .unwrap();
  let projected = GteNewMlp::load(&config, layer_vb.clone())
    .unwrap()
    .forward(&normed)
    .unwrap();
  let expected = LayerNorm::load(layer_vb.pp("mlp_ln"), 16, config.layer_norm_eps)
    .unwrap()
    .forward(&(projected + &normed).unwrap())
    .unwrap()
    .to_vec3::<f32>()
    .unwrap();

  let shift = max_shift(&out, &expected);
  assert!(
    shift < 1e-4,
    "encoder output does not match a hand run ending at the last mlp_ln: {shift}"
  );
}

#[test]
fn the_length_is_the_configured_one() {
  let (_varmap, model) = tiny_model();
  assert_eq!(model.max_position_embeddings(), 64);
}

#[test]
fn an_unsupported_config_is_refused_at_load() {
  let varmap = VarMap::new();
  let vb = VarBuilder::from_varmap(&varmap, DType::F32, &Device::Cpu);
  let mut config = tiny_config();
  config.unpad_inputs = true;
  assert!(matches!(
    GteNewModel::load(config, vb).unwrap_err(),
    crate::error::CanddingError::UnsupportedArchitecture(_)
  ));
}