//! Stationarity tests for time series
#![allow(missing_docs)]

use super::timeseries::TimeSeries;
use ndarray::{Array1, Array2};
use so_core::error::{Error, Result};
use so_linalg;

/// Stationarity test results
#[derive(Debug, Clone)]
pub struct StationarityTestResults {
    pub test_statistic: f64,
    pub p_value: f64,
    pub critical_values: Vec<(String, f64)>,
    pub lags_used: usize,
    pub n_obs: usize,
    pub alternative: String,
    pub stationary: bool,
}

impl StationarityTestResults {
    /// Create summary string
    pub fn summary(&self) -> String {
        let mut summary = String::new();
        summary.push_str(&format!("Test Statistic: {:.4}\n", self.test_statistic));
        summary.push_str(&format!("p-value: {:.4}\n", self.p_value));
        summary.push_str(&format!("Lags Used: {}\n", self.lags_used));
        summary.push_str(&format!("Number of Observations: {}\n", self.n_obs));

        summary.push_str("Critical Values:\n");
        for (level, value) in &self.critical_values {
            summary.push_str(&format!("  {}%: {:.4}\n", level, value));
        }

        summary.push_str(&format!("Alternative: {}\n", self.alternative));
        summary.push_str(&format!(
            "Result: Series is {}stationary\n",
            if self.stationary { "" } else { "not " }
        ));

        summary
    }
}

/// Augmented Dickey-Fuller test
pub struct ADFTest {
    lags: Option<usize>,
    trend: ADFTrend,
    max_lags: usize,
}

/// Trend specification for ADF test
#[derive(Debug, Clone, Copy)]
pub enum ADFTrend {
    /// No constant, no trend
    None,
    /// Constant (drift) only
    Constant,
    /// Constant and linear trend
    Trend,
}

impl Default for ADFTest {
    fn default() -> Self {
        Self {
            lags: None,
            trend: ADFTrend::Constant,
            max_lags: 10,
        }
    }
}

impl ADFTest {
    /// Create new ADF test
    pub fn new() -> Self {
        Self::default()
    }

    /// Set number of lags (if None, use automatic selection)
    pub fn lags(mut self, lags: Option<usize>) -> Self {
        self.lags = lags;
        self
    }

    /// Set trend specification
    pub fn trend(mut self, trend: ADFTrend) -> Self {
        self.trend = trend;
        self
    }

    /// Set maximum lags for automatic selection
    pub fn max_lags(mut self, max_lags: usize) -> Self {
        self.max_lags = max_lags;
        self
    }

    /// Perform ADF test
    pub fn test(&self, ts: &TimeSeries) -> Result<StationarityTestResults> {
        let values = ts.values();
        let n = values.len();

        if n < 20 {
            return Err(Error::DataError(
                "ADF test requires at least 20 observations".to_string(),
            ));
        }

        // Determine optimal lag length
        let lags = match self.lags {
            Some(l) => l.min(self.max_lags.min(n / 4)),
            None => self.select_lags(values),
        };

        // Prepare regression data
        let (y, X) = self.prepare_regression(values, lags);

        // Perform regression
        let (test_stat, p_value) = self.perform_regression(&y, &X, n - lags - 1)?;

        // Critical values from MacKinnon (2010)
        let critical_values = vec![
            ("1".to_string(), -3.43),
            ("5".to_string(), -2.86),
            ("10".to_string(), -2.57),
        ];

        let stationary = test_stat < critical_values[1].1; // 5% critical value

        Ok(StationarityTestResults {
            test_statistic: test_stat,
            p_value,
            critical_values,
            lags_used: lags,
            n_obs: n,
            alternative: "stationary".to_string(),
            stationary,
        })
    }

    /// Select optimal lag length using AIC
    fn select_lags(&self, values: &Array1<f64>) -> usize {
        let n = values.len();
        let max_lags = self.max_lags.min(n / 4);

        let mut best_aic = f64::INFINITY;
        let mut best_lags = 0;

        for lags in 0..=max_lags {
            let (y, X) = self.prepare_regression(values, lags);
            if y.len() > 0 && y.len() > X.ncols() {
                if let Ok((_, _, aic)) = self.calculate_aic(&y, &X) {
                    if aic < best_aic {
                        best_aic = aic;
                        best_lags = lags;
                    }
                }
            }
        }

        best_lags
    }

    /// Prepare regression data for ADF test
    fn prepare_regression(&self, values: &Array1<f64>, lags: usize) -> (Array1<f64>, Array2<f64>) {
        let n = values.len();
        let reg_n = n - lags - 1;

        // Dependent variable: Δy_t
        let mut y = Array1::zeros(reg_n);
        for i in 0..reg_n {
            let idx = i + lags + 1;
            y[i] = values[idx] - values[idx - 1];
        }

        // Independent variables
        let n_cols = match self.trend {
            ADFTrend::None => lags + 1,     // y_{t-1} + lags of Δy
            ADFTrend::Constant => lags + 2, // + constant
            ADFTrend::Trend => lags + 3,    // + constant + trend
        };

        let mut X = Array2::zeros((reg_n, n_cols));

        // Fill design matrix
        for i in 0..reg_n {
            let mut col = 0;
            let idx = i + lags;

            // y_{t-1}
            X[(i, col)] = values[idx];
            col += 1;

            // Lags of Δy
            for j in 0..lags {
                X[(i, col)] = values[idx - j] - values[idx - j - 1];
                col += 1;
            }

            // Constant
            if matches!(self.trend, ADFTrend::Constant | ADFTrend::Trend) {
                X[(i, col)] = 1.0;
                col += 1;
            }

            // Trend
            if matches!(self.trend, ADFTrend::Trend) {
                X[(i, col)] = (idx + 1) as f64;
            }
        }

        (y, X)
    }

    /// Perform regression and calculate test statistic
    fn perform_regression(
        &self,
        y: &Array1<f64>,
        X: &Array2<f64>,
        df_resid: usize,
    ) -> Result<(f64, f64)> {
        // Solve normal equations
        let XtX = X.t().dot(X);
        let Xty = X.t().dot(y);

        let beta = so_linalg::solve(&XtX, &Xty)
            .map_err(|e| Error::LinearAlgebraError(format!("ADF regression failed: {}", e)))?;

        // Calculate residuals and variance
        let y_hat = X.dot(&beta);
        let residuals = y - &y_hat;
        let rss: f64 = residuals.iter().map(|&r| r.powi(2)).sum();
        let sigma2 = rss / df_resid as f64;

        // Test statistic for y_{t-1} coefficient (first element)
        let se = (sigma2 / XtX[(0, 0)]).sqrt();
        let test_stat = beta[0] / se;

        // Approximate p-value using MacKinnon (2010) approximation
        let p_value = self.mackinnon_pvalue(test_stat);

        Ok((test_stat, p_value))
    }

    /// Calculate AIC for lag selection
    fn calculate_aic(&self, y: &Array1<f64>, X: &Array2<f64>) -> Result<(f64, f64, f64)> {
        let XtX = X.t().dot(X);
        let Xty = X.t().dot(y);

        let beta = so_linalg::solve(&XtX, &Xty)
            .map_err(|e| Error::LinearAlgebraError(format!("Regression failed: {}", e)))?;

        let y_hat = X.dot(&beta);
        let residuals = y - &y_hat;
        let rss: f64 = residuals.iter().map(|&r| r.powi(2)).sum();

        let n = y.len();
        let k = X.ncols();

        let sigma2 = rss / (n - k) as f64;
        let log_lik =
            -0.5 * n as f64 * (2.0 * std::f64::consts::PI * sigma2).ln() - rss / (2.0 * sigma2);
        let aic = 2.0 * k as f64 - 2.0 * log_lik;
        let bic = (n as f64).ln() * k as f64 - 2.0 * log_lik;

        Ok((log_lik, aic, bic))
    }

    /// MacKinnon (2010) approximate p-values for ADF test
    fn mackinnon_pvalue(&self, test_stat: f64) -> f64 {
        // Simplified approximation - in practice should use full MacKinnon tables
        let x = test_stat;

        // Coefficients for constant trend case
        let inf_crit = -2.86; // 5% critical value

        if x < inf_crit {
            // More extreme than critical value
            0.01 // Approximate p-value < 0.05
        } else {
            0.5 // Approximate p-value > 0.05
        }
    }
}

/// KPSS test for stationarity
pub struct KPSSTest {
    lags: Option<usize>,
    trend: KPSSTrend,
}

/// Trend specification for KPSS test
#[derive(Debug, Clone, Copy)]
pub enum KPSSTrend {
    /// Level stationary
    Level,
    /// Trend stationary
    Trend,
}

impl Default for KPSSTest {
    fn default() -> Self {
        Self {
            lags: None,
            trend: KPSSTrend::Level,
        }
    }
}

impl KPSSTest {
    /// Create new KPSS test
    pub fn new() -> Self {
        Self::default()
    }

    /// Set number of lags
    pub fn lags(mut self, lags: Option<usize>) -> Self {
        self.lags = lags;
        self
    }

    /// Set trend specification
    pub fn trend(mut self, trend: KPSSTrend) -> Self {
        self.trend = trend;
        self
    }

    /// Perform KPSS test
    pub fn test(&self, ts: &TimeSeries) -> Result<StationarityTestResults> {
        let values = ts.values();
        let n = values.len();

        if n < 10 {
            return Err(Error::DataError(
                "KPSS test requires at least 10 observations".to_string(),
            ));
        }

        // Determine lag length
        let lags = self.lags.unwrap_or_else(|| self.default_lags(n));

        // Calculate test statistic
        let test_stat = self.calculate_statistic(values, lags);

        // Critical values for KPSS test
        let critical_values = match self.trend {
            KPSSTrend::Level => vec![
                ("0.10".to_string(), 0.347),
                ("0.05".to_string(), 0.463),
                ("0.025".to_string(), 0.574),
                ("0.01".to_string(), 0.739),
            ],
            KPSSTrend::Trend => vec![
                ("0.10".to_string(), 0.119),
                ("0.05".to_string(), 0.146),
                ("0.025".to_string(), 0.176),
                ("0.01".to_string(), 0.216),
            ],
        };

        // Determine if stationary (null hypothesis is stationary)
        let stationary = test_stat < critical_values[1].1; // Compare to 5% critical value

        Ok(StationarityTestResults {
            test_statistic: test_stat,
            p_value: self.calculate_pvalue(test_stat),
            critical_values,
            lags_used: lags,
            n_obs: n,
            alternative: "non-stationary".to_string(),
            stationary,
        })
    }

    /// Default lag length based on Schwert (1989)
    fn default_lags(&self, n: usize) -> usize {
        ((12.0 * (n as f64 / 100.0).powf(0.25)) as usize).min(n / 4)
    }

    /// Calculate KPSS test statistic
    fn calculate_statistic(&self, values: &Array1<f64>, lags: usize) -> f64 {
        let n = values.len();

        // Calculate partial sums
        let mean = values.mean().unwrap_or(0.0);
        let mut partial_sums = Array1::zeros(n);
        let mut sum = 0.0;

        for i in 0..n {
            sum += values[i] - mean;
            partial_sums[i] = sum;
        }

        // Calculate denominator (long-run variance)
        let s2 = self.long_run_variance(&partial_sums, lags);

        // Test statistic
        partial_sums.iter().map(|&s| s.powi(2)).sum::<f64>() / (n as f64 * s2)
    }

    /// Calculate long-run variance using Bartlett kernel
    fn long_run_variance(&self, series: &Array1<f64>, lags: usize) -> f64 {
        let n = series.len();
        let mean = series.mean().unwrap_or(0.0);

        // Sample autocovariances
        let mut gamma = vec![0.0; lags + 1];

        for lag in 0..=lags {
            let mut sum = 0.0;
            for t in lag..n {
                sum += (series[t] - mean) * (series[t - lag] - mean);
            }
            gamma[lag] = sum / n as f64;
        }

        // Bartlett kernel weights
        let mut s2 = gamma[0];
        for lag in 1..=lags {
            let weight = 1.0 - lag as f64 / lags as f64;
            s2 += 2.0 * weight * gamma[lag];
        }

        s2.max(0.0)
    }

    /// Calculate approximate p-value
    fn calculate_pvalue(&self, test_stat: f64) -> f64 {
        // Simplified approximation
        if test_stat < 0.1 {
            0.9
        } else if test_stat < 0.5 {
            0.5
        } else {
            0.1
        }
    }
}

/// Phillips-Perron test
pub struct PPTest {
    lags: usize,
}

impl Default for PPTest {
    fn default() -> Self {
        Self { lags: 4 }
    }
}

impl PPTest {
    /// Create new PP test
    pub fn new() -> Self {
        Self::default()
    }

    /// Set number of lags
    pub fn lags(mut self, lags: usize) -> Self {
        self.lags = lags;
        self
    }

    /// Perform Phillips-Perron test
    pub fn test(&self, ts: &TimeSeries) -> Result<StationarityTestResults> {
        // Implementation would be similar to ADF but with non-parametric correction
        // For now, return placeholder
        Ok(StationarityTestResults {
            test_statistic: -2.5,
            p_value: 0.12,
            critical_values: vec![
                ("1".to_string(), -3.43),
                ("5".to_string(), -2.86),
                ("10".to_string(), -2.57),
            ],
            lags_used: self.lags,
            n_obs: ts.len(),
            alternative: "stationary".to_string(),
            stationary: false,
        })
    }
}

/// Trait for all stationarity tests
pub trait StationarityTest {
    fn test(&self, ts: &TimeSeries) -> Result<StationarityTestResults>;
}

impl StationarityTest for ADFTest {
    fn test(&self, ts: &TimeSeries) -> Result<StationarityTestResults> {
        self.test(ts)
    }
}

impl StationarityTest for KPSSTest {
    fn test(&self, ts: &TimeSeries) -> Result<StationarityTestResults> {
        self.test(ts)
    }
}

impl StationarityTest for PPTest {
    fn test(&self, ts: &TimeSeries) -> Result<StationarityTestResults> {
        self.test(ts)
    }
}
