candding 0.1.0

use std::collections::HashMap;

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

use super::*;
use crate::error::CanddingError;

const HIDDEN: usize = 4;
const VOCAB: usize = 16;
/// The learned table's rows: `pad_token_id + 1 = 2` skipped by the numbering, 6 usable positions.
const ROWS: usize = 8;
const PAD: u32 = 1;

fn config(position_embedding_type: Option<&str>) -> XlmRobertaFlashConfig {
  XlmRobertaFlashConfig {
    vocab_size: VOCAB,
    hidden_size: HIDDEN,
    num_hidden_layers: 0,
    num_attention_heads: 1,
    intermediate_size: 1,
    max_position_embeddings: ROWS,
    position_embedding_type: position_embedding_type.map(str::to_string),
    pad_token_id: PAD as usize,
    ..tiny_config()
  }
}

/// Closed forms, so the oracle reads the same numbers the tensors hold and no two position rows are alike.
fn word_row(id: usize) -> [f32; HIDDEN] {
  let x = id as f32;
  [0.1 * x, -0.05 * x, 0.02 * x * x, 0.3]
}

fn position_row(row: usize) -> [f32; HIDDEN] {
  let x = row as f32;
  [x, 0.1 * x * x, -0.5 * x, 1.0 / (x + 1.0)]
}

const TOKEN_TYPE_ROW: [f32; HIDDEN] = [0.3, -0.2, 0.1, 0.05];

fn table(rows: usize, row: impl Fn(usize) -> [f32; HIDDEN]) -> Tensor {
  let values: Vec<f32> = (0..rows).flat_map(row).collect();
  Tensor::from_vec(values, (rows, HIDDEN), &Device::Cpu).unwrap()
}

fn tensors(with_table: bool) -> HashMap<String, Tensor> {
  let mut map = HashMap::new();
  map.insert(
    "embeddings.word_embeddings.weight".to_string(),
    table(VOCAB, word_row),
  );
  map.insert(
    "embeddings.token_type_embeddings.weight".to_string(),
    table(1, |_| TOKEN_TYPE_ROW),
  );
  if with_table {
    map.insert(
      "embeddings.position_embeddings.weight".to_string(),
      table(ROWS, position_row),
    );
  }
  map.insert(
    "emb_ln.weight".to_string(),
    Tensor::ones(HIDDEN, DType::F32, &Device::Cpu).unwrap(),
  );
  map.insert(
    "emb_ln.bias".to_string(),
    Tensor::zeros(HIDDEN, DType::F32, &Device::Cpu).unwrap(),
  );
  map
}

fn load(
  config: XlmRobertaFlashConfig,
  map: HashMap<String, Tensor>,
  declared: Option<PositionScheme>,
) -> Result<XlmRobertaFlashModel> {
  XlmRobertaFlashModel::load(
    config,
    VarBuilder::from_tensors(map, DType::F32, &Device::Cpu),
    declared,
  )
}

fn layer_norm(x: [f64; HIDDEN]) -> [f64; HIDDEN] {
  let mean = x.iter().sum::<f64>() / HIDDEN as f64;
  let variance = x.iter().map(|v| (v - mean).powi(2)).sum::<f64>() / HIDDEN as f64;
  x.map(|v| (v - mean) / (variance + 1e-5).sqrt())
}

fn oracle(id: u32, row: usize) -> [f64; HIDDEN] {
  let (word, position) = (word_row(id as usize), position_row(row));
  layer_norm(std::array::from_fn(|k| {
    word[k] as f64 + position[k] as f64 + TOKEN_TYPE_ROW[k] as f64
  }))
}

/// Row `cumsum(ids != pad) * (ids != pad) + pad`: a real `<pad>` reads row `pad` and shifts every later token down one.
#[test]
fn the_absolute_scheme_adds_the_row_each_position_is_numbered_to() {
  let model = load(
    config(Some("absolute")),
    tensors(true),
    Some(PositionScheme::Absolute),
  )
  .unwrap();
  let rows = [vec![0u32, 5, 9, 2], vec![0u32, 7], vec![0u32, PAD, 9, 2]];
  let table_rows = [[2, 3, 4, 5], [2, 3, 1, 1], [2, 1, 3, 4]];
  let out = run(&model, &rows);
  for (b, row) in rows.iter().enumerate() {
    for (p, actual) in out[b].iter().enumerate() {
      let id = row.get(p).copied().unwrap_or(PAD);
      let table_row = table_rows[b][p];
      let expected = oracle(id, table_row);
      for (k, (a, e)) in actual.iter().zip(expected).enumerate() {
        assert!(
          (*a as f64 - e).abs() < 1e-5,
          "row {b} position {p} dim {k}: {a} against table row {table_row}'s {e}"
        );
      }
    }
  }
}

#[test]
fn the_absolute_scheme_requires_the_position_table() {
  let err = load(
    config(Some("absolute")),
    tensors(false),
    Some(PositionScheme::Absolute),
  )
  .unwrap_err();
  assert!(matches!(err, CanddingError::Candle(_)), "{err}");
}

/// 6 of 8 rows here and 1024 of 1026 for the reranker; the rotary scheme indexes no table.
#[test]
fn the_absolute_scheme_reports_the_tables_sequence_limit() {
  let absolute = load(
    config(Some("absolute")),
    tensors(true),
    Some(PositionScheme::Absolute),
  )
  .unwrap();
  assert_eq!(absolute.max_position_embeddings(), ROWS - PAD as usize - 1);
  let rotary = load(
    config(Some("rotary")),
    tensors(false),
    Some(PositionScheme::Rotary),
  )
  .unwrap();
  assert_eq!(rotary.max_position_embeddings(), ROWS);
}

#[test]
fn a_scheme_the_config_does_not_name_is_refused_before_any_tensor_is_read() {
  for (declared, named) in [
    (Some(PositionScheme::Absolute), Some("rotary")),
    (Some(PositionScheme::Rotary), Some("absolute")),
    (None, Some("absolute")),
    (Some(PositionScheme::Absolute), None),
  ] {
    let err = load(config(named), HashMap::new(), declared).unwrap_err();
    assert!(
      matches!(err, CanddingError::PositionSchemeDisagreesWithConfig { .. }),
      "{declared:?} against {named:?}: {err}"
    );
  }
}

#[test]
fn a_scheme_the_family_does_not_implement_is_refused_at_load() {
  let err = load(
    config(Some("alibi")),
    tensors(true),
    Some(PositionScheme::Absolute),
  )
  .unwrap_err();
  assert!(
    matches!(err, CanddingError::UnsupportedArchitecture(ref m) if m.contains("alibi")),
    "{err}"
  );
}