//! Native type replacements for candle types.
//!
//! These types provide the same API surface that downstream crates previously
//! used from Candle, backed by cudarc primitives and `GpuTensor`.
//!
//! | Old candle type        | Native replacement     |
//! |------------------------|------------------------|
//! | `candle_core::Device`  | [`NativeDevice`]       |
//! | `candle_core::DType`   | [`NativeDType`]        |
//! | `candle_core::Tensor`  | [`NativeTensor`]       |

use serde::{Deserialize, Serialize};
use std::sync::Arc;

use crate::MLError;

// cudarc is a direct dependency.
#[cfg(feature = "cuda")]
use cudarc;

// ---------------------------------------------------------------------------
// NativeDevice
// ---------------------------------------------------------------------------

/// Device selector replacing `candle_core::Device`.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum NativeDevice {
    /// Host CPU.
    Cpu,
    /// CUDA GPU with the given ordinal (0-indexed).
    Cuda(usize),
}

impl NativeDevice {
    /// Returns `true` if this is a CUDA device.
    pub fn is_cuda(&self) -> bool {
        matches!(self, NativeDevice::Cuda(_))
    }

    /// Returns `true` if this is the CPU device.
    pub fn is_cpu(&self) -> bool {
        matches!(self, NativeDevice::Cpu)
    }

    /// Get the CUDA device ordinal, or `None` for CPU.
    pub fn cuda_ordinal(&self) -> Option<usize> {
        match self {
            NativeDevice::Cuda(id) => Some(*id),
            NativeDevice::Cpu => None,
        }
    }

    /// Check whether two devices refer to the same physical device.
    pub fn same_device(&self, other: &NativeDevice) -> bool {
        self == other
    }
}

impl std::fmt::Display for NativeDevice {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            NativeDevice::Cpu => write!(f, "cpu"),
            NativeDevice::Cuda(id) => write!(f, "cuda:{}", id),
        }
    }
}

// ---------------------------------------------------------------------------
// NativeDType
// ---------------------------------------------------------------------------

/// Data type enum replacing `candle_core::DType`.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum NativeDType {
    /// 32-bit IEEE 754 float. Default compute type.
    F32,
    /// 16-bit brain float. Preferred for H100/A100 tensor core training.
    BF16,
    /// 16-bit IEEE 754 float. Used for inference compression.
    F16,
    /// 64-bit IEEE 754 float. Used at precision boundaries only.
    F64,
}

impl NativeDType {
    /// Size in bytes of a single element of this type.
    pub fn size_in_bytes(&self) -> usize {
        match self {
            NativeDType::F32 => 4,
            NativeDType::BF16 => 2,
            NativeDType::F16 => 2,
            NativeDType::F64 => 8,
        }
    }

    /// Returns the training dtype: BF16 on CUDA, F32 on CPU.
    pub fn training_dtype(device: &NativeDevice) -> Self {
        match device {
            NativeDevice::Cuda(_) => NativeDType::BF16,
            NativeDevice::Cpu => NativeDType::F32,
        }
    }

    /// Whether this type is a half-precision format (BF16 or F16).
    pub fn is_half(&self) -> bool {
        matches!(self, NativeDType::BF16 | NativeDType::F16)
    }

    /// Memory multiplier relative to F32 (e.g. 0.5 for half types).
    pub fn memory_multiplier(&self) -> f64 {
        match self {
            NativeDType::F32 => 1.0,
            NativeDType::BF16 | NativeDType::F16 => 0.5,
            NativeDType::F64 => 2.0,
        }
    }
}

impl std::fmt::Display for NativeDType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            NativeDType::F32 => write!(f, "f32"),
            NativeDType::BF16 => write!(f, "bf16"),
            NativeDType::F16 => write!(f, "f16"),
            NativeDType::F64 => write!(f, "f64"),
        }
    }
}

// ---------------------------------------------------------------------------
// NativeTensor
// ---------------------------------------------------------------------------

/// GPU-resident tensor replacing `candle_core::Tensor`.
///
/// Wraps `GpuTensor` from cuda_autograd with metadata (dtype, device).
#[cfg(feature = "cuda")]
pub struct NativeTensor {
    /// The underlying GPU tensor.
    inner: crate::cuda_autograd::GpuTensor,
    /// Logical data type (metadata for checkpoint/serialization; actual GPU data is always f32).
    dtype: NativeDType,
    /// Device this tensor lives on.
    device: NativeDevice,
    /// CUDA stream for GPU operations.
    stream: Arc<cudarc::driver::CudaStream>,
}

#[cfg(feature = "cuda")]
impl std::fmt::Debug for NativeTensor {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "NativeTensor {{ shape: {:?}, dtype: {}, device: {}, numel: {} }}",
            self.inner.shape(),
            self.dtype,
            self.device,
            self.inner.numel(),
        )
    }
}

#[cfg(feature = "cuda")]
impl NativeTensor {
    /// Create a `NativeTensor` from an existing `GpuTensor`.
    pub fn from_gpu_tensor(
        inner: crate::cuda_autograd::GpuTensor,
        dtype: NativeDType,
        device: NativeDevice,
        stream: Arc<cudarc::driver::CudaStream>,
    ) -> Self {
        Self {
            inner,
            dtype,
            device,
            stream,
        }
    }

    /// Allocate a zero-filled tensor of the given shape.
    pub fn zeros(
        shape: &[usize],
        dtype: NativeDType,
        device: NativeDevice,
        stream: &Arc<cudarc::driver::CudaStream>,
    ) -> Result<Self, MLError> {
        let inner = crate::cuda_autograd::GpuTensor::zeros(shape, stream)?;
        Ok(Self {
            inner,
            dtype,
            device,
            stream: Arc::clone(stream),
        })
    }

    /// Allocate a tensor filled with ones.
    pub fn ones(
        shape: &[usize],
        dtype: NativeDType,
        device: NativeDevice,
        stream: &Arc<cudarc::driver::CudaStream>,
    ) -> Result<Self, MLError> {
        let n: usize = shape.iter().product();
        let host = vec![1.0_f32; n];
        let inner = crate::cuda_autograd::GpuTensor::from_host(&host, shape.to_vec(), stream)?;
        Ok(Self {
            inner,
            dtype,
            device,
            stream: Arc::clone(stream),
        })
    }

    /// Create a tensor from a host f32 slice.
    pub fn from_slice(
        data: &[f32],
        shape: &[usize],
        dtype: NativeDType,
        device: NativeDevice,
        stream: &Arc<cudarc::driver::CudaStream>,
    ) -> Result<Self, MLError> {
        let inner =
            crate::cuda_autograd::GpuTensor::from_host(data, shape.to_vec(), stream)?;
        Ok(Self {
            inner,
            dtype,
            device,
            stream: Arc::clone(stream),
        })
    }

    /// Create a tensor from a host `Vec<f32>`.
    pub fn from_vec(
        data: Vec<f32>,
        shape: &[usize],
        dtype: NativeDType,
        device: NativeDevice,
        stream: &Arc<cudarc::driver::CudaStream>,
    ) -> Result<Self, MLError> {
        Self::from_slice(&data, shape, dtype, device, stream)
    }

    /// Tensor dimensions.
    pub fn dims(&self) -> &[usize] {
        self.inner.shape()
    }

    /// Alias for `dims()`.
    pub fn shape(&self) -> &[usize] {
        self.inner.shape()
    }

    /// Number of dimensions.
    pub fn rank(&self) -> usize {
        self.inner.ndim()
    }

    /// Total number of elements.
    pub fn elem_count(&self) -> usize {
        self.inner.numel()
    }

    /// Logical data type.
    pub fn dtype(&self) -> NativeDType {
        self.dtype
    }

    /// Device this tensor resides on.
    pub fn device(&self) -> &NativeDevice {
        &self.device
    }

    /// Reference to the CUDA stream.
    pub fn stream(&self) -> &Arc<cudarc::driver::CudaStream> {
        &self.stream
    }

    /// Download tensor data to a flat `Vec<f32>`.
    pub fn to_vec1(&self) -> Result<Vec<f32>, MLError> {
        self.inner.to_host(&self.stream)
    }

    /// Download tensor data as a 2D `Vec<Vec<f32>>`.
    pub fn to_vec2(&self) -> Result<Vec<Vec<f32>>, MLError> {
        let shape = self.inner.shape();
        if shape.len() != 2 {
            return Err(MLError::DimensionMismatch {
                expected: 2,
                actual: shape.len(),
            });
        }
        let rows = shape[0];
        let cols = shape[1];
        let flat = self.inner.to_host(&self.stream)?;
        let mut result = Vec::with_capacity(rows);
        for r in 0..rows {
            let start = r * cols;
            let end = start + cols;
            let row = flat.get(start..end).ok_or_else(|| {
                MLError::DimensionMismatch {
                    expected: end,
                    actual: flat.len(),
                }
            })?;
            result.push(row.to_vec());
        }
        Ok(result)
    }

    /// Reshape without copying data.
    pub fn reshape(self, new_shape: Vec<usize>) -> Result<Self, MLError> {
        let stream = self.stream;
        let dtype = self.dtype;
        let device = self.device;
        let inner = self.inner.reshape(new_shape)?;
        Ok(Self {
            inner,
            dtype,
            device,
            stream,
        })
    }

    /// Borrow the underlying `GpuTensor`.
    pub fn gpu_tensor(&self) -> &crate::cuda_autograd::GpuTensor {
        &self.inner
    }

    /// Mutably borrow the underlying `GpuTensor`.
    pub fn gpu_tensor_mut(&mut self) -> &mut crate::cuda_autograd::GpuTensor {
        &mut self.inner
    }

    /// Consume this wrapper and return the inner `GpuTensor`.
    pub fn into_gpu_tensor(self) -> crate::cuda_autograd::GpuTensor {
        self.inner
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    #[test]
    fn test_native_device_properties() {
        let cpu = NativeDevice::Cpu;
        assert!(cpu.is_cpu());
        assert!(!cpu.is_cuda());
        assert_eq!(cpu.cuda_ordinal(), None);
        assert_eq!(cpu.to_string(), "cpu");

        let gpu = NativeDevice::Cuda(0);
        assert!(!gpu.is_cpu());
        assert!(gpu.is_cuda());
        assert_eq!(gpu.cuda_ordinal(), Some(0));
        assert_eq!(gpu.to_string(), "cuda:0");
    }

    #[test]
    fn test_native_device_equality() {
        assert_eq!(NativeDevice::Cpu, NativeDevice::Cpu);
        assert_eq!(NativeDevice::Cuda(0), NativeDevice::Cuda(0));
        assert_ne!(NativeDevice::Cpu, NativeDevice::Cuda(0));
        assert_ne!(NativeDevice::Cuda(0), NativeDevice::Cuda(1));
        assert!(NativeDevice::Cpu.same_device(&NativeDevice::Cpu));
        assert!(!NativeDevice::Cpu.same_device(&NativeDevice::Cuda(0)));
    }

    #[test]
    fn test_native_dtype_properties() {
        assert_eq!(NativeDType::F32.size_in_bytes(), 4);
        assert_eq!(NativeDType::BF16.size_in_bytes(), 2);
        assert_eq!(NativeDType::F16.size_in_bytes(), 2);
        assert_eq!(NativeDType::F64.size_in_bytes(), 8);

        assert!(!NativeDType::F32.is_half());
        assert!(NativeDType::BF16.is_half());
        assert!(NativeDType::F16.is_half());
        assert!(!NativeDType::F64.is_half());

        assert!((NativeDType::F32.memory_multiplier() - 1.0).abs() < f64::EPSILON);
        assert!((NativeDType::BF16.memory_multiplier() - 0.5).abs() < f64::EPSILON);
        assert!((NativeDType::F64.memory_multiplier() - 2.0).abs() < f64::EPSILON);
    }

    #[test]
    fn test_native_dtype_training_dtype() {
        assert_eq!(
            NativeDType::training_dtype(&NativeDevice::Cuda(0)),
            NativeDType::BF16
        );
        assert_eq!(
            NativeDType::training_dtype(&NativeDevice::Cpu),
            NativeDType::F32
        );
    }

    #[test]
    fn test_native_dtype_display() {
        assert_eq!(NativeDType::F32.to_string(), "f32");
        assert_eq!(NativeDType::BF16.to_string(), "bf16");
        assert_eq!(NativeDType::F16.to_string(), "f16");
        assert_eq!(NativeDType::F64.to_string(), "f64");
    }
}
