candding 0.1.0

use std::collections::HashMap;
use std::path::PathBuf;

use candle_core::DType;
use candle_core::Device;
use candle_core::Tensor;
use tempfile::tempdir;

use super::*;

fn pickle_fixture(name: &str) -> PathBuf {
  PathBuf::from(env!("CARGO_MANIFEST_DIR"))
    .join("tests/fixtures/pickle")
    .join(name)
}

/// One tensor per stored dtype, a scalar and an empty one, each loaded back through `var_builder` at every dtype.
#[test]
fn the_header_size_is_what_the_loaded_tensors_take() {
  let dir = tempdir().unwrap();
  let path = dir.path().join("model.safetensors");
  let tensors = HashMap::from([
    (
      "f32".to_string(),
      Tensor::zeros((3, 5), DType::F32, &Device::Cpu).unwrap(),
    ),
    (
      "f16".to_string(),
      Tensor::zeros((2, 7, 3), DType::F16, &Device::Cpu).unwrap(),
    ),
    (
      "bf16".to_string(),
      Tensor::zeros(11, DType::BF16, &Device::Cpu).unwrap(),
    ),
    (
      "scalar".to_string(),
      Tensor::new(1.5f32, &Device::Cpu).unwrap(),
    ),
    (
      "empty".to_string(),
      Tensor::zeros((0, 4), DType::F32, &Device::Cpu).unwrap(),
    ),
  ]);
  candle_core::safetensors::save(&tensors, &path).unwrap();
  for dtype in [DType::F32, DType::F16, DType::BF16] {
    let files = std::slice::from_ref(&path);
    let vb = unsafe { crate::weights::var_builder(files, dtype, &Device::Cpu).unwrap() };
    let loaded: usize = tensors
      .keys()
      .map(|name| {
        let tensor = vb.get_unchecked(name).unwrap();
        tensor.elem_count() * tensor.dtype().size_in_bytes()
      })
      .sum();
    assert_eq!(
      resident_bytes(&path, dtype).unwrap(),
      loaded as u64,
      "{dtype:?}"
    );
  }
}

/// Only the header is read: a file that declares a tensor and holds none of its bytes is still sized.
#[test]
fn a_header_alone_is_enough_and_metadata_is_not_a_tensor() {
  let dir = tempdir().unwrap();
  let path = dir.path().join("model.safetensors");
  let header = br#"{"__metadata__":{"format":"pt"},"w":{"dtype":"BF16","shape":[1024,1024,1024],"data_offsets":[0,2147483648]}}"#;
  let mut bytes = (header.len() as u64).to_le_bytes().to_vec();
  bytes.extend_from_slice(header);
  std::fs::write(&path, bytes).unwrap();
  assert_eq!(resident_bytes(&path, DType::F32).unwrap(), 4 << 30);
  assert_eq!(resident_bytes(&path, DType::F16).unwrap(), 2 << 30);
}

#[test]
fn a_truncated_or_oversized_header_is_refused_by_name() {
  let dir = tempdir().unwrap();
  let path = dir.path().join("model.safetensors");
  for bytes in [
    vec![1u8, 2, 3],
    [64u64.to_le_bytes().as_slice(), b"{}"].concat(),
    u64::MAX.to_le_bytes().to_vec(),
  ] {
    std::fs::write(&path, &bytes).unwrap();
    let err = resident_bytes(&path, DType::F32).unwrap_err().to_string();
    assert!(
      err.contains("model.safetensors") && err.contains("safetensors header"),
      "{err}"
    );
  }
}

/// `two_tensors.bin` stores float32 and `linear_head.bin` float16, so each scales its file size from its own dtype.
#[test]
fn a_pickle_scales_its_file_size_from_the_dtype_it_stores() {
  let f32_pickle = pickle_fixture("two_tensors.bin");
  let f16_pickle = pickle_fixture("linear_head.bin");
  let f32_len = std::fs::metadata(&f32_pickle).unwrap().len();
  let f16_len = std::fs::metadata(&f16_pickle).unwrap().len();
  assert_eq!(resident_bytes(&f32_pickle, DType::F32).unwrap(), f32_len);
  assert_eq!(
    resident_bytes(&f32_pickle, DType::F16).unwrap(),
    f32_len / 2
  );
  assert_eq!(
    resident_bytes(&f16_pickle, DType::F32).unwrap(),
    f16_len * 2
  );
  assert_eq!(resident_bytes(&f16_pickle, DType::BF16).unwrap(), f16_len);
}

#[test]
fn a_npy_file_scales_its_size_from_the_dtype_its_header_names() {
  let dir = tempdir().unwrap();
  let path = dir.path().join("matrix.npy");
  let dict = b"{'descr': '<f8', 'fortran_order': False, 'shape': (3,), }\n";
  let mut bytes = b"\x93NUMPY\x01\x00".to_vec();
  bytes.extend_from_slice(&(dict.len() as u16).to_le_bytes());
  bytes.extend_from_slice(dict);
  bytes.extend_from_slice(&[0u8; 24]);
  std::fs::write(&path, &bytes).unwrap();
  assert_eq!(
    resident_bytes(&path, DType::F32).unwrap(),
    bytes.len() as u64 / 2
  );
}

#[test]
fn scaling_converts_between_element_widths_without_overflow() {
  assert_eq!(scale(10, DType::BF16, DType::F32), 20);
  assert_eq!(scale(10, DType::F32, DType::F16), 5);
  assert_eq!(scale(u64::MAX, DType::F16, DType::F32), u64::MAX);
}