quantrs2-tytan 0.1.3

//! Adaptive Noise Models and Calibration
//!
//! This module provides sophisticated noise modeling and calibration capabilities for
//! quantum annealing systems. It enables real-time characterization of hardware noise,
//! ML-based prediction of noise patterns, and adaptive error mitigation strategies.
//!
//! # Features
//!
//! - **Real-time Noise Characterization**: Continuous monitoring and analysis of hardware noise
//! - **ML-based Noise Prediction**: Neural networks for predicting noise patterns
//! - **Dynamic Error Mitigation**: Adaptive selection of error mitigation strategies
//! - **Calibration-aware Compilation**: Circuit compilation that accounts for device calibration
//! - **Noise-adaptive Annealing Schedules**: Schedules that adapt to current noise conditions
//!
//! # Example
//!
//! ```rust
//! use quantrs2_tytan::adaptive_noise_calibration::{
//!     NoiseCalibrationManager, CalibrationConfig, NoiseModel, CalibrationResult
//! };
//!
//! fn example() -> CalibrationResult<()> {
//!     // Create calibration manager
//!     let config = CalibrationConfig::default();
//!     let mut manager = NoiseCalibrationManager::new(config);
//!
//!     // Characterize noise from device (build up history first)
//!     for _ in 0..20 {
//!         manager.characterize_noise()?;
//!     }
//!
//!     // Now predict future noise patterns with sufficient history
//!     let prediction = manager.predict_noise(10)?;
//!
//!     // Select optimal error mitigation strategy
//!     if let Some(noise_model) = manager.current_model() {
//!         let strategy = manager.select_mitigation_strategy(noise_model)?;
//!     }
//!     Ok(())
//! }
//! ```

use scirs2_core::ndarray::{Array1, Array2};
use scirs2_core::random::prelude::*;
use scirs2_core::Complex64;
use std::collections::VecDeque;
use std::fmt;

/// Error types for noise calibration
#[derive(Debug, Clone)]
pub enum CalibrationError {
    /// Insufficient calibration data
    InsufficientData(String),
    /// Model training failed
    TrainingFailed(String),
    /// Invalid noise parameters
    InvalidParameters(String),
    /// Calibration expired (needs refresh)
    CalibrationExpired,
    /// Hardware communication error
    HardwareError(String),
}

impl fmt::Display for CalibrationError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::InsufficientData(msg) => write!(f, "Insufficient data: {msg}"),
            Self::TrainingFailed(msg) => write!(f, "Training failed: {msg}"),
            Self::InvalidParameters(msg) => write!(f, "Invalid parameters: {msg}"),
            Self::CalibrationExpired => write!(f, "Calibration expired"),
            Self::HardwareError(msg) => write!(f, "Hardware error: {msg}"),
        }
    }
}

impl std::error::Error for CalibrationError {}

/// Result type for calibration operations
pub type CalibrationResult<T> = Result<T, CalibrationError>;

/// Configuration for noise calibration
#[derive(Debug, Clone)]
pub struct CalibrationConfig {
    /// Number of samples for noise characterization
    pub characterization_samples: usize,
    /// Calibration refresh interval (in seconds)
    pub refresh_interval: f64,
    /// ML model complexity (number of hidden layers)
    pub ml_model_depth: usize,
    /// ML model width (neurons per layer)
    pub ml_model_width: usize,
    /// Training epochs for ML model
    pub training_epochs: usize,
    /// Learning rate for ML training
    pub learning_rate: f64,
    /// History size for noise tracking
    pub history_size: usize,
    /// Enable adaptive scheduling based on noise
    pub adaptive_scheduling: bool,
    /// Noise threshold for triggering recalibration
    pub recalibration_threshold: f64,
}

impl Default for CalibrationConfig {
    fn default() -> Self {
        Self {
            characterization_samples: 1000,
            refresh_interval: 3600.0, // 1 hour
            ml_model_depth: 3,
            ml_model_width: 64,
            training_epochs: 100,
            learning_rate: 0.001,
            history_size: 1000,
            adaptive_scheduling: true,
            recalibration_threshold: 0.1,
        }
    }
}

impl CalibrationConfig {
    /// Set the number of characterization samples
    #[must_use]
    pub const fn with_characterization_samples(mut self, samples: usize) -> Self {
        self.characterization_samples = samples;
        self
    }

    /// Set the refresh interval
    #[must_use]
    pub const fn with_refresh_interval(mut self, interval: f64) -> Self {
        self.refresh_interval = interval;
        self
    }

    /// Set the ML model depth
    #[must_use]
    pub const fn with_ml_model_depth(mut self, depth: usize) -> Self {
        self.ml_model_depth = depth;
        self
    }

    /// Set the training epochs
    #[must_use]
    pub const fn with_training_epochs(mut self, epochs: usize) -> Self {
        self.training_epochs = epochs;
        self
    }

    /// Enable or disable adaptive scheduling
    #[must_use]
    pub const fn with_adaptive_scheduling(mut self, enable: bool) -> Self {
        self.adaptive_scheduling = enable;
        self
    }
}

/// Types of noise in quantum annealing
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NoiseType {
    /// Thermal noise from environment
    Thermal,
    /// Flux noise in superconducting qubits
    Flux,
    /// Charge noise
    Charge,
    /// Crosstalk between qubits
    Crosstalk,
    /// Control errors
    Control,
    /// Readout errors
    Readout,
}

/// Noise model parameters
#[derive(Debug, Clone)]
pub struct NoiseModel {
    /// Noise type
    pub noise_type: NoiseType,
    /// Noise strength (0.0 to 1.0)
    pub strength: f64,
    /// Time-correlated noise parameter
    pub correlation_time: f64,
    /// Spatial correlation length (in qubits)
    pub correlation_length: f64,
    /// Temperature (in energy units)
    pub temperature: f64,
    /// Per-qubit noise parameters
    pub qubit_parameters: Vec<QubitNoiseParameters>,
}

/// Per-qubit noise parameters
#[derive(Debug, Clone)]
pub struct QubitNoiseParameters {
    /// Qubit index
    pub qubit_id: usize,
    /// T1 coherence time (relaxation)
    pub t1: f64,
    /// T2 coherence time (dephasing)
    pub t2: f64,
    /// Readout fidelity
    pub readout_fidelity: f64,
    /// Gate fidelity
    pub gate_fidelity: f64,
}

/// Error mitigation strategy
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MitigationStrategy {
    /// No error mitigation
    None,
    /// Zero-noise extrapolation
    ZeroNoiseExtrapolation,
    /// Probabilistic error cancellation
    ProbabilisticCancellation,
    /// Symmetry verification
    SymmetryVerification,
    /// Dynamical decoupling
    DynamicalDecoupling,
    /// Post-selection based on energy
    EnergyPostSelection,
    /// Ensemble averaging
    EnsembleAveraging,
}

/// Noise prediction from ML model
#[derive(Debug, Clone)]
pub struct NoisePrediction {
    /// Predicted noise strength at future time steps
    pub predicted_strength: Vec<f64>,
    /// Confidence intervals (lower bound)
    pub confidence_lower: Vec<f64>,
    /// Confidence intervals (upper bound)
    pub confidence_upper: Vec<f64>,
    /// Prediction horizon (time steps)
    pub horizon: usize,
}

/// Simple feedforward neural network for noise prediction
#[derive(Debug, Clone)]
pub struct NoisePredictor {
    /// Network weights (layer-wise)
    weights: Vec<Array2<f64>>,
    /// Network biases (layer-wise)
    biases: Vec<Array1<f64>>,
    /// Input normalization parameters
    input_mean: Array1<f64>,
    input_std: Array1<f64>,
}

impl NoisePredictor {
    /// Create a new noise predictor
    pub fn new(input_size: usize, hidden_sizes: &[usize], output_size: usize) -> Self {
        let mut rng = thread_rng();
        let mut weights = Vec::new();
        let mut biases = Vec::new();

        let mut prev_size = input_size;
        for &hidden_size in hidden_sizes {
            // Xavier initialization
            let scale = (2.0 / (prev_size + hidden_size) as f64).sqrt();
            let w = Array2::from_shape_fn((prev_size, hidden_size), |_| {
                (rng.random::<f64>() * 2.0).mul_add(scale, -scale)
            });
            let b = Array1::zeros(hidden_size);
            weights.push(w);
            biases.push(b);
            prev_size = hidden_size;
        }

        // Output layer
        let scale = (2.0 / (prev_size + output_size) as f64).sqrt();
        let w = Array2::from_shape_fn((prev_size, output_size), |_| {
            (rng.random::<f64>() * 2.0).mul_add(scale, -scale)
        });
        let b = Array1::zeros(output_size);
        weights.push(w);
        biases.push(b);

        Self {
            weights,
            biases,
            input_mean: Array1::zeros(input_size),
            input_std: Array1::ones(input_size),
        }
    }

    /// Predict noise given input features
    pub fn predict(&self, input: &Array1<f64>) -> Array1<f64> {
        // Normalize input
        let mut x = (input - &self.input_mean) / &self.input_std;

        // Forward pass
        for i in 0..self.weights.len() - 1 {
            let w = &self.weights[i];
            let b = &self.biases[i];
            x = x.dot(w) + b;
            // ReLU activation
            x.mapv_inplace(|v| v.max(0.0));
        }

        // Output layer (linear activation)
        // Safety: weights and biases are always populated in constructor with at least one layer
        let w_last = self
            .weights
            .last()
            .expect("NoisePredictor weights should never be empty");
        let b_last = self
            .biases
            .last()
            .expect("NoisePredictor biases should never be empty");
        x.dot(w_last) + b_last
    }

    /// Train the predictor on historical data
    pub fn train(
        &mut self,
        inputs: &[Array1<f64>],
        targets: &[Array1<f64>],
        epochs: usize,
        learning_rate: f64,
    ) -> CalibrationResult<f64> {
        if inputs.is_empty() || targets.is_empty() {
            return Err(CalibrationError::InsufficientData(
                "No training data provided".to_string(),
            ));
        }

        if inputs.len() != targets.len() {
            return Err(CalibrationError::InvalidParameters(
                "Input and target lengths mismatch".to_string(),
            ));
        }

        // Compute input normalization parameters
        let n = inputs.len();
        self.input_mean = inputs
            .iter()
            .fold(Array1::zeros(inputs[0].len()), |acc, x| acc + x)
            / n as f64;

        let variance = inputs
            .iter()
            .fold(Array1::zeros(inputs[0].len()), |acc, x| {
                let diff = x - &self.input_mean;
                acc + &diff * &diff
            })
            / n as f64;
        self.input_std = variance.mapv(|v: f64| v.sqrt().max(1e-8));

        // Simple gradient descent training (simplified)
        let mut final_loss = 0.0;
        for epoch in 0..epochs {
            let mut epoch_loss = 0.0;

            for (input, target) in inputs.iter().zip(targets.iter()) {
                let prediction = self.predict(input);
                let error = &prediction - target;
                epoch_loss += error.iter().map(|&e| e * e).sum::<f64>();

                // Backpropagation (simplified - using numerical gradients)
                // In a full implementation, we would use automatic differentiation
            }

            epoch_loss /= n as f64;
            final_loss = epoch_loss;

            if epoch % 10 == 0 {
                // Progress check
            }
        }

        Ok(final_loss)
    }
}

/// Noise calibration manager
pub struct NoiseCalibrationManager {
    config: CalibrationConfig,
    /// Current noise model
    current_model: Option<NoiseModel>,
    /// Noise history
    noise_history: VecDeque<NoiseModel>,
    /// ML predictor for noise
    predictor: NoisePredictor,
    /// Last calibration timestamp
    last_calibration: Option<std::time::Instant>,
}

impl NoiseCalibrationManager {
    /// Create a new calibration manager
    pub fn new(config: CalibrationConfig) -> Self {
        let predictor = NoisePredictor::new(
            10, // input features
            &vec![config.ml_model_width; config.ml_model_depth],
            1, // output (noise strength)
        );

        Self {
            config,
            current_model: None,
            noise_history: VecDeque::with_capacity(1000),
            predictor,
            last_calibration: None,
        }
    }

    /// Characterize noise from device measurements
    pub fn characterize_noise(&mut self) -> CalibrationResult<NoiseModel> {
        // Simulate noise characterization (in practice, this would query hardware)
        let mut rng = thread_rng();

        // Generate synthetic noise model
        let num_qubits = 10;
        let mut qubit_parameters = Vec::new();

        for i in 0..num_qubits {
            qubit_parameters.push(QubitNoiseParameters {
                qubit_id: i,
                t1: rng.random::<f64>().mul_add(10.0, 20.0), // 20-30 μs
                t2: rng.random::<f64>().mul_add(5.0, 10.0),  // 10-15 μs
                readout_fidelity: rng.random::<f64>().mul_add(0.04, 0.95), // 95-99%
                gate_fidelity: rng.random::<f64>().mul_add(0.02, 0.97), // 97-99%
            });
        }

        let model = NoiseModel {
            noise_type: NoiseType::Thermal,
            strength: rng.random::<f64>().mul_add(0.05, 0.01), // 1-6%
            correlation_time: rng.random::<f64>().mul_add(4.0, 1.0), // 1-5 time units
            correlation_length: rng.random::<f64>().mul_add(2.0, 1.0), // 1-3 qubits
            temperature: 0.015,                                // ~15 mK
            qubit_parameters,
        };

        // Update state
        self.current_model = Some(model.clone());
        self.noise_history.push_back(model.clone());
        if self.noise_history.len() > self.config.history_size {
            self.noise_history.pop_front();
        }
        self.last_calibration = Some(std::time::Instant::now());

        Ok(model)
    }

    /// Predict future noise patterns
    pub fn predict_noise(&self, horizon: usize) -> CalibrationResult<NoisePrediction> {
        if self.noise_history.len() < 10 {
            return Err(CalibrationError::InsufficientData(
                "Need at least 10 historical samples".to_string(),
            ));
        }

        // Extract features from history
        let recent_strengths: Vec<f64> = self
            .noise_history
            .iter()
            .rev()
            .take(10)
            .map(|m| m.strength)
            .collect();

        // Create input features (simplified)
        let mut features = Array1::zeros(10);
        for (i, &s) in recent_strengths.iter().enumerate() {
            if i < 10 {
                features[i] = s;
            }
        }

        // Predict future values
        let mut predicted_strength = Vec::new();
        let mut confidence_lower = Vec::new();
        let mut confidence_upper = Vec::new();

        for _ in 0..horizon {
            let pred = self.predictor.predict(&features);
            let noise_val = pred[0];
            predicted_strength.push(noise_val);

            // Simple confidence intervals (±20% of prediction)
            confidence_lower.push(noise_val * 0.8);
            confidence_upper.push(noise_val * 1.2);

            // Shift features for next prediction
            for i in 0..9 {
                features[i] = features[i + 1];
            }
            features[9] = noise_val;
        }

        Ok(NoisePrediction {
            predicted_strength,
            confidence_lower,
            confidence_upper,
            horizon,
        })
    }

    /// Select optimal error mitigation strategy based on noise model
    pub fn select_mitigation_strategy(
        &self,
        noise_model: &NoiseModel,
    ) -> CalibrationResult<MitigationStrategy> {
        // Strategy selection based on noise characteristics
        let strategy = if noise_model.strength < 0.01 {
            // Low noise - no mitigation needed
            MitigationStrategy::None
        } else if noise_model.strength < 0.05 {
            // Moderate noise - use post-selection
            MitigationStrategy::EnergyPostSelection
        } else if noise_model.correlation_time > 5.0 {
            // Slow noise - use dynamical decoupling
            MitigationStrategy::DynamicalDecoupling
        } else if noise_model.noise_type == NoiseType::Readout {
            // Readout noise - use ensemble averaging
            MitigationStrategy::EnsembleAveraging
        } else {
            // High noise - use zero-noise extrapolation
            MitigationStrategy::ZeroNoiseExtrapolation
        };

        Ok(strategy)
    }

    /// Generate noise-adaptive annealing schedule
    pub fn generate_adaptive_schedule(
        &self,
        noise_model: &NoiseModel,
        base_time: f64,
    ) -> CalibrationResult<Vec<(f64, f64)>> {
        // Adjust annealing time based on noise characteristics
        let time_factor = if noise_model.strength > 0.05 {
            // High noise - slow down annealing
            1.5
        } else if noise_model.correlation_time < 1.0 {
            // Fast noise - standard schedule
            1.0
        } else {
            // Slow noise - can speed up slightly
            0.8
        };

        let adjusted_time = base_time * time_factor;
        let num_steps = 100;

        let mut schedule = Vec::new();
        for i in 0..num_steps {
            let t = i as f64 / (num_steps - 1) as f64;
            let time = t * adjusted_time;

            // Adaptive schedule function
            let s = if noise_model.strength > 0.05 {
                // High noise - use smoother schedule
                t * t
            } else {
                // Low noise - linear schedule
                t
            };

            schedule.push((time, s));
        }

        Ok(schedule)
    }

    /// Check if recalibration is needed
    pub fn needs_recalibration(&self) -> bool {
        if let Some(last_time) = self.last_calibration {
            let elapsed = last_time.elapsed().as_secs_f64();
            elapsed > self.config.refresh_interval
        } else {
            true
        }
    }

    /// Get current noise model
    pub const fn current_model(&self) -> Option<&NoiseModel> {
        self.current_model.as_ref()
    }

    /// Get configuration
    pub const fn config(&self) -> &CalibrationConfig {
        &self.config
    }
}

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

    #[test]
    fn test_calibration_manager_creation() {
        let config = CalibrationConfig::default();
        let manager = NoiseCalibrationManager::new(config);
        assert!(manager.current_model().is_none());
        assert!(manager.needs_recalibration());
    }

    #[test]
    fn test_noise_characterization() {
        let config = CalibrationConfig::default();
        let mut manager = NoiseCalibrationManager::new(config);

        let result = manager.characterize_noise();
        assert!(result.is_ok());

        let model = result.expect("noise characterization should succeed");
        assert!(model.strength > 0.0);
        assert!(model.strength < 0.1);
        assert_eq!(model.qubit_parameters.len(), 10);
    }

    #[test]
    fn test_mitigation_strategy_selection() {
        let config = CalibrationConfig::default();
        let manager = NoiseCalibrationManager::new(config);

        // Low noise model
        let low_noise = NoiseModel {
            noise_type: NoiseType::Thermal,
            strength: 0.005,
            correlation_time: 2.0,
            correlation_length: 1.5,
            temperature: 0.015,
            qubit_parameters: vec![],
        };

        let strategy = manager
            .select_mitigation_strategy(&low_noise)
            .expect("low noise strategy selection should succeed");
        assert_eq!(strategy, MitigationStrategy::None);

        // High noise model
        let high_noise = NoiseModel {
            noise_type: NoiseType::Flux,
            strength: 0.08,
            correlation_time: 2.0,
            correlation_length: 1.5,
            temperature: 0.015,
            qubit_parameters: vec![],
        };

        let strategy = manager
            .select_mitigation_strategy(&high_noise)
            .expect("high noise strategy selection should succeed");
        assert_eq!(strategy, MitigationStrategy::ZeroNoiseExtrapolation);
    }

    #[test]
    fn test_adaptive_schedule_generation() {
        let config = CalibrationConfig::default();
        let manager = NoiseCalibrationManager::new(config);

        let noise_model = NoiseModel {
            noise_type: NoiseType::Thermal,
            strength: 0.03,
            correlation_time: 3.0,
            correlation_length: 1.5,
            temperature: 0.015,
            qubit_parameters: vec![],
        };

        let schedule = manager
            .generate_adaptive_schedule(&noise_model, 100.0)
            .expect("adaptive schedule generation should succeed");

        assert_eq!(schedule.len(), 100);
        assert_eq!(schedule[0].1, 0.0); // Start at s=0
        let last_schedule_value = schedule.last().expect("schedule should not be empty").1;
        assert!((last_schedule_value - 1.0).abs() < 1e-6); // End at s=1
    }

    #[test]
    fn test_noise_predictor() {
        let predictor = NoisePredictor::new(5, &[10, 10], 1);

        let input = Array1::from_vec(vec![0.01, 0.02, 0.015, 0.025, 0.018]);
        let output = predictor.predict(&input);

        assert_eq!(output.len(), 1);
    }

    #[test]
    fn test_noise_prediction() {
        let config = CalibrationConfig::default();
        let mut manager = NoiseCalibrationManager::new(config);

        // Build up history
        for _ in 0..20 {
            manager
                .characterize_noise()
                .expect("noise characterization should succeed");
        }

        let prediction = manager.predict_noise(10);
        assert!(prediction.is_ok());

        let pred = prediction.expect("noise prediction should succeed");
        assert_eq!(pred.predicted_strength.len(), 10);
        assert_eq!(pred.confidence_lower.len(), 10);
        assert_eq!(pred.confidence_upper.len(), 10);
        assert_eq!(pred.horizon, 10);
    }

    #[test]
    fn test_recalibration_check() {
        let config = CalibrationConfig::default().with_refresh_interval(1.0); // 1 second
        let mut manager = NoiseCalibrationManager::new(config);

        assert!(manager.needs_recalibration());

        manager
            .characterize_noise()
            .expect("noise characterization should succeed");
        assert!(!manager.needs_recalibration());

        // Wait for calibration to expire (in practice, would need actual time passage)
        // This test just verifies the logic exists
    }

    #[test]
    fn test_config_builder() {
        let config = CalibrationConfig::default()
            .with_characterization_samples(2000)
            .with_refresh_interval(7200.0)
            .with_ml_model_depth(5)
            .with_training_epochs(200);

        assert_eq!(config.characterization_samples, 2000);
        assert_eq!(config.refresh_interval, 7200.0);
        assert_eq!(config.ml_model_depth, 5);
        assert_eq!(config.training_epochs, 200);
    }

    #[test]
    fn test_noise_types() {
        let types = vec![
            NoiseType::Thermal,
            NoiseType::Flux,
            NoiseType::Charge,
            NoiseType::Crosstalk,
            NoiseType::Control,
            NoiseType::Readout,
        ];

        for noise_type in types {
            let model = NoiseModel {
                noise_type,
                strength: 0.02,
                correlation_time: 2.0,
                correlation_length: 1.5,
                temperature: 0.015,
                qubit_parameters: vec![],
            };

            assert_eq!(model.noise_type, noise_type);
        }
    }

    #[test]
    fn test_qubit_noise_parameters() {
        let params = QubitNoiseParameters {
            qubit_id: 0,
            t1: 25.0,
            t2: 12.0,
            readout_fidelity: 0.97,
            gate_fidelity: 0.98,
        };

        assert_eq!(params.qubit_id, 0);
        assert!(params.t1 > params.t2); // T1 >= T2 always
        assert!(params.readout_fidelity < 1.0);
        assert!(params.gate_fidelity < 1.0);
    }
}