candding 0.1.0

use candle_core::DType;
use candle_core::Device;
use candle_core::Tensor;

use super::*;
use crate::embed::synthetic;
use crate::embed::synthetic::HIDDEN;
use crate::models::MultiVectorHead;
use crate::multi_vector::MaskedRowFill;
use crate::multi_vector::ScoreReduction;
use crate::registry::Backend;
use crate::registry::BackendSupport;
use crate::registry::DtypeBackendSupport;
use crate::registry::Family;
use crate::registry::Support;
use crate::registry::WeightsFormat;

/// Reads the private field: `batch_size` has no public getter.
#[test]
fn default_batch_size_is_eight() {
  let builder = MultiVectorTextEmbeddingBuilder::new("org/unregistered");
  assert_eq!(builder.batch_size, 8);
}

fn stanford_descriptor() -> MultiVectorDescriptor {
  MultiVectorDescriptor {
    id: "synthetic/zero-colbert".to_string(),
    family: Family::Bert,
    head: MultiVectorHead::StanfordColbert,
    default_max_length: 32,
    keeps_cls_row: true,
    projection_in: HIDDEN,
    projection_out: 2,
    projection_bias: false,
    normalize: true,
    score_reduction: ScoreReduction::Sum,
    query_marker: Some("[unused0] ".to_string()),
    document_marker: Some("[unused1] ".to_string()),
    query_length: Some(8),
    attends_to_mask_tokens: false,
    masked_document_row_fill: MaskedRowFill::DtypeMinimum,
    position_scheme: None,
    backends: BackendSupport::untested(),
    dtype_backends: DtypeBackendSupport::untested(),
    license: "unknown".to_string(),
    size_gb: 0.0,
    phase: u8::MAX,
    weights: WeightsFormat::Safetensors,
    gated: false,
  }
}

#[test]
fn build_loads_the_weights_in_the_dtype_it_resolves() {
  let mut tensors = synthetic::zero_bert("", 19);
  tensors.insert("linear.weight".to_string(), synthetic::zeros(&[2, HIDDEN]));
  let config = synthetic::bert_config(19);
  let dir = synthetic::checkpoint(&config, "tiny_tokenizer_markers.json", &tensors);
  for (requested, expected) in [(None, DType::F32), (Some(DType::F16), DType::F16)] {
    let mut builder = MultiVectorTextEmbeddingBuilder::new(dir.path().to_str().unwrap())
      .descriptor(stanford_descriptor())
      .device(Device::Cpu);
    if let Some(dtype) = requested {
      builder = builder.dtype(dtype);
    }
    let model = builder.build().unwrap();
    assert_eq!(model.dtype(), expected, "requested {requested:?}");
    assert_eq!(model.weights_dtype(), expected, "requested {requested:?}");
  }
}

/// Every hidden state is the LayerNorm bias `[3, 4, 0, 0]` and the projection keeps two components, so every row is `[3, 4]`.
#[test]
fn rows_are_unit_length_unless_normalize_is_turned_off() {
  let mut tensors = synthetic::zero_bert("", 19);
  let bias = Tensor::new(&[3f32, 4.0, 0.0, 0.0], &Device::Cpu).unwrap();
  tensors.insert("embeddings.LayerNorm.bias".to_string(), bias);
  let keep_two = Tensor::new(&[[1f32, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0]], &Device::Cpu).unwrap();
  tensors.insert("linear.weight".to_string(), keep_two);
  let config = synthetic::bert_config(19);
  let dir = synthetic::checkpoint(&config, "tiny_tokenizer_markers.json", &tensors);
  for (normalize, expected) in [(None, [0.6, 0.8]), (Some(false), [3.0, 4.0])] {
    let mut builder = MultiVectorTextEmbeddingBuilder::new(dir.path().to_str().unwrap())
      .descriptor(stanford_descriptor())
      .device(Device::Cpu);
    if let Some(normalize) = normalize {
      builder = builder.normalize(normalize);
    }
    let embedding = builder
      .build()
      .unwrap()
      .passage_embed(&["a b"], None)
      .unwrap()
      .remove(0);
    assert!(embedding.row_count() > 0);
    for row in embedding.rows() {
      let close = row.iter().zip(expected).all(|(a, e)| (a - e).abs() < 1e-6);
      assert!(close, "{normalize:?}: {row:?} against {expected:?}");
    }
  }
}

/// The spec is an empty directory, so a refusal moved below `hub::resolve` or `loader::discover` surfaces as `MissingFile`.
#[test]
fn a_refused_dtype_is_refused_before_anything_is_fetched() {
  let dir = tempfile::tempdir().unwrap();
  let refused = MultiVectorTextEmbeddingBuilder::new(dir.path().to_str().unwrap())
    .descriptor(stanford_descriptor())
    .device(Device::Cpu)
    .dtype(DType::BF16)
    .build()
    .err();
  assert!(
    matches!(
      refused,
      Some(CanddingError::UnsupportedOnDevice {
        backend: Backend::Cpu,
        dtype: DType::BF16,
        ..
      })
    ),
    "{refused:?}"
  );
}

/// The dtype check reads the device's own cell: refused where only the CPU's refuses, and past the check, into the
/// empty directory, where only the accelerators' do.
#[test]
fn the_dtype_check_reads_the_device_s_own_cell() {
  let (v, u) = (Support::Verified, Support::Unsupported);
  for (f16, refused) in [
    (
      BackendSupport {
        cpu: u,
        metal: v,
        cuda: v,
      },
      true,
    ),
    (
      BackendSupport {
        cpu: v,
        metal: u,
        cuda: u,
      },
      false,
    ),
  ] {
    let dir = tempfile::tempdir().unwrap();
    let mut descriptor = stanford_descriptor();
    descriptor.dtype_backends.f16 = f16;
    let err = MultiVectorTextEmbeddingBuilder::new(dir.path().to_str().unwrap())
      .descriptor(descriptor)
      .device(Device::Cpu)
      .dtype(DType::F16)
      .build()
      .err()
      .expect("an empty directory loads nothing");
    let on_the_cpu = matches!(
      err,
      CanddingError::UnsupportedOnDevice {
        backend: Backend::Cpu,
        dtype: DType::F16,
        ..
      }
    );
    assert_eq!(on_the_cpu, refused, "{f16:?}: {err}");
    assert!(
      !matches!(err, CanddingError::UnsupportedDtype { .. }),
      "{f16:?}: {err}"
    );
  }
}