//! Tensor data types.
//!
//! T2 only needs `F32`; `BF16` was added in T12 for mixed-precision training
//! (bf16 linears / activations, fp32 master weights — see
//! `docs/11-bf16-mixed-precision.md`). The enum + `TensorDType` trait keep call
//! sites dtype-polymorphic.

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum DType {
    F32,
    /// bfloat16: 1 sign / 8 exponent / 7 mantissa. Same exponent range as f32
    /// (so no loss scaling needed), ~2-3 decimal digits. The T12 AMP compute type.
    BF16,
    /// 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::BF16 => 2,
            DType::I32 => 4,
        }
    }

    pub fn name(self) -> &'static str {
        match self {
            DType::F32 => "f32",
            DType::BF16 => "bf16",
            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 half::bf16 {
    const DTYPE: DType = DType::BF16;
    fn to_f64(self) -> f64 {
        self.to_f64()
    }
    fn from_f64(v: f64) -> Self {
        half::bf16::from_f64(v)
    }
}

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
    }
}
