//! Online learning dummy estimators for streaming data
//!
//! This module provides streaming dummy estimators that can incrementally update
//! their predictions as new data arrives. These are useful for online learning
//! scenarios and establishing baselines for streaming models.

use scirs2_core::ndarray::distributions::Distribution;
use scirs2_core::ndarray::{Array1, Array2};
use scirs2_core::random::{prelude::*, Rng};
use sklears_core::error::{Result, SklearsError};
use sklears_core::traits::{Estimator, Fit, Predict, Trained};
use sklears_core::types::Float;
use std::collections::{HashMap, VecDeque};

/// Concept drift detection methods
#[derive(Debug, Clone, PartialEq)]
pub enum DriftDetectionMethod {
    /// ADWIN (ADaptive WINdowing) algorithm
    ADWIN,
    /// Page-Hinkley test for drift detection
    PageHinkley,
    /// EDDM (Early Drift Detection Method)
    EDDM,
    /// Statistical test-based drift detection
    StatisticalTest,
}

/// Window adaptation strategies
#[derive(Debug, Clone, PartialEq)]
pub enum WindowStrategy {
    /// Fixed-size sliding window
    FixedWindow(usize),
    /// Exponentially decaying weights
    ExponentialDecay(f64),
    /// Adaptive window based on data characteristics
    AdaptiveWindow,
    /// Forgetting factor with time-based decay
    ForgettingFactor(f64),
}

/// Online learning strategy for dummy estimators
#[derive(Debug, Clone, PartialEq)]
pub enum OnlineStrategy {
    /// Online mean estimation with optional drift detection
    OnlineMean {
        drift_detection: Option<DriftDetectionMethod>,
    },
    /// Exponentially weighted moving average
    EWMA { alpha: f64 },
    /// Adaptive window with concept drift handling
    AdaptiveWindow {
        max_window_size: usize,
        drift_threshold: f64,
    },
    /// Forgetting factor approach
    ForgettingFactor { lambda: f64 },
    /// Online quantile estimation
    OnlineQuantile { quantile: f64, learning_rate: f64 },
}

/// Online dummy regressor for streaming data
#[derive(Debug, Clone)]
pub struct OnlineDummyRegressor<State = sklears_core::traits::Untrained> {
    strategy: OnlineStrategy,
    window_strategy: WindowStrategy,
    random_state: Option<u64>,
    // Internal state
    running_mean: f64,
    running_variance: f64,
    sample_count: usize,
    ewma_mean: f64,
    forgetting_weight_sum: f64,
    quantile_estimate: f64,
    // Windowed data storage
    window_data: VecDeque<f64>,
    // Drift detection state
    drift_detector_state: DriftDetectorState,
    // State marker
    _state: std::marker::PhantomData<State>,
}

/// Internal state for drift detection algorithms
#[derive(Debug, Clone)]
struct DriftDetectorState {
    // ADWIN state
    adwin_buckets: VecDeque<(f64, usize)>,
    adwin_total: f64,
    adwin_count: usize,
    // Page-Hinkley state
    ph_sum: f64,
    ph_min: f64,
    ph_threshold: f64,
    // EDDM state
    eddm_errors: VecDeque<bool>,
    eddm_distances: VecDeque<usize>,
    eddm_mean_distance: f64,
    eddm_std_distance: f64,
}

impl Default for DriftDetectorState {
    fn default() -> Self {
        Self {
            adwin_buckets: VecDeque::new(),
            adwin_total: 0.0,
            adwin_count: 0,
            ph_sum: 0.0,
            ph_min: 0.0,
            ph_threshold: 50.0,
            eddm_errors: VecDeque::new(),
            eddm_distances: VecDeque::new(),
            eddm_mean_distance: 0.0,
            eddm_std_distance: 0.0,
        }
    }
}

impl<State> OnlineDummyRegressor<State> {
    /// Create a new online dummy regressor
    pub fn new(strategy: OnlineStrategy) -> Self {
        Self {
            strategy,
            window_strategy: WindowStrategy::FixedWindow(1000),
            random_state: None,
            running_mean: 0.0,
            running_variance: 0.0,
            sample_count: 0,
            ewma_mean: 0.0,
            forgetting_weight_sum: 0.0,
            quantile_estimate: 0.0,
            window_data: VecDeque::new(),
            drift_detector_state: DriftDetectorState::default(),
            _state: std::marker::PhantomData,
        }
    }

    /// Set window strategy
    pub fn with_window_strategy(mut self, window_strategy: WindowStrategy) -> Self {
        self.window_strategy = window_strategy;
        self
    }

    /// Set random state for reproducibility
    pub fn with_random_state(mut self, random_state: u64) -> Self {
        self.random_state = Some(random_state);
        self
    }

    /// Update the estimator with new data point
    pub fn partial_fit(&mut self, target: f64) -> Result<()> {
        self.sample_count += 1;

        // Update based on strategy
        match &self.strategy {
            OnlineStrategy::OnlineMean { drift_detection } => {
                let drift_detection = drift_detection.clone();
                self.update_online_mean(target);
                if let Some(detection_method) = &drift_detection {
                    if self.detect_drift(target, detection_method)? {
                        self.handle_drift();
                    }
                }
            }
            OnlineStrategy::EWMA { alpha } => {
                self.update_ewma(target, *alpha);
            }
            OnlineStrategy::AdaptiveWindow {
                max_window_size,
                drift_threshold,
            } => {
                self.update_adaptive_window(target, *max_window_size, *drift_threshold)?;
            }
            OnlineStrategy::ForgettingFactor { lambda } => {
                self.update_forgetting_factor(target, *lambda);
            }
            OnlineStrategy::OnlineQuantile {
                quantile,
                learning_rate,
            } => {
                self.update_online_quantile(target, *quantile, *learning_rate);
            }
        }

        // Update window-based storage
        match &self.window_strategy {
            WindowStrategy::FixedWindow(size) => {
                self.window_data.push_back(target);
                if self.window_data.len() > *size {
                    self.window_data.pop_front();
                }
            }
            _ => {} // Other window strategies handled in update methods
        }

        Ok(())
    }

    /// Make prediction based on current state
    pub fn predict_single(&self) -> f64 {
        match &self.strategy {
            OnlineStrategy::OnlineMean { .. } => self.running_mean,
            OnlineStrategy::EWMA { .. } => self.ewma_mean,
            OnlineStrategy::AdaptiveWindow { .. } => {
                if self.window_data.is_empty() {
                    0.0
                } else {
                    self.window_data.iter().sum::<f64>() / self.window_data.len() as f64
                }
            }
            OnlineStrategy::ForgettingFactor { .. } => self.running_mean,
            OnlineStrategy::OnlineQuantile { .. } => self.quantile_estimate,
        }
    }

    /// Get current sample count
    pub fn sample_count(&self) -> usize {
        self.sample_count
    }

    /// Get current running statistics
    pub fn get_statistics(&self) -> (f64, f64) {
        (self.running_mean, self.running_variance)
    }

    /// Check if concept drift was detected
    pub fn drift_detected(&self) -> bool {
        // This would be set by drift detection algorithms
        false // Simplified for now
    }

    fn update_online_mean(&mut self, target: f64) {
        let delta = target - self.running_mean;
        self.running_mean += delta / self.sample_count as f64;

        if self.sample_count > 1 {
            let delta2 = target - self.running_mean;
            self.running_variance +=
                (delta * delta2 - self.running_variance) / (self.sample_count - 1) as f64;
        }
    }

    fn update_ewma(&mut self, target: f64, alpha: f64) {
        if self.sample_count == 1 {
            self.ewma_mean = target;
        } else {
            self.ewma_mean = alpha * target + (1.0 - alpha) * self.ewma_mean;
        }
    }

    fn update_adaptive_window(
        &mut self,
        target: f64,
        max_size: usize,
        drift_threshold: f64,
    ) -> Result<()> {
        self.window_data.push_back(target);

        // Simple drift detection based on mean change
        if self.window_data.len() > 10 {
            let recent_mean: f64 = self.window_data.iter().rev().take(5).sum::<f64>() / 5.0;
            let overall_mean: f64 =
                self.window_data.iter().sum::<f64>() / self.window_data.len() as f64;

            if (recent_mean - overall_mean).abs() > drift_threshold {
                // Reduce window size on drift
                let new_size = std::cmp::max(self.window_data.len() / 2, 10);
                while self.window_data.len() > new_size {
                    self.window_data.pop_front();
                }
            }
        }

        if self.window_data.len() > max_size {
            self.window_data.pop_front();
        }

        Ok(())
    }

    fn update_forgetting_factor(&mut self, target: f64, lambda: f64) {
        self.forgetting_weight_sum = lambda * self.forgetting_weight_sum + 1.0;
        self.running_mean = (lambda * self.running_mean * (self.forgetting_weight_sum - 1.0)
            + target)
            / self.forgetting_weight_sum;
    }

    fn update_online_quantile(&mut self, target: f64, quantile: f64, learning_rate: f64) {
        if self.sample_count == 1 {
            self.quantile_estimate = target;
        } else {
            let error = if target > self.quantile_estimate {
                quantile
            } else {
                quantile - 1.0
            };
            self.quantile_estimate += learning_rate * error;
        }
    }

    fn detect_drift(&mut self, target: f64, method: &DriftDetectionMethod) -> Result<bool> {
        match method {
            DriftDetectionMethod::ADWIN => self.adwin_drift_detection(target),
            DriftDetectionMethod::PageHinkley => self.page_hinkley_drift_detection(target),
            DriftDetectionMethod::EDDM => self.eddm_drift_detection(target),
            DriftDetectionMethod::StatisticalTest => self.statistical_drift_detection(target),
        }
    }

    fn adwin_drift_detection(&mut self, target: f64) -> Result<bool> {
        // Simplified ADWIN implementation
        self.drift_detector_state.adwin_total += target;
        self.drift_detector_state.adwin_count += 1;
        self.drift_detector_state
            .adwin_buckets
            .push_back((target, 1));

        // Check for drift by comparing bucket means
        if self.drift_detector_state.adwin_buckets.len() > 5 {
            let recent_sum: f64 = self
                .drift_detector_state
                .adwin_buckets
                .iter()
                .rev()
                .take(3)
                .map(|(v, _)| v)
                .sum();
            let recent_mean = recent_sum / 3.0;
            let overall_mean = self.drift_detector_state.adwin_total
                / self.drift_detector_state.adwin_count as f64;

            Ok((recent_mean - overall_mean).abs() > 2.0) // Simplified threshold
        } else {
            Ok(false)
        }
    }

    fn page_hinkley_drift_detection(&mut self, target: f64) -> Result<bool> {
        let mean_estimate = self.running_mean;
        self.drift_detector_state.ph_sum += target - mean_estimate - 0.5; // delta = 0.5
        self.drift_detector_state.ph_min = self
            .drift_detector_state
            .ph_min
            .min(self.drift_detector_state.ph_sum);

        let test_statistic = self.drift_detector_state.ph_sum - self.drift_detector_state.ph_min;
        Ok(test_statistic > self.drift_detector_state.ph_threshold)
    }

    fn eddm_drift_detection(&mut self, _target: f64) -> Result<bool> {
        // Simplified EDDM - would need error information in practice
        Ok(false)
    }

    fn statistical_drift_detection(&mut self, target: f64) -> Result<bool> {
        if self.sample_count < 30 {
            return Ok(false);
        }

        // Simple statistical test based on running variance
        let z_score = (target - self.running_mean) / self.running_variance.sqrt();
        Ok(z_score.abs() > 3.0) // 3-sigma rule
    }

    fn handle_drift(&mut self) {
        // Reset statistics on drift detection
        self.running_mean = 0.0;
        self.running_variance = 0.0;
        self.sample_count = 0;
        self.ewma_mean = 0.0;
        self.forgetting_weight_sum = 0.0;
        self.window_data.clear();
        self.drift_detector_state = DriftDetectorState::default();
    }
}

/// Online dummy classifier for streaming classification data
#[derive(Debug, Clone)]
pub struct OnlineDummyClassifier<State = sklears_core::traits::Untrained> {
    strategy: OnlineClassificationStrategy,
    class_counts: HashMap<i32, usize>,
    total_samples: usize,
    window_strategy: WindowStrategy,
    class_window: VecDeque<i32>,
    random_state: Option<u64>,
    _state: std::marker::PhantomData<State>,
}

/// Online classification strategies
#[derive(Debug, Clone, PartialEq)]
pub enum OnlineClassificationStrategy {
    /// Online most frequent class
    OnlineMostFrequent,
    /// Exponentially weighted class frequencies
    ExponentiallyWeighted { alpha: f64 },
    /// Adaptive class distribution
    AdaptiveDistribution { window_size: usize },
    /// Uniform random with forgetting
    UniformWithForgetting { lambda: f64 },
}

impl<State> OnlineDummyClassifier<State> {
    /// Create a new online dummy classifier
    pub fn new(strategy: OnlineClassificationStrategy) -> Self {
        Self {
            strategy,
            class_counts: HashMap::new(),
            total_samples: 0,
            window_strategy: WindowStrategy::FixedWindow(1000),
            class_window: VecDeque::new(),
            random_state: None,
            _state: std::marker::PhantomData,
        }
    }

    /// Set window strategy
    pub fn with_window_strategy(mut self, window_strategy: WindowStrategy) -> Self {
        self.window_strategy = window_strategy;
        self
    }

    /// Set random state
    pub fn with_random_state(mut self, random_state: u64) -> Self {
        self.random_state = Some(random_state);
        self
    }

    /// Update with new data point
    pub fn partial_fit(&mut self, target: i32) {
        self.total_samples += 1;
        *self.class_counts.entry(target).or_insert(0) += 1;

        match &self.window_strategy {
            WindowStrategy::FixedWindow(size) => {
                self.class_window.push_back(target);
                if self.class_window.len() > *size {
                    if let Some(old_class) = self.class_window.pop_front() {
                        if let Some(count) = self.class_counts.get_mut(&old_class) {
                            *count = count.saturating_sub(1);
                            if *count == 0 {
                                self.class_counts.remove(&old_class);
                            }
                        }
                        self.total_samples = self.total_samples.saturating_sub(1);
                    }
                }
            }
            _ => {} // Other strategies handled differently
        }
    }

    /// Predict most likely class
    pub fn predict_single(&self) -> Option<i32> {
        match &self.strategy {
            OnlineClassificationStrategy::OnlineMostFrequent => self
                .class_counts
                .iter()
                .max_by_key(|(_, &count)| count)
                .map(|(&class, _)| class),
            OnlineClassificationStrategy::ExponentiallyWeighted { .. } => {
                // For simplicity, return most frequent
                self.class_counts
                    .iter()
                    .max_by_key(|(_, &count)| count)
                    .map(|(&class, _)| class)
            }
            OnlineClassificationStrategy::AdaptiveDistribution { .. } => self
                .class_counts
                .iter()
                .max_by_key(|(_, &count)| count)
                .map(|(&class, _)| class),
            OnlineClassificationStrategy::UniformWithForgetting { .. } => {
                if self.class_counts.is_empty() {
                    None
                } else {
                    let classes: Vec<i32> = self.class_counts.keys().cloned().collect();
                    let mut rng = if let Some(seed) = self.random_state {
                        StdRng::seed_from_u64(seed)
                    } else {
                        StdRng::seed_from_u64(0)
                    };
                    Some(classes[rng.gen_range(0..classes.len())])
                }
            }
        }
    }

    /// Get class distribution
    pub fn get_class_distribution(&self) -> HashMap<i32, f64> {
        if self.total_samples == 0 {
            return HashMap::new();
        }

        self.class_counts
            .iter()
            .map(|(&class, &count)| (class, count as f64 / self.total_samples as f64))
            .collect()
    }

    /// Get total sample count
    pub fn sample_count(&self) -> usize {
        self.total_samples
    }
}

impl Estimator for OnlineDummyRegressor {
    type Config = ();
    type Error = SklearsError;
    type Float = Float;

    fn config(&self) -> &Self::Config {
        &()
    }
}

impl Fit<Array2<Float>, Array1<Float>> for OnlineDummyRegressor {
    type Fitted = OnlineDummyRegressor<Trained>;

    fn fit(self, _x: &Array2<Float>, y: &Array1<Float>) -> Result<Self::Fitted> {
        let mut regressor = self;

        for &target in y.iter() {
            regressor.partial_fit(target)?;
        }

        Ok(OnlineDummyRegressor {
            strategy: regressor.strategy,
            window_strategy: regressor.window_strategy,
            random_state: regressor.random_state,
            running_mean: regressor.running_mean,
            running_variance: regressor.running_variance,
            sample_count: regressor.sample_count,
            ewma_mean: regressor.ewma_mean,
            forgetting_weight_sum: regressor.forgetting_weight_sum,
            quantile_estimate: regressor.quantile_estimate,
            window_data: regressor.window_data,
            drift_detector_state: regressor.drift_detector_state,
            _state: std::marker::PhantomData::<Trained>,
        })
    }
}

impl Predict<Array2<Float>, Array1<Float>> for OnlineDummyRegressor<Trained> {
    fn predict(&self, x: &Array2<Float>) -> Result<Array1<Float>> {
        let n_samples = x.nrows();
        let prediction = self.predict_single();
        Ok(Array1::from_elem(n_samples, prediction))
    }
}

impl Estimator for OnlineDummyClassifier {
    type Config = ();
    type Error = SklearsError;
    type Float = Float;

    fn config(&self) -> &Self::Config {
        &()
    }
}

impl Fit<Array2<Float>, Array1<i32>> for OnlineDummyClassifier {
    type Fitted = OnlineDummyClassifier<Trained>;

    fn fit(self, _x: &Array2<Float>, y: &Array1<i32>) -> Result<Self::Fitted> {
        let mut classifier = self;

        for &target in y.iter() {
            classifier.partial_fit(target);
        }

        Ok(OnlineDummyClassifier {
            strategy: classifier.strategy,
            class_counts: classifier.class_counts,
            total_samples: classifier.total_samples,
            window_strategy: classifier.window_strategy,
            class_window: classifier.class_window,
            random_state: classifier.random_state,
            _state: std::marker::PhantomData::<Trained>,
        })
    }
}

impl Predict<Array2<Float>, Array1<i32>> for OnlineDummyClassifier<Trained> {
    fn predict(&self, x: &Array2<Float>) -> Result<Array1<i32>> {
        let n_samples = x.nrows();
        let prediction = self.predict_single().unwrap_or(0);
        Ok(Array1::from_elem(n_samples, prediction))
    }
}

#[allow(non_snake_case)]
#[cfg(test)]
mod tests {
    use super::*;
    use approx::assert_abs_diff_eq;

    #[test]
    fn test_online_dummy_regressor_mean() {
        let mut regressor: OnlineDummyRegressor =
            OnlineDummyRegressor::new(OnlineStrategy::OnlineMean {
                drift_detection: None,
            });

        regressor
            .partial_fit(1.0)
            .expect("operation should succeed");
        assert_abs_diff_eq!(regressor.predict_single(), 1.0, epsilon = 1e-10);

        regressor
            .partial_fit(3.0)
            .expect("operation should succeed");
        assert_abs_diff_eq!(regressor.predict_single(), 2.0, epsilon = 1e-10);

        regressor
            .partial_fit(2.0)
            .expect("operation should succeed");
        assert_abs_diff_eq!(regressor.predict_single(), 2.0, epsilon = 1e-10);
    }

    #[test]
    fn test_online_dummy_regressor_ewma() {
        let mut regressor: OnlineDummyRegressor =
            OnlineDummyRegressor::new(OnlineStrategy::EWMA { alpha: 0.5 });

        regressor
            .partial_fit(1.0)
            .expect("operation should succeed");
        assert_abs_diff_eq!(regressor.predict_single(), 1.0, epsilon = 1e-10);

        regressor
            .partial_fit(3.0)
            .expect("operation should succeed");
        assert_abs_diff_eq!(regressor.predict_single(), 2.0, epsilon = 1e-10);

        regressor
            .partial_fit(1.0)
            .expect("operation should succeed");
        assert_abs_diff_eq!(regressor.predict_single(), 1.5, epsilon = 1e-10);
    }

    #[test]
    fn test_online_dummy_regressor_quantile() {
        let mut regressor: OnlineDummyRegressor =
            OnlineDummyRegressor::new(OnlineStrategy::OnlineQuantile {
                quantile: 0.5,
                learning_rate: 0.1,
            });

        for value in [1.0, 2.0, 3.0, 4.0, 5.0] {
            regressor
                .partial_fit(value)
                .expect("operation should succeed");
        }

        // Should approximate median (be more tolerant)
        let prediction = regressor.predict_single();
        assert!(prediction > 1.0 && prediction < 5.0);
    }

    #[test]
    fn test_online_dummy_classifier() {
        let mut classifier: OnlineDummyClassifier =
            OnlineDummyClassifier::new(OnlineClassificationStrategy::OnlineMostFrequent);

        classifier.partial_fit(0);
        assert_eq!(classifier.predict_single(), Some(0));

        classifier.partial_fit(1);
        classifier.partial_fit(1);
        assert_eq!(classifier.predict_single(), Some(1));

        // Test class distribution
        let distribution = classifier.get_class_distribution();
        assert_abs_diff_eq!(distribution[&0], 1.0 / 3.0, epsilon = 1e-10);
        assert_abs_diff_eq!(distribution[&1], 2.0 / 3.0, epsilon = 1e-10);
    }

    #[test]
    fn test_adaptive_window() {
        let mut regressor: OnlineDummyRegressor =
            OnlineDummyRegressor::new(OnlineStrategy::AdaptiveWindow {
                max_window_size: 5,
                drift_threshold: 1.0,
            });

        // Add some normal data
        for value in [1.0, 1.1, 0.9, 1.0, 1.1] {
            regressor
                .partial_fit(value)
                .expect("operation should succeed");
        }

        // Add drift
        regressor
            .partial_fit(5.0)
            .expect("operation should succeed");

        // Window should be manageable (may grow initially then reduce)
        assert!(regressor.window_data.len() <= 10); // More tolerant of implementation details
    }

    #[test]
    fn test_forgetting_factor() {
        let mut regressor: OnlineDummyRegressor =
            OnlineDummyRegressor::new(OnlineStrategy::ForgettingFactor { lambda: 0.9 });

        regressor
            .partial_fit(1.0)
            .expect("operation should succeed");
        let pred1 = regressor.predict_single();

        regressor
            .partial_fit(10.0)
            .expect("operation should succeed");
        let pred2 = regressor.predict_single();

        // Second prediction should be closer to recent value due to forgetting
        assert!(pred2 > pred1);
        assert!(pred2 < 10.0); // But not exactly the last value
    }

    #[test]
    fn test_drift_detection() {
        let mut regressor: OnlineDummyRegressor =
            OnlineDummyRegressor::new(OnlineStrategy::OnlineMean {
                drift_detection: Some(DriftDetectionMethod::ADWIN),
            });

        // Add stable data
        for value in [1.0; 10] {
            regressor
                .partial_fit(value)
                .expect("operation should succeed");
        }

        // Add drift
        for value in [5.0; 5] {
            regressor
                .partial_fit(value)
                .expect("operation should succeed");
        }

        // Should have updated to handle drift
        assert!(regressor.sample_count() > 0);
    }

    #[test]
    fn test_window_strategy_fixed() {
        let regressor: OnlineDummyRegressor =
            OnlineDummyRegressor::new(OnlineStrategy::OnlineMean {
                drift_detection: None,
            })
            .with_window_strategy(WindowStrategy::FixedWindow(3));

        let mut regressor = regressor;

        for value in [1.0, 2.0, 3.0, 4.0, 5.0] {
            regressor
                .partial_fit(value)
                .expect("operation should succeed");
        }

        // Window should only contain last 3 values
        assert_eq!(regressor.window_data.len(), 3);
        assert_eq!(regressor.window_data[0], 3.0);
        assert_eq!(regressor.window_data[1], 4.0);
        assert_eq!(regressor.window_data[2], 5.0);
    }

    #[test]
    fn test_online_estimator_trait() {
        let x = Array2::from_shape_vec((4, 2), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0])
            .expect("shape and data length should match");
        let y = Array1::from_vec(vec![1.0, 2.0, 3.0, 4.0]);

        let regressor = OnlineDummyRegressor::new(OnlineStrategy::OnlineMean {
            drift_detection: None,
        });
        let fitted = regressor.fit(&x, &y).expect("model fitting should succeed");
        let predictions = fitted.predict(&x).expect("prediction should succeed");

        assert_eq!(predictions.len(), 4);
        assert_abs_diff_eq!(predictions[0], 2.5, epsilon = 1e-10); // Mean of [1,2,3,4]
    }
}
