//! Correlation regime detection and contagion monitoring.
//!
//! Live correlation matrices, regime shift detection, contagion scoring,
//! and decorrelation finding for prediction market portfolios.

use pyo3::prelude::*;

/// A single entry in a correlation matrix.
#[pyclass(frozen, get_all)]
#[derive(Debug, Clone)]
pub struct CorrelationEntry {
    pub market_a: String,
    pub market_b: String,
    pub correlation: f64,
    pub num_observations: usize,
}

#[pymethods]
impl CorrelationEntry {
    #[new]
    fn new(market_a: String, market_b: String, correlation: f64, num_observations: usize) -> Self {
        Self { market_a, market_b, correlation, num_observations }
    }

    fn __repr__(&self) -> String {
        format!("CorrelationEntry({}<->{}, corr={:.4}, n={})",
            self.market_a, self.market_b, self.correlation, self.num_observations)
    }
}

/// Result of a regime shift detection.
#[pyclass(frozen, get_all)]
#[derive(Debug, Clone)]
pub struct RegimeShiftResult {
    pub shifted: bool,
    pub old_mean_corr: f64,
    pub new_mean_corr: f64,
    pub change_magnitude: f64,
    pub confidence: f64,
}

#[pymethods]
impl RegimeShiftResult {
    fn __repr__(&self) -> String {
        format!("RegimeShiftResult(shifted={}, change={:+.4}, conf={:.4})",
            self.shifted, self.change_magnitude, self.confidence)
    }
}

/// Contagion alert for correlated market stress.
#[pyclass(frozen, get_all)]
#[derive(Debug, Clone)]
pub struct ContagionAlert {
    pub source_market: String,
    pub affected_markets: Vec<String>,
    pub contagion_score: f64,
    pub avg_correlation: f64,
}

#[pymethods]
impl ContagionAlert {
    fn __repr__(&self) -> String {
        format!("ContagionAlert(source='{}', score={:.4}, affected={})",
            self.source_market, self.contagion_score, self.affected_markets.len())
    }
}

/// Result of decorrelation analysis.
#[pyclass(frozen, get_all)]
#[derive(Debug, Clone)]
pub struct DecorrelationResult {
    pub market_id: String,
    pub avg_correlation: f64,
    pub max_correlation: f64,
    pub diversification_score: f64,
}

#[pymethods]
impl DecorrelationResult {
    fn __repr__(&self) -> String {
        format!("DecorrelationResult(market='{}', avg_corr={:.4}, div_score={:.4})",
            self.market_id, self.avg_correlation, self.diversification_score)
    }
}

#[inline]
fn finite_or_zero(v: f64) -> f64 {
    if v.is_finite() { v } else { 0.0 }
}

/// Compute a live correlation matrix from return series.
///
/// Each row of `returns_matrix` is the return series for one market.
/// Returns vector of CorrelationEntry for all unique pairs.
#[pyfunction]
pub fn live_correlation_matrix(
    market_ids: Vec<String>,
    returns_matrix: Vec<Vec<f64>>,
) -> Vec<CorrelationEntry> {
    let n = market_ids.len().min(returns_matrix.len());
    let mut entries = Vec::new();

    for i in 0..n {
        for j in (i + 1)..n {
            let corr = pearson_correlation(&returns_matrix[i], &returns_matrix[j]);
            let obs = returns_matrix[i].len().min(returns_matrix[j].len());
            entries.push(CorrelationEntry {
                market_a: market_ids[i].clone(),
                market_b: market_ids[j].clone(),
                correlation: corr,
                num_observations: obs,
            });
        }
    }

    entries
}

/// Detect a regime shift by comparing recent vs historical correlation levels.
///
/// Uses a simple z-score test on mean correlations.
#[pyfunction]
#[pyo3(signature = (historical_corrs, recent_corrs, threshold=2.0))]
pub fn detect_regime_shift(
    historical_corrs: Vec<f64>,
    recent_corrs: Vec<f64>,
    threshold: f64,
) -> RegimeShiftResult {
    if historical_corrs.is_empty() || recent_corrs.is_empty() {
        return RegimeShiftResult {
            shifted: false,
            old_mean_corr: 0.0,
            new_mean_corr: 0.0,
            change_magnitude: 0.0,
            confidence: 0.0,
        };
    }

    let old_mean = mean(&historical_corrs);
    let new_mean = mean(&recent_corrs);
    let old_std = std_dev(&historical_corrs);

    let change = new_mean - old_mean;
    let confidence = if old_std > 1e-10 {
        finite_or_zero((change.abs() / old_std).min(10.0))
    } else {
        if change.abs() > 1e-10 { 10.0 } else { 0.0 }
    };

    RegimeShiftResult {
        shifted: confidence > threshold,
        old_mean_corr: finite_or_zero(old_mean),
        new_mean_corr: finite_or_zero(new_mean),
        change_magnitude: finite_or_zero(change),
        confidence,
    }
}

/// Compute contagion score for a source market.
///
/// Measures how strongly correlated a market is with others,
/// weighted by the strength of recent price moves.
#[pyfunction]
#[pyo3(signature = (source_market, market_ids, correlations, price_changes, threshold=0.5))]
pub fn contagion_score(
    source_market: String,
    market_ids: Vec<String>,
    correlations: Vec<f64>,
    price_changes: Vec<f64>,
    threshold: f64,
) -> ContagionAlert {
    let n = market_ids.len().min(correlations.len()).min(price_changes.len());
    let mut affected = Vec::new();
    let mut total_corr = 0.0;
    let mut count = 0;

    for i in 0..n {
        if market_ids[i] == source_market {
            continue;
        }
        let corr = correlations[i];
        if corr.abs() >= threshold {
            affected.push(market_ids[i].clone());
            total_corr += corr.abs();
            count += 1;
        }
    }

    let avg_corr = if count > 0 {
        finite_or_zero(total_corr / count as f64)
    } else {
        0.0
    };

    let source_change = price_changes
        .iter()
        .zip(market_ids.iter())
        .find(|(_, id)| **id == source_market)
        .map(|(c, _)| c.abs())
        .unwrap_or(0.0);

    let score = finite_or_zero(avg_corr * source_change * count as f64);

    ContagionAlert {
        source_market,
        affected_markets: affected,
        contagion_score: score,
        avg_correlation: avg_corr,
    }
}

/// Find markets with lowest average correlation (best diversifiers).
#[pyfunction]
#[pyo3(signature = (market_ids, returns_matrix, top_n=5))]
pub fn find_decorrelated(
    market_ids: Vec<String>,
    returns_matrix: Vec<Vec<f64>>,
    top_n: usize,
) -> Vec<DecorrelationResult> {
    let n = market_ids.len().min(returns_matrix.len());
    let mut results: Vec<DecorrelationResult> = Vec::new();

    for i in 0..n {
        let mut corrs = Vec::new();
        for j in 0..n {
            if i == j {
                continue;
            }
            let c = pearson_correlation(&returns_matrix[i], &returns_matrix[j]);
            corrs.push(c.abs());
        }

        let avg = if corrs.is_empty() { 0.0 } else { mean(&corrs) };
        let max = corrs.iter().cloned().fold(0.0f64, f64::max);
        let div_score = finite_or_zero(1.0 - avg);

        results.push(DecorrelationResult {
            market_id: market_ids[i].clone(),
            avg_correlation: finite_or_zero(avg),
            max_correlation: finite_or_zero(max),
            diversification_score: div_score,
        });
    }

    results.sort_by(|a, b| b.diversification_score.partial_cmp(&a.diversification_score)
        .unwrap_or(std::cmp::Ordering::Equal));
    results.truncate(top_n);
    results
}

/// Incremental O(1) correlation update using Welford-style running stats.
///
/// Given existing stats (n, sum_x, sum_y, sum_xy, sum_x2, sum_y2) and new
/// data point (x, y), returns updated correlation.
#[pyfunction]
pub fn incremental_correlation_update(
    n: u64,
    sum_x: f64,
    sum_y: f64,
    sum_xy: f64,
    sum_x2: f64,
    sum_y2: f64,
    new_x: f64,
    new_y: f64,
) -> (u64, f64, f64, f64, f64, f64, f64) {
    let n_new = n + 1;
    let sx = sum_x + new_x;
    let sy = sum_y + new_y;
    let sxy = sum_xy + new_x * new_y;
    let sx2 = sum_x2 + new_x * new_x;
    let sy2 = sum_y2 + new_y * new_y;

    let nf = n_new as f64;
    let num = nf * sxy - sx * sy;
    let den_x = (nf * sx2 - sx * sx).max(0.0).sqrt();
    let den_y = (nf * sy2 - sy * sy).max(0.0).sqrt();

    let corr = if den_x > 1e-10 && den_y > 1e-10 {
        finite_or_zero((num / (den_x * den_y)).clamp(-1.0, 1.0))
    } else {
        0.0
    };

    (n_new, sx, sy, sxy, sx2, sy2, corr)
}

// ---- helpers ----

fn pearson_correlation(x: &[f64], y: &[f64]) -> f64 {
    let n = x.len().min(y.len());
    if n < 2 {
        return 0.0;
    }

    let mx = x.iter().take(n).sum::<f64>() / n as f64;
    let my = y.iter().take(n).sum::<f64>() / n as f64;

    let mut cov = 0.0;
    let mut vx = 0.0;
    let mut vy = 0.0;

    for i in 0..n {
        let dx = x[i] - mx;
        let dy = y[i] - my;
        cov += dx * dy;
        vx += dx * dx;
        vy += dy * dy;
    }

    if vx < 1e-20 || vy < 1e-20 {
        return 0.0;
    }

    finite_or_zero((cov / (vx.sqrt() * vy.sqrt())).clamp(-1.0, 1.0))
}

fn mean(xs: &[f64]) -> f64 {
    if xs.is_empty() { return 0.0; }
    xs.iter().sum::<f64>() / xs.len() as f64
}

fn std_dev(xs: &[f64]) -> f64 {
    if xs.len() < 2 { return 0.0; }
    let m = mean(xs);
    let var = xs.iter().map(|x| (x - m).powi(2)).sum::<f64>() / (xs.len() - 1) as f64;
    var.sqrt()
}

pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<CorrelationEntry>()?;
    m.add_class::<RegimeShiftResult>()?;
    m.add_class::<ContagionAlert>()?;
    m.add_class::<DecorrelationResult>()?;
    m.add_function(wrap_pyfunction!(live_correlation_matrix, m)?)?;
    m.add_function(wrap_pyfunction!(detect_regime_shift, m)?)?;
    m.add_function(wrap_pyfunction!(contagion_score, m)?)?;
    m.add_function(wrap_pyfunction!(find_decorrelated, m)?)?;
    m.add_function(wrap_pyfunction!(incremental_correlation_update, m)?)?;
    Ok(())
}

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

    #[test]
    fn test_pearson_perfect_positive() {
        let x = vec![1.0, 2.0, 3.0, 4.0, 5.0];
        let y = vec![2.0, 4.0, 6.0, 8.0, 10.0];
        assert!((pearson_correlation(&x, &y) - 1.0).abs() < 1e-10);
    }

    #[test]
    fn test_pearson_perfect_negative() {
        let x = vec![1.0, 2.0, 3.0, 4.0, 5.0];
        let y = vec![10.0, 8.0, 6.0, 4.0, 2.0];
        assert!((pearson_correlation(&x, &y) - (-1.0)).abs() < 1e-10);
    }

    #[test]
    fn test_pearson_uncorrelated() {
        let x = vec![1.0, -1.0, 1.0, -1.0];
        let y = vec![1.0, 1.0, -1.0, -1.0];
        assert!(pearson_correlation(&x, &y).abs() < 0.1);
    }

    #[test]
    fn test_pearson_short() {
        assert_eq!(pearson_correlation(&[1.0], &[2.0]), 0.0);
        assert_eq!(pearson_correlation(&[], &[]), 0.0);
    }

    #[test]
    fn test_live_correlation_matrix() {
        let ids = vec!["A".into(), "B".into(), "C".into()];
        let returns = vec![
            vec![0.01, 0.02, -0.01, 0.03],
            vec![0.01, 0.02, -0.01, 0.03], // identical to A
            vec![-0.01, -0.02, 0.01, -0.03], // negative of A
        ];
        let entries = live_correlation_matrix(ids, returns);
        assert_eq!(entries.len(), 3); // 3 pairs: AB, AC, BC
        assert!((entries[0].correlation - 1.0).abs() < 1e-10); // A-B
        assert!((entries[1].correlation - (-1.0)).abs() < 1e-10); // A-C
    }

    #[test]
    fn test_detect_regime_shift_no_shift() {
        let hist = vec![0.3, 0.35, 0.28, 0.32, 0.31];
        let recent = vec![0.33, 0.29, 0.34];
        let result = detect_regime_shift(hist, recent, 2.0);
        assert!(!result.shifted);
    }

    #[test]
    fn test_detect_regime_shift_with_shift() {
        let hist = vec![0.3, 0.35, 0.28, 0.32, 0.31];
        let recent = vec![0.8, 0.85, 0.9]; // big jump
        let result = detect_regime_shift(hist, recent, 2.0);
        assert!(result.shifted);
        assert!(result.change_magnitude > 0.0);
    }

    #[test]
    fn test_detect_regime_shift_empty() {
        let result = detect_regime_shift(vec![], vec![0.5], 2.0);
        assert!(!result.shifted);
    }

    #[test]
    fn test_contagion_score_basic() {
        let alert = contagion_score(
            "A".into(),
            vec!["A".into(), "B".into(), "C".into()],
            vec![1.0, 0.8, 0.2], // B highly correlated with A, C not
            vec![0.05, 0.04, 0.01],
            0.5,
        );
        assert_eq!(alert.affected_markets.len(), 1); // only B
        assert!(alert.contagion_score > 0.0);
    }

    #[test]
    fn test_find_decorrelated() {
        let ids = vec!["A".into(), "B".into(), "C".into()];
        let returns = vec![
            vec![0.01, 0.02, -0.01],
            vec![0.01, 0.02, -0.01], // same as A
            vec![-0.02, 0.01, 0.03], // different
        ];
        let results = find_decorrelated(ids, returns, 2);
        assert!(!results.is_empty());
        // C should have highest diversification score
        assert_eq!(results[0].market_id, "C");
    }

    #[test]
    fn test_incremental_correlation_update() {
        let (n, sx, sy, sxy, sx2, sy2, corr) =
            incremental_correlation_update(0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 2.0);
        assert_eq!(n, 1);
        assert!((sx - 1.0).abs() < 1e-10);
        assert!((sy - 2.0).abs() < 1e-10);

        // Add more points
        let (n, sx, sy, sxy, sx2, sy2, corr) =
            incremental_correlation_update(n, sx, sy, sxy, sx2, sy2, 2.0, 4.0);
        assert_eq!(n, 2);
        assert!((corr - 1.0).abs() < 1e-10); // perfect correlation

        let (n, _sx, _sy, _sxy, _sx2, _sy2, corr) =
            incremental_correlation_update(n, sx, sy, sxy, sx2, sy2, 3.0, 6.0);
        assert_eq!(n, 3);
        assert!((corr - 1.0).abs() < 1e-10);
    }
}
