//! Integration layer for RTX Science with existing RTX ecosystem
//!
//! This module provides seamless integration with other RTX components,
//! including tensor operations, autograd, distributed computing, and GPU acceleration.

use crate::error::{Result, ScienceError};
use crate::types::MemoryPool; // DistributedContext temporarily disabled
use async_trait::async_trait;
use rtx_autograd::Variable;
use rtx_tensor::{DType, Device, Tensor};
use std::collections::HashMap;

/// RTX Device abstraction for scientific computing
#[derive(Debug, Clone)]
pub struct RTXDevice {
    /// Underlying RTX device
    pub device: Device,
    /// Memory pool for efficient allocation
    pub memory_pool: Option<MemoryPool>,
    /// Device capabilities
    pub capabilities: DeviceCapabilities,
}

/// Scientific tensor with enhanced functionality
#[derive(Debug, Clone)]
pub struct ScientificTensor {
    /// Underlying RTX tensor
    pub tensor: Tensor,
    /// Physical units (if applicable)
    pub units: Option<String>,
    /// Uncertainty/error bars
    pub uncertainty: Option<Tensor>,
    /// Metadata for scientific context
    pub metadata: HashMap<String, String>,
}

/// Automatic differentiation wrapper for scientific computing
pub struct AutoDiff {
    /// Variable tracking for gradients
    pub variables: HashMap<String, Variable>,
    /// Computation graph context
    pub graph_context: Option<String>,
}

/// Gradient computation utilities
pub struct GradientCompute {
    /// Device for computations
    device: Device,
    /// Precision settings
    precision: ComputePrecision,
}

/// Device capabilities for scientific computing
#[derive(Debug, Clone)]
pub struct DeviceCapabilities {
    /// Supports double precision
    pub double_precision: bool,
    /// Supports complex numbers
    pub complex_support: bool,
    /// CUDA compute capability (if applicable)
    pub cuda_compute: Option<String>,
    /// Memory bandwidth (GB/s)
    pub memory_bandwidth: Option<f64>,
    /// Peak FLOPS
    pub peak_flops: Option<f64>,
}

/// Computation precision settings
#[derive(Debug, Clone)]
pub enum ComputePrecision {
    /// Single precision (32-bit)
    Single,
    /// Double precision (64-bit)
    Double,
    /// Mixed precision
    Mixed,
    /// Adaptive precision
    Adaptive { tolerance: f64 },
}

/// Validation metrics for scientific models
pub struct ValidationMetrics {
    /// R-squared coefficient
    pub r_squared: f64,
    /// Mean Absolute Error
    pub mae: f64,
    /// Root Mean Square Error
    pub rmse: f64,
    /// Mean Absolute Percentage Error
    pub mape: f64,
    /// Physics-specific metrics
    pub physics_metrics: HashMap<String, f64>,
}

/// Benchmark suite for scientific computing performance
pub struct BenchmarkSuite {
    /// Available benchmarks
    benchmarks: Vec<Box<dyn ScientificBenchmark + Send + Sync>>,
    /// Results storage
    results: HashMap<String, BenchmarkResult>,
}

/// Scientific benchmark trait
#[async_trait]
pub trait ScientificBenchmark: Send + Sync {
    /// Run the benchmark
    async fn run(&self, device: &RTXDevice) -> Result<BenchmarkResult>;

    /// Benchmark name
    fn name(&self) -> &str;

    /// Benchmark description
    fn description(&self) -> &str;
}

/// Benchmark result
#[derive(Debug, Clone)]
pub struct BenchmarkResult {
    /// Benchmark name
    pub name: String,
    /// Execution time (seconds)
    pub execution_time: f64,
    /// Memory usage (bytes)
    pub memory_usage: u64,
    /// FLOPS achieved
    pub flops: Option<f64>,
    /// Accuracy metrics (if applicable)
    pub accuracy: Option<f64>,
    /// Additional metrics
    pub additional_metrics: HashMap<String, f64>,
}

/// Data loader for scientific datasets
pub struct DataLoader {
    /// Batch size
    pub batch_size: usize,
    /// Device for data loading
    pub device: Device,
    /// Shuffle data
    pub shuffle: bool,
    /// Number of workers for parallel loading
    pub num_workers: usize,
}

impl RTXDevice {
    /// Create RTX device wrapper
    #[must_use]
    pub fn new(device: Device) -> Self {
        let capabilities = DeviceCapabilities::detect(&device);

        Self {
            device,
            memory_pool: None,
            capabilities,
        }
    }

    /// Initialize with memory pool
    #[must_use]
    pub fn with_memory_pool(mut self, pool: MemoryPool) -> Self {
        self.memory_pool = Some(pool);
        self
    }

    /// Check if device supports double precision
    #[must_use]
    pub fn supports_double_precision(&self) -> bool {
        self.capabilities.double_precision
    }

    /// Get optimal tensor dtype for computations
    #[must_use]
    pub fn optimal_dtype(&self, precision: &ComputePrecision) -> DType {
        match precision {
            ComputePrecision::Single => DType::F32,
            ComputePrecision::Double if self.supports_double_precision() => DType::F64,
            ComputePrecision::Double => DType::F32, // Fallback
            ComputePrecision::Mixed => DType::F16,  // Use half precision for mixed
            ComputePrecision::Adaptive { .. } => DType::F32, // Default to single
        }
    }

    /// Allocate scientific tensor with metadata
    pub fn allocate_scientific_tensor(
        &self,
        shape: &[usize],
        _dtype: DType,
        units: Option<String>,
    ) -> Result<ScientificTensor> {
        let tensor = Tensor::zeros(shape, &self.device)?;

        Ok(ScientificTensor {
            tensor,
            units,
            uncertainty: None,
            metadata: HashMap::new(),
        })
    }

    /// Create tensor from scientific data with validation
    pub fn tensor_from_data<T>(
        &self,
        data: &[T],
        shape: &[usize],
        units: Option<String>,
    ) -> Result<ScientificTensor>
    where
        T: Clone + Into<f32>,
    {
        // Validate data integrity
        let expected_len: usize = shape.iter().product();
        if data.len() != expected_len {
            return Err(ScienceError::data_validation(
                "Data length mismatch",
                "data",
                format!("length {expected_len}"),
                format!("length {}", data.len()),
            ));
        }

        // Check for invalid values (NaN, Inf)
        let float_data: Vec<f32> = data.iter().cloned().map(std::convert::Into::into).collect();
        for (i, &value) in float_data.iter().enumerate() {
            if !value.is_finite() {
                return Err(ScienceError::data_validation(
                    format!("Invalid value at index {i}: {value}"),
                    "data_validation",
                    "finite numbers",
                    "NaN or Inf",
                ));
            }
        }

        let tensor = Tensor::from_slice(&float_data, shape, &self.device)?;

        Ok(ScientificTensor {
            tensor,
            units,
            uncertainty: None,
            metadata: HashMap::new(),
        })
    }
}

impl DeviceCapabilities {
    /// Detect device capabilities
    #[must_use]
    pub fn detect(device: &Device) -> Self {
        let double_precision = match device {
            Device::Cpu => true,                        // CPU supports double precision
            Device::Cuda(_) => true,                    // All CUDA GPUs support double precision
            Device::Rocm(_) | Device::Metal(_) => true, // Assume support
        };

        let complex_support = true; // RTX supports complex numbers

        // Detect CUDA compute capability if applicable
        let cuda_compute = match device {
            Device::Cpu => None,
            Device::Cuda(id) => Some(format!("8.{id}")), // Placeholder
            Device::Rocm(_) | Device::Metal(_) => None,
        };

        // Estimate memory bandwidth (placeholder values)
        let memory_bandwidth = match device {
            Device::Cpu => Some(100.0),      // ~100 GB/s for modern CPU
            Device::Cuda(_) => Some(900.0),  // ~900 GB/s for modern GPU
            Device::Rocm(_) => Some(800.0),  // Estimate for ROCm
            Device::Metal(_) => Some(400.0), // Estimate for Metal
        };

        // Estimate peak FLOPS (placeholder values)
        let peak_flops = match device {
            Device::Cpu => Some(1e12),       // ~1 TFLOPS for CPU
            Device::Cuda(_) => Some(50e12),  // ~50 TFLOPS for modern GPU
            Device::Rocm(_) => Some(45e12),  // Estimate for ROCm
            Device::Metal(_) => Some(10e12), // Estimate for Metal
        };

        Self {
            double_precision,
            complex_support,
            cuda_compute,
            memory_bandwidth,
            peak_flops,
        }
    }
}

impl ScientificTensor {
    /// Create scientific tensor from RTX tensor
    #[must_use]
    pub fn from_tensor(tensor: Tensor, units: Option<String>) -> Self {
        Self {
            tensor,
            units,
            uncertainty: None,
            metadata: HashMap::new(),
        }
    }

    /// Add uncertainty/error bars
    pub fn with_uncertainty(mut self, uncertainty: Tensor) -> Result<Self> {
        // Validate uncertainty tensor matches data tensor shape
        if uncertainty.shape() != self.tensor.shape() {
            return Err(ScienceError::data_validation(
                "Uncertainty tensor shape mismatch",
                "uncertainty",
                format!("{:?}", self.tensor.shape()),
                format!("{:?}", uncertainty.shape()),
            ));
        }

        self.uncertainty = Some(uncertainty);
        Ok(self)
    }

    /// Add metadata
    #[must_use]
    pub fn with_metadata(mut self, key: String, value: String) -> Self {
        self.metadata.insert(key, value);
        self
    }

    /// Get tensor with units validation
    pub fn tensor_with_units(&self, expected_units: Option<&str>) -> Result<&Tensor> {
        if let Some(expected) = expected_units {
            if let Some(ref actual) = self.units {
                if actual != expected {
                    return Err(ScienceError::data_validation(
                        "Unit mismatch",
                        "units",
                        expected.to_string(),
                        actual.clone(),
                    ));
                }
            } else {
                return Err(ScienceError::data_validation(
                    "Missing units",
                    "units",
                    expected.to_string(),
                    "none".to_string(),
                ));
            }
        }

        Ok(&self.tensor)
    }

    /// Convert units (simplified implementation)
    pub fn convert_units(&mut self, target_units: &str) -> Result<()> {
        if let Some(ref current_units) = self.units {
            let conversion_factor = get_unit_conversion_factor(current_units, target_units)?;

            self.tensor = self.tensor.mul_scalar(conversion_factor as f32)?;

            if let Some(ref mut uncertainty) = self.uncertainty {
                *uncertainty = uncertainty.mul_scalar(conversion_factor as f32)?;
            }

            self.units = Some(target_units.to_string());
        }

        Ok(())
    }

    /// Statistical summary
    pub fn statistical_summary(&self) -> Result<StatisticalSummary> {
        let mean = f64::from(self.tensor.mean(&[], false)?.to_scalar::<f32>()?);
        let std_dev = f64::from(
            self.tensor
                .std(Some(&[]), false, true)?
                .to_scalar::<f32>()?,
        );
        // For now, use mean as a placeholder for min/max until reduction ops are available
        let min_val = mean - 2.0 * std_dev; // Approximate min
        let max_val = mean + 2.0 * std_dev; // Approximate max

        Ok(StatisticalSummary {
            mean,
            std_dev,
            min: min_val,
            max: max_val,
            units: self.units.clone(),
            sample_size: self.tensor.numel(),
        })
    }
}

/// Statistical summary of scientific data
#[derive(Debug, Clone)]
pub struct StatisticalSummary {
    /// Mean value
    pub mean: f64,
    /// Standard deviation
    pub std_dev: f64,
    /// Minimum value
    pub min: f64,
    /// Maximum value
    pub max: f64,
    /// Data units
    pub units: Option<String>,
    /// Number of data points
    pub sample_size: usize,
}

impl AutoDiff {
    /// Create new `AutoDiff` context
    #[must_use]
    pub fn new() -> Self {
        Self {
            variables: HashMap::new(),
            graph_context: None,
        }
    }

    /// Register variable for gradient tracking
    pub fn register_variable(&mut self, name: String, tensor: Tensor) -> Result<()> {
        let variable = Variable::new(tensor, true);
        self.variables.insert(name, variable);
        Ok(())
    }

    /// Get variable by name
    #[must_use]
    pub fn get_variable(&self, name: &str) -> Option<&Variable> {
        self.variables.get(name)
    }

    /// Compute gradients with respect to all variables
    ///
    /// NOTE: This is a stub implementation. Real gradient computation requires
    /// the Autodiff<B> decorator pattern. Currently returns zero gradients.
    pub async fn compute_gradients(&self, loss: &Variable) -> Result<HashMap<String, Tensor>> {
        // Perform backward pass (stub returns empty HashMap)
        let _grad_map = loss.backward();

        // For now, return zero gradients for all variables
        // Real implementation would extract gradients from the autograd graph
        let mut gradients = HashMap::new();
        for (name, variable) in &self.variables {
            let zero_grad = Tensor::zeros(variable.shape().as_slice(), variable.device())?;
            gradients.insert(name.clone(), zero_grad);
        }

        Ok(gradients)
    }
}

impl GradientCompute {
    /// Create gradient computer
    #[must_use]
    pub fn new(device: Device, precision: ComputePrecision) -> Self {
        Self { device, precision }
    }

    /// Compute numerical gradient using finite differences
    pub async fn numerical_gradient(
        &self,
        f: &dyn Fn(&Tensor) -> Result<Tensor>,
        x: &Tensor,
        h: f64,
    ) -> Result<Tensor> {
        let h_tensor = Tensor::full(x.shape().dims(), h as f32, &self.device)?;

        let f_plus = f(&x.add(&h_tensor)?)?;
        let f_minus = f(&x.subtract(&h_tensor)?)?;

        let gradient = f_plus.subtract(&f_minus)?.div(&h_tensor.mul_scalar(2.0)?)?;

        Ok(gradient)
    }

    /// Validate analytical vs numerical gradients
    pub async fn validate_gradients(
        &self,
        analytical: &Tensor,
        numerical: &Tensor,
        tolerance: f64,
    ) -> Result<bool> {
        let diff = analytical.subtract(numerical)?;
        let max_diff = f64::from(diff.abs()?.max()?.to_scalar::<f32>()?);

        Ok(max_diff < tolerance)
    }
}

impl ValidationMetrics {
    /// Compute metrics from predictions and targets
    pub fn compute(predictions: &Tensor, targets: &Tensor) -> Result<Self> {
        let pred_vec = predictions.to_cpu()?;
        let target_vec = targets.to_cpu()?;

        if pred_vec.len() != target_vec.len() {
            return Err(ScienceError::data_validation(
                "Prediction and target lengths mismatch",
                "validation",
                format!("length {}", target_vec.len()),
                format!("length {}", pred_vec.len()),
            ));
        }

        // Compute R-squared
        let target_mean = target_vec.iter().sum::<f32>() / target_vec.len() as f32;
        let ss_tot: f32 = target_vec.iter().map(|&y| (y - target_mean).powi(2)).sum();
        let ss_res: f32 = pred_vec
            .iter()
            .zip(target_vec.iter())
            .map(|(&pred, &target)| (target - pred).powi(2))
            .sum();
        let r_squared = 1.0 - (ss_res / ss_tot);

        // Compute MAE
        let mae = pred_vec
            .iter()
            .zip(target_vec.iter())
            .map(|(&pred, &target)| (pred - target).abs())
            .sum::<f32>()
            / pred_vec.len() as f32;

        // Compute RMSE
        let mse = pred_vec
            .iter()
            .zip(target_vec.iter())
            .map(|(&pred, &target)| (pred - target).powi(2))
            .sum::<f32>()
            / pred_vec.len() as f32;
        let rmse = mse.sqrt();

        // Compute MAPE
        let mape = pred_vec
            .iter()
            .zip(target_vec.iter())
            .map(|(&pred, &target)| {
                if target.abs() > 1e-8 {
                    ((pred - target) / target).abs()
                } else {
                    0.0
                }
            })
            .sum::<f32>()
            * 100.0
            / pred_vec.len() as f32;

        Ok(Self {
            r_squared: f64::from(r_squared),
            mae: f64::from(mae),
            rmse: f64::from(rmse),
            mape: f64::from(mape),
            physics_metrics: HashMap::new(),
        })
    }

    /// Add physics-specific metric
    #[must_use]
    pub fn add_physics_metric(mut self, name: String, value: f64) -> Self {
        self.physics_metrics.insert(name, value);
        self
    }
}

impl BenchmarkSuite {
    /// Create new benchmark suite
    #[must_use]
    pub fn new() -> Self {
        Self {
            benchmarks: Vec::new(),
            results: HashMap::new(),
        }
    }

    /// Add benchmark
    #[must_use]
    pub fn add_benchmark(mut self, benchmark: Box<dyn ScientificBenchmark + Send + Sync>) -> Self {
        self.benchmarks.push(benchmark);
        self
    }

    /// Run all benchmarks
    pub async fn run_all(&mut self, device: &RTXDevice) -> Result<()> {
        for benchmark in &self.benchmarks {
            let result = benchmark.run(device).await?;
            self.results.insert(result.name.clone(), result);
        }
        Ok(())
    }

    /// Get results summary
    #[must_use]
    pub fn results_summary(&self) -> HashMap<String, String> {
        let mut summary = HashMap::new();

        for (name, result) in &self.results {
            summary.insert(
                format!("{name}_time"),
                format!("{:.4} s", result.execution_time),
            );
            summary.insert(
                format!("{name}_memory"),
                format!("{:.2} MB", result.memory_usage as f64 / 1024.0 / 1024.0),
            );
            if let Some(flops) = result.flops {
                summary.insert(format!("{name}_flops"), format!("{flops:.2e} FLOPS"));
            }
        }

        summary
    }
}

impl DataLoader {
    /// Create new data loader
    #[must_use]
    pub fn new(batch_size: usize, device: Device) -> Self {
        Self {
            batch_size,
            device,
            shuffle: true,
            num_workers: 1,
        }
    }

    /// Load batch of scientific data
    pub async fn load_batch(&self, data: &[ScientificTensor]) -> Result<Vec<ScientificTensor>> {
        if data.len() < self.batch_size {
            return Ok(data.to_vec());
        }

        // Simple batch loading (in practice would be more sophisticated)
        let batch = data.iter().take(self.batch_size).cloned().collect();
        Ok(batch)
    }

    /// Set number of parallel workers
    #[must_use]
    pub fn with_workers(mut self, num_workers: usize) -> Self {
        self.num_workers = num_workers;
        self
    }

    /// Enable/disable shuffling
    #[must_use]
    pub fn with_shuffle(mut self, shuffle: bool) -> Self {
        self.shuffle = shuffle;
        self
    }
}

/// Simple unit conversion (placeholder implementation)
fn get_unit_conversion_factor(from_units: &str, to_units: &str) -> Result<f32> {
    if from_units == to_units {
        return Ok(1.0);
    }

    // Placeholder conversions
    match (from_units, to_units) {
        ("m", "cm") => Ok(100.0),
        ("cm", "m") => Ok(0.01),
        ("kg", "g") => Ok(1000.0),
        ("g", "kg") => Ok(0.001),
        ("K", "C") => Ok(1.0), // Temperature difference conversion
        ("eV", "J") => Ok(1.602_176_6e-19),
        ("J", "eV") => Ok(6.241_509e18),
        _ => Err(ScienceError::data_validation(
            format!("Unknown unit conversion: {from_units} to {to_units}"),
            "units",
            to_units.to_string(),
            from_units.to_string(),
        )),
    }
}

/// Matrix multiplication benchmark
pub struct MatMulBenchmark {
    pub size: usize,
}

#[async_trait]
impl ScientificBenchmark for MatMulBenchmark {
    async fn run(&self, device: &RTXDevice) -> Result<BenchmarkResult> {
        let start = std::time::Instant::now();

        let a = Tensor::randn(&[self.size, self.size], &device.device)?;
        let b = Tensor::randn(&[self.size, self.size], &device.device)?;
        let _c = a.matmul(&b)?;

        let elapsed = start.elapsed().as_secs_f64();
        let flops = 2.0 * (self.size as f64).powi(3); // 2 * n^3 operations

        Ok(BenchmarkResult {
            name: format!("MatMul{}x{}", self.size, self.size),
            execution_time: elapsed,
            memory_usage: (3 * self.size * self.size * 4) as u64, // 3 matrices * f32
            flops: Some(flops / elapsed),
            accuracy: None,
            additional_metrics: HashMap::new(),
        })
    }

    fn name(&self) -> &'static str {
        "Matrix Multiplication"
    }

    fn description(&self) -> &'static str {
        "Benchmark matrix multiplication performance"
    }
}

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

    #[test]
    fn test_rtx_device_creation() {
        let device = Device::cpu();
        let rtx_device = RTXDevice::new(device);
        assert!(rtx_device.supports_double_precision());
    }

    #[tokio::test]
    async fn test_scientific_tensor() -> Result<()> {
        let device = Device::cpu();
        let rtx_device = RTXDevice::new(device);

        let data = vec![1.0, 2.0, 3.0, 4.0];
        let tensor = rtx_device.tensor_from_data(&data, &[2, 2], Some("m".to_string()))?;

        assert_eq!(tensor.units, Some("m".to_string()));
        Ok(())
    }

    #[test]
    fn test_validation_metrics() -> Result<()> {
        let device = Device::cpu();
        let predictions = Tensor::from_slice(&[1.0, 2.0, 3.0], &[3], &device)?;
        let targets = Tensor::from_slice(&[1.1, 1.9, 3.1], &[3], &device)?;

        let metrics = ValidationMetrics::compute(&predictions, &targets)?;
        assert!(metrics.r_squared > 0.0);
        assert!(metrics.mae < 0.5);

        Ok(())
    }

    #[tokio::test]
    async fn test_benchmark_suite() -> Result<()> {
        let device = RTXDevice::new(Device::cpu());
        let mut suite =
            BenchmarkSuite::new().add_benchmark(Box::new(MatMulBenchmark { size: 100 }));

        suite.run_all(&device).await?;
        let summary = suite.results_summary();
        assert!(!summary.is_empty());

        Ok(())
    }
}