//! Element type traits for tensor data.
//!
//! These traits define the numeric types that can be stored in tensors.
//! Each backend specifies which element types it supports.

use half::{bf16, f16};
// num-traits import for potential future use
#[allow(unused_imports)]
use num_traits::ToPrimitive;
use serde::{Deserialize, Serialize};
use std::fmt::{Debug, Display};

/// Trait for floating-point element types.
///
/// Supported types:
/// - `f32` - Standard single precision
/// - `f64` - Double precision
/// - `f16` - Half precision (IEEE 754)
/// - `bf16` - Brain floating point (Google format)
///
/// # Performance Notes
///
/// - Use `f32` as the default for training
/// - Use `f16`/`bf16` for inference to reduce memory
/// - `bf16` has better dynamic range than `f16` but lower precision
pub trait FloatElement:
    Clone
    + Copy
    + Debug
    + Default
    + Display
    + PartialEq
    + PartialOrd
    + Send
    + Sync
    + Serialize
    + for<'de> Deserialize<'de>
    + 'static
{
    /// The name of this element type.
    fn name() -> &'static str;

    /// Number of bytes per element.
    fn bytes() -> usize;

    /// Create from f64.
    fn from_f64(val: f64) -> Self;

    /// Convert to f64.
    fn to_f64(self) -> f64;

    /// Create from f32.
    fn from_f32(val: f32) -> Self;

    /// Convert to f32.
    fn to_f32(self) -> f32;

    /// Zero value.
    fn zero() -> Self;

    /// One value.
    fn one() -> Self;

    /// Minimum representable value.
    fn min_value() -> Self;

    /// Maximum representable value.
    fn max_value() -> Self;

    /// Not a number.
    fn nan() -> Self;

    /// Positive infinity.
    fn infinity() -> Self;

    /// Negative infinity.
    fn neg_infinity() -> Self;

    /// Check if NaN.
    fn is_nan(self) -> bool;

    /// Check if infinite.
    fn is_infinite(self) -> bool;

    /// Epsilon value.
    fn epsilon() -> Self;
}

impl FloatElement for f32 {
    #[inline]
    fn name() -> &'static str {
        "f32"
    }
    #[inline]
    fn bytes() -> usize {
        4
    }
    #[inline]
    fn from_f64(val: f64) -> Self {
        val as f32
    }
    #[inline]
    fn to_f64(self) -> f64 {
        self as f64
    }
    #[inline]
    fn from_f32(val: f32) -> Self {
        val
    }
    #[inline]
    fn to_f32(self) -> f32 {
        self
    }
    #[inline]
    fn zero() -> Self {
        0.0
    }
    #[inline]
    fn one() -> Self {
        1.0
    }
    #[inline]
    fn min_value() -> Self {
        f32::MIN
    }
    #[inline]
    fn max_value() -> Self {
        f32::MAX
    }
    #[inline]
    fn nan() -> Self {
        f32::NAN
    }
    #[inline]
    fn infinity() -> Self {
        f32::INFINITY
    }
    #[inline]
    fn neg_infinity() -> Self {
        f32::NEG_INFINITY
    }
    #[inline]
    fn is_nan(self) -> bool {
        f32::is_nan(self)
    }
    #[inline]
    fn is_infinite(self) -> bool {
        f32::is_infinite(self)
    }
    #[inline]
    fn epsilon() -> Self {
        f32::EPSILON
    }
}

impl FloatElement for f64 {
    #[inline]
    fn name() -> &'static str {
        "f64"
    }
    #[inline]
    fn bytes() -> usize {
        8
    }
    #[inline]
    fn from_f64(val: f64) -> Self {
        val
    }
    #[inline]
    fn to_f64(self) -> f64 {
        self
    }
    #[inline]
    fn from_f32(val: f32) -> Self {
        val as f64
    }
    #[inline]
    fn to_f32(self) -> f32 {
        self as f32
    }
    #[inline]
    fn zero() -> Self {
        0.0
    }
    #[inline]
    fn one() -> Self {
        1.0
    }
    #[inline]
    fn min_value() -> Self {
        f64::MIN
    }
    #[inline]
    fn max_value() -> Self {
        f64::MAX
    }
    #[inline]
    fn nan() -> Self {
        f64::NAN
    }
    #[inline]
    fn infinity() -> Self {
        f64::INFINITY
    }
    #[inline]
    fn neg_infinity() -> Self {
        f64::NEG_INFINITY
    }
    #[inline]
    fn is_nan(self) -> bool {
        f64::is_nan(self)
    }
    #[inline]
    fn is_infinite(self) -> bool {
        f64::is_infinite(self)
    }
    #[inline]
    fn epsilon() -> Self {
        f64::EPSILON
    }
}

impl FloatElement for f16 {
    #[inline]
    fn name() -> &'static str {
        "f16"
    }
    #[inline]
    fn bytes() -> usize {
        2
    }
    #[inline]
    fn from_f64(val: f64) -> Self {
        f16::from_f64(val)
    }
    #[inline]
    fn to_f64(self) -> f64 {
        self.to_f64()
    }
    #[inline]
    fn from_f32(val: f32) -> Self {
        f16::from_f32(val)
    }
    #[inline]
    fn to_f32(self) -> f32 {
        self.to_f32()
    }
    #[inline]
    fn zero() -> Self {
        f16::ZERO
    }
    #[inline]
    fn one() -> Self {
        f16::ONE
    }
    #[inline]
    fn min_value() -> Self {
        f16::MIN
    }
    #[inline]
    fn max_value() -> Self {
        f16::MAX
    }
    #[inline]
    fn nan() -> Self {
        f16::NAN
    }
    #[inline]
    fn infinity() -> Self {
        f16::INFINITY
    }
    #[inline]
    fn neg_infinity() -> Self {
        f16::NEG_INFINITY
    }
    #[inline]
    fn is_nan(self) -> bool {
        self.is_nan()
    }
    #[inline]
    fn is_infinite(self) -> bool {
        self.is_infinite()
    }
    #[inline]
    fn epsilon() -> Self {
        f16::EPSILON
    }
}

impl FloatElement for bf16 {
    #[inline]
    fn name() -> &'static str {
        "bf16"
    }
    #[inline]
    fn bytes() -> usize {
        2
    }
    #[inline]
    fn from_f64(val: f64) -> Self {
        bf16::from_f64(val)
    }
    #[inline]
    fn to_f64(self) -> f64 {
        self.to_f64()
    }
    #[inline]
    fn from_f32(val: f32) -> Self {
        bf16::from_f32(val)
    }
    #[inline]
    fn to_f32(self) -> f32 {
        self.to_f32()
    }
    #[inline]
    fn zero() -> Self {
        bf16::ZERO
    }
    #[inline]
    fn one() -> Self {
        bf16::ONE
    }
    #[inline]
    fn min_value() -> Self {
        bf16::MIN
    }
    #[inline]
    fn max_value() -> Self {
        bf16::MAX
    }
    #[inline]
    fn nan() -> Self {
        bf16::NAN
    }
    #[inline]
    fn infinity() -> Self {
        bf16::INFINITY
    }
    #[inline]
    fn neg_infinity() -> Self {
        bf16::NEG_INFINITY
    }
    #[inline]
    fn is_nan(self) -> bool {
        self.is_nan()
    }
    #[inline]
    fn is_infinite(self) -> bool {
        self.is_infinite()
    }
    #[inline]
    fn epsilon() -> Self {
        bf16::EPSILON
    }
}

/// Trait for integer element types.
///
/// Supported types:
/// - `i8`, `i16`, `i32`, `i64` - Signed integers
/// - `u8`, `u16`, `u32`, `u64` - Unsigned integers
pub trait IntElement:
    Clone
    + Copy
    + Debug
    + Default
    + Display
    + PartialEq
    + Eq
    + PartialOrd
    + Ord
    + Send
    + Sync
    + Serialize
    + for<'de> Deserialize<'de>
    + 'static
{
    /// The name of this element type.
    fn name() -> &'static str;

    /// Number of bytes per element.
    fn bytes() -> usize;

    /// Create from i64.
    fn from_i64(val: i64) -> Self;

    /// Convert to i64.
    fn to_i64(self) -> i64;

    /// Zero value.
    fn zero() -> Self;

    /// One value.
    fn one() -> Self;

    /// Minimum value.
    fn min_value() -> Self;

    /// Maximum value.
    fn max_value() -> Self;
}

macro_rules! impl_int_element {
    ($t:ty, $name:expr) => {
        impl IntElement for $t {
            fn name() -> &'static str {
                $name
            }
            fn bytes() -> usize {
                std::mem::size_of::<$t>()
            }
            fn from_i64(val: i64) -> Self {
                val as $t
            }
            fn to_i64(self) -> i64 {
                self as i64
            }
            fn zero() -> Self {
                0
            }
            fn one() -> Self {
                1
            }
            fn min_value() -> Self {
                <$t>::MIN
            }
            fn max_value() -> Self {
                <$t>::MAX
            }
        }
    };
}

impl_int_element!(i8, "i8");
impl_int_element!(i16, "i16");
impl_int_element!(i32, "i32");
impl_int_element!(i64, "i64");
impl_int_element!(u8, "u8");
impl_int_element!(u16, "u16");
impl_int_element!(u32, "u32");
impl_int_element!(u64, "u64");

/// Trait for boolean element types.
///
/// Used for mask operations and comparisons.
pub trait BoolElement:
    Clone + Copy + Debug + Default + PartialEq + Eq + Send + Sync + 'static
{
    /// The name of this element type.
    fn name() -> &'static str;

    /// True value.
    fn true_val() -> Self;

    /// False value.
    fn false_val() -> Self;

    /// Convert to bool.
    fn to_bool(self) -> bool;

    /// Create from bool.
    fn from_bool(val: bool) -> Self;
}

impl BoolElement for bool {
    #[inline]
    fn name() -> &'static str {
        "bool"
    }
    #[inline]
    fn true_val() -> Self {
        true
    }
    #[inline]
    fn false_val() -> Self {
        false
    }
    #[inline]
    fn to_bool(self) -> bool {
        self
    }
    #[inline]
    fn from_bool(val: bool) -> Self {
        val
    }
}

/// Boolean stored as u8 (useful for GPU where bool isn't native).
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct BoolU8(pub u8);

impl BoolElement for BoolU8 {
    #[inline]
    fn name() -> &'static str {
        "bool_u8"
    }
    #[inline]
    fn true_val() -> Self {
        BoolU8(1)
    }
    #[inline]
    fn false_val() -> Self {
        BoolU8(0)
    }
    #[inline]
    fn to_bool(self) -> bool {
        self.0 != 0
    }
    #[inline]
    fn from_bool(val: bool) -> Self {
        BoolU8(u8::from(val))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_float_element_f32() {
        assert_eq!(f32::name(), "f32");
        assert_eq!(f32::bytes(), 4);
        assert_eq!(f32::from_f64(1.5), 1.5f32);
        assert!(f32::nan().is_nan());
    }

    #[test]
    fn test_float_element_bf16() {
        assert_eq!(bf16::name(), "bf16");
        assert_eq!(bf16::bytes(), 2);
        let val = bf16::from_f32(1.5);
        assert!((val.to_f32() - 1.5).abs() < 0.01);
    }

    #[test]
    fn test_int_element() {
        assert_eq!(i32::name(), "i32");
        assert_eq!(i32::bytes(), 4);
        assert_eq!(i32::from_i64(42), 42i32);
    }

    #[test]
    fn test_bool_element() {
        assert!(bool::true_val());
        assert!(!bool::false_val());
        assert_eq!(BoolU8::true_val(), BoolU8(1));
        assert!(BoolU8::true_val().to_bool());
    }
}