//! Precision conversion utilities (Candle-free).
//!
//! Precision types are now metadata-only. Actual GPU data is always F32 in
//! `cuda_autograd::GpuTensor`. Precision conversion (BF16/F16 storage) is
//! handled at the CUDA kernel level via mixed-precision patterns.

use serde::{Deserialize, Serialize};

use crate::MLError;

/// Precision type for model weights and activations.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum PrecisionType {
    /// 32-bit floating point (baseline).
    Float32,
    /// 16-bit floating point (50% memory reduction).
    Float16,
    /// Brain float 16 (50% memory reduction, better for training).
    BFloat16,
}

impl PrecisionType {
    /// Get memory multiplier relative to float32.
    pub fn memory_multiplier(&self) -> f64 {
        match self {
            PrecisionType::Float32 => 1.0,
            PrecisionType::Float16 => 0.5,
            PrecisionType::BFloat16 => 0.5,
        }
    }

    /// Get bytes per element.
    pub fn bytes_per_element(&self) -> usize {
        match self {
            PrecisionType::Float32 => 4,
            PrecisionType::Float16 => 2,
            PrecisionType::BFloat16 => 2,
        }
    }
}

/// Precision converter (metadata-only tracking).
///
/// With the Candle removal, this no longer does actual dtype conversion on
/// tensors. It tracks conversion statistics for reporting purposes.
pub struct PrecisionConverter {
    target_precision: PrecisionType,
    conversions: usize,
    memory_saved_mb: f64,
}

impl std::fmt::Debug for PrecisionConverter {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("PrecisionConverter")
            .field("target_precision", &self.target_precision)
            .field("conversions", &self.conversions)
            .field("memory_saved_mb", &self.memory_saved_mb)
            .finish()
    }
}

impl PrecisionConverter {
    /// Create a new precision converter.
    pub fn new(target_precision: PrecisionType) -> Self {
        Self {
            target_precision,
            conversions: 0,
            memory_saved_mb: 0.0,
        }
    }

    /// Record a conversion for tracking purposes.
    ///
    /// `elem_count` is the number of f32 elements being conceptually converted.
    pub fn record_conversion(&mut self, elem_count: usize) {
        self.conversions += 1;
        let original_bytes = elem_count * 4; // f32
        let converted_bytes = elem_count * self.target_precision.bytes_per_element();
        let saved_mb = (original_bytes.saturating_sub(converted_bytes)) as f64 / 1_048_576.0;
        self.memory_saved_mb += saved_mb;
    }

    /// Get conversion statistics.
    pub fn get_stats(&self) -> ConversionStats {
        ConversionStats {
            conversions: self.conversions,
            memory_saved_mb: self.memory_saved_mb,
            target_precision: self.target_precision,
        }
    }

    /// Reset statistics.
    pub fn reset_stats(&mut self) {
        self.conversions = 0;
        self.memory_saved_mb = 0.0;
    }
}

/// Statistics about precision conversions.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ConversionStats {
    /// Number of tensors converted.
    pub conversions: usize,
    /// Total memory saved (MB).
    pub memory_saved_mb: f64,
    /// Target precision.
    pub target_precision: PrecisionType,
}

/// Accuracy metrics for precision conversion validation.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AccuracyMetrics {
    /// Mean Absolute Error.
    pub mae: f64,
    /// Mean Squared Error.
    pub mse: f64,
    /// Root Mean Squared Error.
    pub rmse: f64,
    /// Mean Relative Error.
    pub mean_relative_error: f64,
    /// Maximum Absolute Error.
    pub max_absolute_error: f64,
}

impl AccuracyMetrics {
    /// Check if accuracy degradation is within acceptable threshold.
    pub fn is_acceptable(&self, threshold_percent: f64) -> bool {
        self.mean_relative_error * 100.0 < threshold_percent
    }
}

/// Validate precision accuracy between two f32 arrays.
pub fn validate_precision_accuracy_cpu(
    original: &[f32],
    converted: &[f32],
) -> Result<AccuracyMetrics, MLError> {
    if original.len() != converted.len() {
        return Err(MLError::DimensionMismatch {
            expected: original.len(),
            actual: converted.len(),
        });
    }
    if original.is_empty() {
        return Ok(AccuracyMetrics {
            mae: 0.0,
            mse: 0.0,
            rmse: 0.0,
            mean_relative_error: 0.0,
            max_absolute_error: 0.0,
        });
    }

    let n = original.len() as f64;
    let mut sum_abs_diff = 0.0_f64;
    let mut sum_sq_diff = 0.0_f64;
    let mut sum_rel_diff = 0.0_f64;
    let mut max_abs = 0.0_f64;

    for (o, c) in original.iter().zip(converted.iter()) {
        let diff = (*o as f64 - *c as f64).abs();
        sum_abs_diff += diff;
        sum_sq_diff += diff * diff;
        if o.abs() > 1e-10 {
            sum_rel_diff += diff / o.abs() as f64;
        }
        if diff > max_abs {
            max_abs = diff;
        }
    }

    let mae = sum_abs_diff / n;
    let mse = sum_sq_diff / n;

    Ok(AccuracyMetrics {
        mae,
        mse,
        rmse: mse.sqrt(),
        mean_relative_error: sum_rel_diff / n,
        max_absolute_error: max_abs,
    })
}

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

    #[test]
    fn test_precision_types() {
        assert_eq!(PrecisionType::Float32.bytes_per_element(), 4);
        assert_eq!(PrecisionType::Float16.bytes_per_element(), 2);
        assert_eq!(PrecisionType::BFloat16.bytes_per_element(), 2);
    }

    #[test]
    fn test_memory_multiplier() {
        assert_eq!(PrecisionType::Float32.memory_multiplier(), 1.0);
        assert_eq!(PrecisionType::Float16.memory_multiplier(), 0.5);
        assert_eq!(PrecisionType::BFloat16.memory_multiplier(), 0.5);
    }
}
