//! Portable alpha / market-neutral strategy decomposition.
//!
//! Factor decomposition, beta-neutral adjustments, and incremental beta tracking
//! for prediction market portfolios.

use pyo3::prelude::*;

/// Single factor exposure measurement.
#[pyclass(frozen, get_all)]
#[derive(Debug, Clone)]
pub struct FactorExposure {
    pub factor_name: String,
    pub beta: f64,
    pub r_squared: f64,
    pub residual_vol: f64,
}

#[pymethods]
impl FactorExposure {
    #[new]
    fn new(factor_name: String, beta: f64, r_squared: f64, residual_vol: f64) -> Self {
        Self { factor_name, beta, r_squared, residual_vol }
    }

    fn __repr__(&self) -> String {
        format!("FactorExposure(factor='{}', beta={:.4}, R2={:.4})",
            self.factor_name, self.beta, self.r_squared)
    }
}

/// Alpha/beta decomposition result.
#[pyclass(frozen, get_all)]
#[derive(Debug, Clone)]
pub struct AlphaDecomposition {
    pub alpha: f64,
    pub beta: f64,
    pub r_squared: f64,
    pub residual_returns: Vec<f64>,
    pub tracking_error: f64,
    pub information_ratio: f64,
}

#[pymethods]
impl AlphaDecomposition {
    fn __repr__(&self) -> String {
        format!("AlphaDecomposition(alpha={:.4}, beta={:.4}, IR={:.4})",
            self.alpha, self.beta, self.information_ratio)
    }
}

/// Target neutrality specification.
#[pyclass(frozen, get_all)]
#[derive(Debug, Clone)]
pub struct NeutralityTarget {
    pub factor_name: String,
    pub target_beta: f64,
    pub tolerance: f64,
}

#[pymethods]
impl NeutralityTarget {
    #[new]
    #[pyo3(signature = (factor_name, target_beta=0.0, tolerance=0.05))]
    fn new(factor_name: String, target_beta: f64, tolerance: f64) -> Self {
        Self { factor_name, target_beta, tolerance }
    }

    fn __repr__(&self) -> String {
        format!("NeutralityTarget(factor='{}', target={:.4}, tol={:.4})",
            self.factor_name, self.target_beta, self.tolerance)
    }
}

/// Adjustment needed to achieve neutrality.
#[pyclass(frozen, get_all)]
#[derive(Debug, Clone)]
pub struct NeutralityAdjustment {
    pub factor_name: String,
    pub current_beta: f64,
    pub target_beta: f64,
    pub hedge_ratio: f64,
    pub is_within_tolerance: bool,
}

#[pymethods]
impl NeutralityAdjustment {
    fn __repr__(&self) -> String {
        format!("NeutralityAdjustment(factor='{}', hedge_ratio={:.4}, ok={})",
            self.factor_name, self.hedge_ratio, self.is_within_tolerance)
    }
}

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

/// Decompose portfolio returns into alpha + beta * factor_returns.
///
/// Uses OLS regression (Cholesky solve for single factor).
#[pyfunction]
pub fn beta_decompose(
    portfolio_returns: Vec<f64>,
    factor_returns: Vec<f64>,
) -> AlphaDecomposition {
    let n = portfolio_returns.len().min(factor_returns.len());
    if n < 3 {
        return AlphaDecomposition {
            alpha: 0.0, beta: 0.0, r_squared: 0.0,
            residual_returns: portfolio_returns,
            tracking_error: 0.0, information_ratio: 0.0,
        };
    }

    let mean_p = portfolio_returns.iter().take(n).sum::<f64>() / n as f64;
    let mean_f = factor_returns.iter().take(n).sum::<f64>() / n as f64;

    let mut cov_pf = 0.0;
    let mut var_f = 0.0;

    for i in 0..n {
        let dp = portfolio_returns[i] - mean_p;
        let df = factor_returns[i] - mean_f;
        cov_pf += dp * df;
        var_f += df * df;
    }

    let beta = if var_f > 1e-20 {
        finite_or_zero(cov_pf / var_f)
    } else {
        0.0
    };

    let alpha = finite_or_zero(mean_p - beta * mean_f);

    // Residuals
    let residuals: Vec<f64> = (0..n)
        .map(|i| portfolio_returns[i] - alpha - beta * factor_returns[i])
        .collect();

    // R-squared
    let ss_res: f64 = residuals.iter().map(|r| r * r).sum();
    let ss_tot: f64 = portfolio_returns.iter().take(n).map(|p| (p - mean_p).powi(2)).sum();
    let r_squared = if ss_tot > 1e-20 {
        finite_or_zero((1.0 - ss_res / ss_tot).clamp(0.0, 1.0))
    } else {
        0.0
    };

    // Tracking error = std of residuals
    let resid_mean: f64 = residuals.iter().sum::<f64>() / n as f64;
    let te_var: f64 = residuals.iter().map(|r| (r - resid_mean).powi(2)).sum::<f64>() / (n - 1).max(1) as f64;
    let tracking_error = finite_or_zero(te_var.sqrt());

    // Information ratio = alpha / tracking error
    let information_ratio = if tracking_error > 1e-10 {
        finite_or_zero(alpha / tracking_error)
    } else {
        0.0
    };

    AlphaDecomposition {
        alpha,
        beta,
        r_squared,
        residual_returns: residuals,
        tracking_error,
        information_ratio,
    }
}

/// Compute factor exposures for multiple factors.
#[pyfunction]
pub fn compute_factor_exposures(
    portfolio_returns: Vec<f64>,
    factor_names: Vec<String>,
    factor_returns_matrix: Vec<Vec<f64>>,
) -> Vec<FactorExposure> {
    let nf = factor_names.len().min(factor_returns_matrix.len());
    let mut exposures = Vec::with_capacity(nf);

    for i in 0..nf {
        let decomp = beta_decompose(portfolio_returns.clone(), factor_returns_matrix[i].clone());
        let resid_vol = if decomp.residual_returns.len() >= 2 {
            let m = decomp.residual_returns.iter().sum::<f64>() / decomp.residual_returns.len() as f64;
            let v = decomp.residual_returns.iter().map(|r| (r - m).powi(2)).sum::<f64>()
                / (decomp.residual_returns.len() - 1) as f64;
            finite_or_zero(v.sqrt())
        } else {
            0.0
        };

        exposures.push(FactorExposure {
            factor_name: factor_names[i].clone(),
            beta: decomp.beta,
            r_squared: decomp.r_squared,
            residual_vol: resid_vol,
        });
    }

    exposures
}

/// Compute hedge adjustments needed for market neutrality.
#[pyfunction]
pub fn neutrality_adjustment(
    current_betas: Vec<f64>,
    factor_names: Vec<String>,
    target_betas: Vec<f64>,
    tolerances: Vec<f64>,
) -> Vec<NeutralityAdjustment> {
    let n = current_betas.len()
        .min(factor_names.len())
        .min(target_betas.len())
        .min(tolerances.len());

    (0..n)
        .map(|i| {
            let diff = current_betas[i] - target_betas[i];
            let within = diff.abs() <= tolerances[i];
            NeutralityAdjustment {
                factor_name: factor_names[i].clone(),
                current_beta: current_betas[i],
                target_beta: target_betas[i],
                hedge_ratio: finite_or_zero(-diff),
                is_within_tolerance: within,
            }
        })
        .collect()
}

/// Incremental EWMA beta update: O(1) per new observation.
///
/// Returns (new_cov, new_var, new_beta).
#[pyfunction]
#[pyo3(signature = (prev_cov, prev_var, new_port_ret, new_factor_ret, decay=0.94))]
pub fn incremental_beta_update(
    prev_cov: f64,
    prev_var: f64,
    new_port_ret: f64,
    new_factor_ret: f64,
    decay: f64,
) -> (f64, f64, f64) {
    let d = decay.clamp(0.0, 1.0);
    let new_cov = d * prev_cov + (1.0 - d) * new_port_ret * new_factor_ret;
    let new_var = d * prev_var + (1.0 - d) * new_factor_ret * new_factor_ret;
    let new_beta = if new_var > 1e-20 {
        finite_or_zero(new_cov / new_var)
    } else {
        0.0
    };
    (new_cov, new_var, new_beta)
}

pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<FactorExposure>()?;
    m.add_class::<AlphaDecomposition>()?;
    m.add_class::<NeutralityTarget>()?;
    m.add_class::<NeutralityAdjustment>()?;
    m.add_function(wrap_pyfunction!(beta_decompose, m)?)?;
    m.add_function(wrap_pyfunction!(compute_factor_exposures, m)?)?;
    m.add_function(wrap_pyfunction!(neutrality_adjustment, m)?)?;
    m.add_function(wrap_pyfunction!(incremental_beta_update, m)?)?;
    Ok(())
}

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

    #[test]
    fn test_beta_decompose_perfect() {
        // portfolio = 2 * factor
        let factor = vec![0.01, -0.02, 0.03, -0.01, 0.02];
        let port: Vec<f64> = factor.iter().map(|f| 2.0 * f).collect();
        let result = beta_decompose(port, factor);
        assert!((result.beta - 2.0).abs() < 1e-6);
        assert!(result.r_squared > 0.99);
        assert!(result.alpha.abs() < 1e-6);
    }

    #[test]
    fn test_beta_decompose_with_alpha() {
        let factor = vec![0.01, -0.02, 0.03, -0.01, 0.02];
        let port: Vec<f64> = factor.iter().map(|f| 1.0 * f + 0.005).collect();
        let result = beta_decompose(port, factor);
        assert!((result.beta - 1.0).abs() < 1e-6);
        assert!((result.alpha - 0.005).abs() < 0.001);
    }

    #[test]
    fn test_beta_decompose_short() {
        let result = beta_decompose(vec![0.01, 0.02], vec![0.01, 0.02]);
        assert_eq!(result.beta, 0.0);
    }

    #[test]
    fn test_compute_factor_exposures() {
        let port = vec![0.02, -0.01, 0.03, -0.02, 0.01];
        let factor = vec![0.01, -0.005, 0.015, -0.01, 0.005];
        let exposures = compute_factor_exposures(
            port,
            vec!["market".into()],
            vec![factor],
        );
        assert_eq!(exposures.len(), 1);
        assert!(exposures[0].beta > 0.0);
    }

    #[test]
    fn test_neutrality_adjustment() {
        let adjs = neutrality_adjustment(
            vec![1.2, 0.3],
            vec!["market".into(), "vol".into()],
            vec![0.0, 0.0],
            vec![0.05, 0.05],
        );
        assert_eq!(adjs.len(), 2);
        assert!(!adjs[0].is_within_tolerance);
        assert!((adjs[0].hedge_ratio - (-1.2)).abs() < 1e-10);
    }

    #[test]
    fn test_neutrality_within_tolerance() {
        let adjs = neutrality_adjustment(
            vec![0.02],
            vec!["market".into()],
            vec![0.0],
            vec![0.05],
        );
        assert!(adjs[0].is_within_tolerance);
    }

    #[test]
    fn test_incremental_beta_update() {
        let (cov, var, beta) = incremental_beta_update(0.0, 0.0, 0.02, 0.01, 0.94);
        assert!(cov > 0.0);
        assert!(var > 0.0);
        assert!(beta > 0.0);
    }

    #[test]
    fn test_incremental_beta_update_zero_var() {
        let (_, _, beta) = incremental_beta_update(0.0, 0.0, 0.02, 0.0, 0.94);
        assert_eq!(beta, 0.0);
    }
}
