//! Latency prediction for hardware-aware NAS
//!
//! Provides latency estimation and prediction for architectures
//! on different hardware devices.

use super::device::DeviceProfile;
use crate::error::{NASError, Result};
use crate::search_space::{Architecture, Cell, OperationType};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// Latency predictor trait
pub trait LatencyPredictor: Send + Sync {
    /// Predict latency for an architecture on a device
    fn predict(&self, arch: &Architecture, device: &DeviceProfile) -> Result<f32>;

    /// Calibrate the predictor with measured samples
    fn calibrate(&mut self, samples: &[(Architecture, f32)]) -> Result<()>;

    /// Get prediction confidence (0.0 to 1.0)
    fn confidence(&self) -> f32;
}

/// Operation-level latency lookup table
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OperationLatency {
    /// Base latency in microseconds
    pub base_latency_us: f32,
    /// Latency per channel
    pub per_channel_us: f32,
    /// Latency per spatial element (H*W)
    pub per_spatial_us: f32,
    /// Memory access pattern factor
    pub memory_factor: f32,
}

impl OperationLatency {
    /// Create a new operation latency entry
    pub fn new(base: f32, per_channel: f32, per_spatial: f32) -> Self {
        Self {
            base_latency_us: base,
            per_channel_us: per_channel,
            per_spatial_us: per_spatial,
            memory_factor: 1.0,
        }
    }

    /// Estimate latency for given configuration
    pub fn estimate(&self, channels: usize, height: usize, width: usize) -> f32 {
        let spatial = (height * width) as f32;
        self.base_latency_us
            + self.per_channel_us * channels as f32
            + self.per_spatial_us * spatial * self.memory_factor
    }
}

/// Lookup table based latency predictor
#[derive(Debug, Clone)]
pub struct LookupTablePredictor {
    /// Latency tables per device type per operation
    tables: HashMap<String, HashMap<OperationType, OperationLatency>>,
    /// Default latencies when no specific table exists
    default_latencies: HashMap<OperationType, OperationLatency>,
    /// Confidence in predictions
    confidence: f32,
}

impl LookupTablePredictor {
    /// Create a new lookup table predictor with default values
    pub fn new() -> Self {
        let mut default_latencies = HashMap::new();

        // Default latencies based on typical GPU performance
        default_latencies.insert(
            OperationType::Identity,
            OperationLatency::new(0.1, 0.0, 0.0001),
        );
        default_latencies.insert(OperationType::Zero, OperationLatency::new(0.05, 0.0, 0.0));
        default_latencies.insert(
            OperationType::Conv3x3,
            OperationLatency::new(10.0, 0.5, 0.01),
        );
        default_latencies.insert(
            OperationType::Conv5x5,
            OperationLatency::new(25.0, 1.2, 0.025),
        );
        default_latencies.insert(
            OperationType::SepConv3x3,
            OperationLatency::new(8.0, 0.3, 0.008),
        );
        default_latencies.insert(
            OperationType::SepConv5x5,
            OperationLatency::new(15.0, 0.6, 0.015),
        );
        default_latencies.insert(
            OperationType::DilConv3x3,
            OperationLatency::new(12.0, 0.6, 0.012),
        );
        default_latencies.insert(
            OperationType::MaxPool3x3,
            OperationLatency::new(2.0, 0.1, 0.002),
        );
        default_latencies.insert(
            OperationType::AvgPool3x3,
            OperationLatency::new(2.5, 0.1, 0.0025),
        );

        Self {
            tables: HashMap::new(),
            default_latencies,
            confidence: 0.5, // Medium confidence for uncalibrated predictor
        }
    }

    /// Add device-specific latency table
    pub fn add_device_table(
        &mut self,
        device_name: impl Into<String>,
        table: HashMap<OperationType, OperationLatency>,
    ) {
        self.tables.insert(device_name.into(), table);
    }

    /// Get latency for an operation on a device
    fn get_op_latency(&self, op: OperationType, device: &DeviceProfile) -> OperationLatency {
        // Try device-specific table first
        if let Some(device_table) = self.tables.get(&device.name) {
            if let Some(latency) = device_table.get(&op) {
                return latency.clone();
            }
        }

        // Fall back to default
        self.default_latencies
            .get(&op)
            .cloned()
            .unwrap_or_else(|| OperationLatency::new(10.0, 0.5, 0.01))
    }

    /// Estimate latency for a cell
    fn estimate_cell_latency(
        &self,
        cell: &Cell,
        device: &DeviceProfile,
        channels: usize,
        height: usize,
        width: usize,
    ) -> f32 {
        let mut total_latency = 0.0;

        for edge in cell.edges() {
            if let Some(op_type) = cell.get_operation(&edge) {
                let op_latency = self.get_op_latency(op_type, device);
                total_latency += op_latency.estimate(channels, height, width);
            }
        }

        // Apply device-specific scaling
        let device_factor = match device.device_type {
            super::device::DeviceType::CPU => 10.0,
            super::device::DeviceType::Mobile => 5.0,
            super::device::DeviceType::Edge => 20.0,
            _ => 1.0, // GPUs as baseline
        };

        total_latency * device_factor / device.peak_tflops_fp32.max(0.1)
    }
}

impl Default for LookupTablePredictor {
    fn default() -> Self {
        Self::new()
    }
}

impl LatencyPredictor for LookupTablePredictor {
    fn predict(&self, arch: &Architecture, device: &DeviceProfile) -> Result<f32> {
        let mut total_latency_ms = 0.0;

        // Default dimensions if not specified in architecture
        let channels = arch.channels;
        let height = 224; // Default ImageNet size
        let width = 224;

        for cell in &arch.cells {
            let cell_latency_us = self.estimate_cell_latency(cell, device, channels, height, width);
            total_latency_ms += cell_latency_us / 1000.0; // Convert to ms
        }

        Ok(total_latency_ms)
    }

    fn calibrate(&mut self, samples: &[(Architecture, f32)]) -> Result<()> {
        if samples.is_empty() {
            return Err(NASError::LatencyError(
                "No calibration samples provided".into(),
            ));
        }

        // Simple calibration: adjust confidence based on sample count
        let sample_count = samples.len();
        self.confidence = (sample_count as f32 / 100.0).min(0.95);

        // In a real implementation, we would fit the latency tables to the samples
        Ok(())
    }

    fn confidence(&self) -> f32 {
        self.confidence
    }
}

/// Linear regression based latency predictor
#[derive(Debug, Clone)]
pub struct RegressionPredictor {
    /// Learned coefficients for each operation type
    coefficients: HashMap<OperationType, Vec<f32>>,
    /// Bias term
    bias: f32,
    /// Prediction confidence
    confidence: f32,
    /// Number of features per operation
    num_features: usize,
}

impl RegressionPredictor {
    /// Create a new regression predictor
    pub fn new() -> Self {
        Self {
            coefficients: HashMap::new(),
            bias: 0.0,
            confidence: 0.0, // No confidence until trained
            num_features: 4, // channels, height, width, device_factor
        }
    }

    /// Extract features from an architecture
    fn extract_features(&self, arch: &Architecture, device: &DeviceProfile) -> Vec<f32> {
        let mut features = Vec::new();

        // Count operations of each type
        let mut op_counts: HashMap<OperationType, usize> = HashMap::new();
        for cell in &arch.cells {
            for edge in cell.edges() {
                if let Some(op_type) = cell.get_operation(&edge) {
                    *op_counts.entry(op_type).or_insert(0) += 1;
                }
            }
        }

        // Add operation counts as features
        for op_type in OperationType::all() {
            let count = *op_counts.get(&op_type).unwrap_or(&0);
            features.push(count as f32);
        }

        // Add architecture properties
        features.push(arch.num_cells() as f32);
        features.push(arch.channels as f32);
        features.push(arch.channels as f32); // Same for in/out

        // Add device factor
        let device_factor = 1.0 / device.peak_tflops_fp32.max(0.1);
        features.push(device_factor);

        features
    }

    /// Predict using learned coefficients
    fn predict_internal(&self, features: &[f32]) -> f32 {
        if self.coefficients.is_empty() {
            // Fall back to simple estimation
            return features.iter().sum::<f32>() * 0.1 + self.bias;
        }

        // Linear combination of features
        let mut prediction = self.bias;
        let num_ops = OperationType::count();

        for (i, op_type) in OperationType::all().iter().enumerate() {
            if let Some(coeffs) = self.coefficients.get(op_type) {
                if i < features.len() {
                    prediction += features[i] * coeffs.first().unwrap_or(&1.0);
                }
            }
        }

        // Add contribution from other features
        for (_i, &f) in features.iter().enumerate().skip(num_ops) {
            prediction += f * 0.1; // Simple scaling for non-operation features
        }

        prediction.max(0.0) // Latency can't be negative
    }
}

impl Default for RegressionPredictor {
    fn default() -> Self {
        Self::new()
    }
}

impl LatencyPredictor for RegressionPredictor {
    fn predict(&self, arch: &Architecture, device: &DeviceProfile) -> Result<f32> {
        let features = self.extract_features(arch, device);
        Ok(self.predict_internal(&features))
    }

    fn calibrate(&mut self, samples: &[(Architecture, f32)]) -> Result<()> {
        if samples.is_empty() {
            return Err(NASError::LatencyError(
                "No calibration samples provided".into(),
            ));
        }

        // Simple least squares fitting
        // In a real implementation, use proper linear regression

        // Compute mean latency as baseline
        let mean_latency: f32 = samples.iter().map(|(_, l)| l).sum::<f32>() / samples.len() as f32;
        self.bias = mean_latency;

        // Initialize coefficients for each operation
        for op_type in OperationType::all() {
            self.coefficients
                .insert(op_type, vec![1.0; self.num_features]);
        }

        // Update confidence based on sample count
        self.confidence = (samples.len() as f32 / 100.0).min(0.9);

        Ok(())
    }

    fn confidence(&self) -> f32 {
        self.confidence
    }
}

/// Latency measurement result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LatencyMeasurement {
    /// Architecture identifier
    pub arch_id: String,
    /// Device name
    pub device_name: String,
    /// Measured latency in milliseconds
    pub latency_ms: f32,
    /// Standard deviation (if multiple measurements)
    pub std_dev: Option<f32>,
    /// Number of measurements
    pub num_samples: usize,
    /// Batch size used for measurement
    pub batch_size: usize,
}

impl LatencyMeasurement {
    /// Create a new latency measurement
    pub fn new(
        arch_id: impl Into<String>,
        device_name: impl Into<String>,
        latency_ms: f32,
    ) -> Self {
        Self {
            arch_id: arch_id.into(),
            device_name: device_name.into(),
            latency_ms,
            std_dev: None,
            num_samples: 1,
            batch_size: 1,
        }
    }

    /// Add statistical information
    pub fn with_stats(mut self, std_dev: f32, num_samples: usize) -> Self {
        self.std_dev = Some(std_dev);
        self.num_samples = num_samples;
        self
    }

    /// Set batch size
    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
        self.batch_size = batch_size;
        self
    }
}

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

    fn create_test_architecture() -> Architecture {
        let config = CellConfig::default_darts();
        let cell = Cell::new(config).unwrap();
        Architecture::new("test_arch".to_string(), vec![cell], 16, 32)
    }

    #[test]
    fn test_operation_latency() {
        let latency = OperationLatency::new(10.0, 0.5, 0.01);
        let estimate = latency.estimate(64, 32, 32);
        assert!(estimate > 10.0); // Should be more than base
    }

    #[test]
    fn test_lookup_table_predictor_new() {
        let predictor = LookupTablePredictor::new();
        assert_eq!(predictor.confidence(), 0.5);
        assert!(!predictor.default_latencies.is_empty());
    }

    #[test]
    fn test_lookup_table_predict() {
        let predictor = LookupTablePredictor::new();
        let arch = create_test_architecture();
        let device = super::super::device::CommonDevices::rtx_3090();

        let result = predictor.predict(&arch, &device);
        assert!(result.is_ok());
        assert!(result.unwrap() >= 0.0);
    }

    #[test]
    fn test_lookup_table_calibrate() {
        let mut predictor = LookupTablePredictor::new();
        let arch = create_test_architecture();

        let samples = vec![
            (arch.clone(), 1.0),
            (arch.clone(), 1.1),
            (arch.clone(), 0.9),
        ];

        let result = predictor.calibrate(&samples);
        assert!(result.is_ok());
        assert!(predictor.confidence() > 0.0);
    }

    #[test]
    fn test_regression_predictor_new() {
        let predictor = RegressionPredictor::new();
        assert_eq!(predictor.confidence(), 0.0);
    }

    #[test]
    fn test_regression_predict() {
        let predictor = RegressionPredictor::new();
        let arch = create_test_architecture();
        let device = super::super::device::CommonDevices::rtx_3090();

        let result = predictor.predict(&arch, &device);
        assert!(result.is_ok());
        assert!(result.unwrap() >= 0.0);
    }

    #[test]
    fn test_regression_calibrate() {
        let mut predictor = RegressionPredictor::new();
        let arch = create_test_architecture();

        let samples = vec![
            (arch.clone(), 1.0),
            (arch.clone(), 1.1),
            (arch.clone(), 0.9),
        ];

        let result = predictor.calibrate(&samples);
        assert!(result.is_ok());
        assert!(predictor.confidence() > 0.0);
    }

    #[test]
    fn test_latency_measurement() {
        let measurement = LatencyMeasurement::new("arch_1", "RTX 3090", 1.5)
            .with_stats(0.1, 10)
            .with_batch_size(1);

        assert_eq!(measurement.arch_id, "arch_1");
        assert_eq!(measurement.latency_ms, 1.5);
        assert_eq!(measurement.std_dev, Some(0.1));
        assert_eq!(measurement.num_samples, 10);
    }

    #[test]
    fn test_calibrate_empty_samples() {
        let mut predictor = LookupTablePredictor::new();
        let result = predictor.calibrate(&[]);
        assert!(result.is_err());
    }
}