//! Tensor data types.
//!
//! T2 only needs `F32`, but the enum + `TensorDType` trait are structured so
//! half-precision types (F16/BF16) can be added later (T7 mixed precision)
//! without touching call sites.

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum DType {
    F32,
    /// 32-bit signed integers. Used for cross-entropy targets (token ids).
    I32,
}

impl DType {
    pub fn size_bytes(self) -> usize {
        match self {
            DType::F32 => 4,
            DType::I32 => 4,
        }
    }

    pub fn name(self) -> &'static str {
        match self {
            DType::F32 => "f32",
            DType::I32 => "i32",
        }
    }
}

impl std::fmt::Display for DType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(self.name())
    }
}

/// Rust types that can back a tensor. Gives `from_slice`/`as_slice` type safety.
pub trait TensorDType: Copy + Send + Sync + 'static {
    const DTYPE: DType;
    fn to_f64(self) -> f64;
    fn from_f64(v: f64) -> Self;
}

impl TensorDType for f32 {
    const DTYPE: DType = DType::F32;
    fn to_f64(self) -> f64 {
        self as f64
    }
    fn from_f64(v: f64) -> Self {
        v as f32
    }
}

impl TensorDType for i32 {
    const DTYPE: DType = DType::I32;
    fn to_f64(self) -> f64 {
        self as f64
    }
    fn from_f64(v: f64) -> Self {
        v as i32
    }
}
