//! Cross-Impact Matrix estimation for multi-market price impact analysis.
//!
//! Estimates how order flow in one market affects prices in another.
//! Uses OLS cross-regression: returns_i = sum_j beta_{ij} * signed_volume_j + eps.
//! The resulting matrix is projected to positive semi-definite via
//! eigenvalue clipping (negative eigenvalues set to zero).
//!
//! Eigenvalue decomposition is implemented via power iteration with deflation.
//!
//! All math from scratch in Rust -- no external crate dependencies beyond PyO3.
//!
//! Ultra tier required.

use pyo3::prelude::*;
use std::sync::Mutex;

// ===========================================================================
// Helpers
// ===========================================================================

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

// ===========================================================================
// Frozen result type
// ===========================================================================

/// Estimated cross-impact matrix with eigenvalue analysis.
#[pyclass(frozen)]
#[derive(Debug, Clone)]
pub struct CrossImpactMatrix {
    /// Cross-impact matrix (N x N), row i = impact of all markets' flow on market i's return.
    #[pyo3(get)]
    pub matrix: Vec<Vec<f64>>,
    /// Market names corresponding to rows/columns.
    #[pyo3(get)]
    pub market_names: Vec<String>,
    /// Eigenvalues of the cross-impact matrix (sorted descending).
    #[pyo3(get)]
    pub eigenvalues: Vec<f64>,
    /// Whether the matrix is positive semi-definite (after projection).
    #[pyo3(get)]
    pub is_psd: bool,
}

#[pymethods]
impl CrossImpactMatrix {
    fn __repr__(&self) -> String {
        let n = self.matrix.len();
        let top_eig: Vec<String> = self
            .eigenvalues
            .iter()
            .take(3)
            .map(|e| format!("{:.4}", e))
            .collect();
        format!(
            "CrossImpactMatrix(n_markets={}, eigenvalues=[{}{}], psd={})",
            n,
            top_eig.join(", "),
            if self.eigenvalues.len() > 3 { ", ..." } else { "" },
            self.is_psd
        )
    }
}

// ===========================================================================
// Linear algebra helpers
// ===========================================================================

type Mat = Vec<Vec<f64>>;

fn mat_zeros(n: usize, m: usize) -> Mat {
    vec![vec![0.0; m]; n]
}

#[allow(dead_code)]
fn mat_eye(n: usize) -> Mat {
    let mut m = mat_zeros(n, n);
    for i in 0..n {
        m[i][i] = 1.0;
    }
    m
}

fn mat_transpose(a: &Mat) -> Mat {
    if a.is_empty() { return vec![]; }
    let n = a.len();
    let m = a[0].len();
    let mut t = mat_zeros(m, n);
    for i in 0..n {
        for j in 0..m {
            t[j][i] = a[i][j];
        }
    }
    t
}

fn mat_mul(a: &Mat, b: &Mat) -> Mat {
    let n = a.len();
    if n == 0 { return vec![]; }
    let k = a[0].len();
    let m = b[0].len();
    let mut c = mat_zeros(n, m);
    for i in 0..n {
        for j in 0..m {
            let mut s = 0.0;
            for l in 0..k {
                s += a[i][l] * b[l][j];
            }
            c[i][j] = finite_or_zero(s);
        }
    }
    c
}

fn mat_vec_mul(a: &Mat, v: &[f64]) -> Vec<f64> {
    let n = a.len();
    let mut w = vec![0.0; n];
    for i in 0..n {
        let mut s = 0.0;
        for j in 0..a[i].len() {
            s += a[i][j] * v[j];
        }
        w[i] = finite_or_zero(s);
    }
    w
}

fn vec_dot(a: &[f64], b: &[f64]) -> f64 {
    let mut s = 0.0;
    for i in 0..a.len() {
        s += a[i] * b[i];
    }
    finite_or_zero(s)
}

fn vec_norm(a: &[f64]) -> f64 {
    vec_dot(a, a).sqrt()
}

fn vec_scale(a: &[f64], s: f64) -> Vec<f64> {
    a.iter().map(|&x| finite_or_zero(x * s)).collect()
}

#[allow(dead_code)]
fn vec_sub(a: &[f64], b: &[f64]) -> Vec<f64> {
    a.iter().zip(b.iter()).map(|(&x, &y)| x - y).collect()
}

// ===========================================================================
// OLS regression
// ===========================================================================

/// Solve X'X * beta = X'y via Cholesky decomposition.
/// X is (n x p), y is (n,).
fn ols_solve(x: &Mat, y: &[f64]) -> Vec<f64> {
    let p = x[0].len();
    let xtx = mat_mul(&mat_transpose(x), x);
    let xty_mat = mat_transpose(x);
    let xty: Vec<f64> = mat_vec_mul(&xty_mat, y);

    // Regularized: add small ridge to diagonal for stability
    let mut a = xtx;
    for i in 0..p {
        a[i][i] += 1e-8;
    }

    // Cholesky decomposition A = L L^T
    let mut l = mat_zeros(p, p);
    for i in 0..p {
        for j in 0..=i {
            let mut s = 0.0;
            for k in 0..j {
                s += l[i][k] * l[j][k];
            }
            if i == j {
                let diag = a[i][i] - s;
                l[i][j] = if diag > 0.0 { diag.sqrt() } else { 1e-10 };
            } else {
                l[i][j] = if l[j][j].abs() > 1e-20 {
                    (a[i][j] - s) / l[j][j]
                } else {
                    0.0
                };
            }
        }
    }

    // Forward substitution: L z = xty
    let mut z = vec![0.0; p];
    for i in 0..p {
        let mut s = 0.0;
        for j in 0..i {
            s += l[i][j] * z[j];
        }
        z[i] = if l[i][i].abs() > 1e-20 {
            (xty[i] - s) / l[i][i]
        } else {
            0.0
        };
    }

    // Back substitution: L^T beta = z
    let mut beta = vec![0.0; p];
    for i in (0..p).rev() {
        let mut s = 0.0;
        for j in (i + 1)..p {
            s += l[j][i] * beta[j];
        }
        beta[i] = if l[i][i].abs() > 1e-20 {
            finite_or_zero((z[i] - s) / l[i][i])
        } else {
            0.0
        };
    }

    beta
}

// ===========================================================================
// Eigenvalue decomposition (power iteration with deflation)
// ===========================================================================

/// Compute eigenvalues and eigenvectors of a symmetric matrix via power iteration.
/// Returns (eigenvalues, eigenvectors) sorted descending by absolute eigenvalue.
fn eigen_symmetric(a: &Mat) -> (Vec<f64>, Mat) {
    let n = a.len();
    let max_iter = 300;
    let tol = 1e-10;

    let mut eigenvalues = Vec::with_capacity(n);
    let mut eigenvectors = Vec::with_capacity(n);
    let mut deflated = a.clone();

    for _ in 0..n {
        // Power iteration to find dominant eigenvalue/eigenvector
        let mut v = vec![1.0 / (n as f64).sqrt(); n];
        let mut eigenvalue = 0.0;

        for _ in 0..max_iter {
            let av = mat_vec_mul(&deflated, &v);
            let norm = vec_norm(&av);
            if norm < 1e-20 {
                break;
            }
            let new_v = vec_scale(&av, 1.0 / norm);
            eigenvalue = vec_dot(&new_v, &mat_vec_mul(&deflated, &new_v));

            let diff: f64 = new_v
                .iter()
                .zip(v.iter())
                .map(|(a, b)| (a - b).abs())
                .sum();
            v = new_v;
            if diff < tol {
                break;
            }
        }

        eigenvalues.push(finite_or_zero(eigenvalue));
        eigenvectors.push(v.clone());

        // Deflate: A <- A - eigenvalue * v * v^T
        for i in 0..n {
            for j in 0..n {
                deflated[i][j] -= eigenvalue * v[i] * v[j];
                deflated[i][j] = finite_or_zero(deflated[i][j]);
            }
        }
    }

    // Sort by descending absolute eigenvalue
    let mut indices: Vec<usize> = (0..n).collect();
    indices.sort_by(|&a, &b| {
        eigenvalues[b]
            .abs()
            .partial_cmp(&eigenvalues[a].abs())
            .unwrap_or(std::cmp::Ordering::Equal)
    });

    let sorted_eig: Vec<f64> = indices.iter().map(|&i| eigenvalues[i]).collect();
    let sorted_vec: Mat = indices.iter().map(|&i| eigenvectors[i].clone()).collect();

    (sorted_eig, sorted_vec)
}

/// Project matrix to positive semi-definite by clipping negative eigenvalues.
fn project_psd(a: &Mat) -> (Mat, Vec<f64>, bool) {
    let n = a.len();
    let (eigenvalues, eigenvectors) = eigen_symmetric(a);

    let was_psd = eigenvalues.iter().all(|&e| e >= -1e-10);

    // Clip negative eigenvalues to zero
    let clipped: Vec<f64> = eigenvalues.iter().map(|&e| e.max(0.0)).collect();

    // Reconstruct: A = V * diag(lambda_clipped) * V^T
    let mut result = mat_zeros(n, n);
    for k in 0..n {
        if clipped[k] < 1e-20 {
            continue;
        }
        let v = &eigenvectors[k];
        for i in 0..n {
            for j in 0..n {
                result[i][j] += clipped[k] * v[i] * v[j];
            }
        }
    }

    // Ensure finite
    for i in 0..n {
        for j in 0..n {
            result[i][j] = finite_or_zero(result[i][j]);
        }
    }

    (result, clipped, was_psd)
}

// ===========================================================================
// Cross-impact estimation
// ===========================================================================

/// Estimate the N x N cross-impact matrix from returns and signed volumes.
///
/// For each market i, regress returns_i on all signed_volumes:
///   r_i = beta_{i,1} * sv_1 + ... + beta_{i,N} * sv_N + epsilon
///
/// The matrix B = [beta_{ij}] is then projected to PSD.
fn estimate_cross_impact_impl(
    returns_matrix: &[Vec<f64>],
    volume_matrix: &[Vec<f64>],
    names: Vec<String>,
) -> PyResult<CrossImpactMatrix> {
    let n_markets = returns_matrix.len();
    let n_obs = returns_matrix[0].len();

    // Build design matrix X (n_obs x n_markets) from signed volumes
    let mut x = mat_zeros(n_obs, n_markets);
    for j in 0..n_markets {
        for t in 0..n_obs {
            x[t][j] = volume_matrix[j][t];
        }
    }

    // For each market, do OLS: returns_i ~ signed_volumes
    let mut beta_matrix = mat_zeros(n_markets, n_markets);
    for i in 0..n_markets {
        let y: Vec<f64> = returns_matrix[i].clone();
        let beta = ols_solve(&x, &y);
        for j in 0..n_markets {
            beta_matrix[i][j] = if j < beta.len() { beta[j] } else { 0.0 };
        }
    }

    // Symmetrize: B_sym = (B + B^T) / 2
    let bt = mat_transpose(&beta_matrix);
    let mut sym = mat_zeros(n_markets, n_markets);
    for i in 0..n_markets {
        for j in 0..n_markets {
            sym[i][j] = finite_or_zero(0.5 * (beta_matrix[i][j] + bt[i][j]));
        }
    }

    // Project to PSD
    let (psd_matrix, eigenvalues, was_psd) = project_psd(&sym);

    Ok(CrossImpactMatrix {
        matrix: psd_matrix,
        market_names: names,
        eigenvalues,
        is_psd: was_psd,
    })
}

// ===========================================================================
// CrossImpactEstimator (streaming with Mutex)
// ===========================================================================

struct EstimatorInner {
    n_markets: usize,
    lookback: usize,
    returns_buffer: Vec<Vec<f64>>,  // n_markets x varying
    volumes_buffer: Vec<Vec<f64>>,  // n_markets x varying
}

impl EstimatorInner {
    fn new(n_markets: usize, lookback: usize) -> Self {
        Self {
            n_markets,
            lookback,
            returns_buffer: vec![Vec::new(); n_markets],
            volumes_buffer: vec![Vec::new(); n_markets],
        }
    }

    fn add_observation(&mut self, returns: &[f64], signed_volumes: &[f64]) {
        for i in 0..self.n_markets {
            self.returns_buffer[i].push(finite_or_zero(returns[i]));
            self.volumes_buffer[i].push(finite_or_zero(signed_volumes[i]));

            // Enforce lookback window
            if self.returns_buffer[i].len() > self.lookback {
                self.returns_buffer[i].remove(0);
                self.volumes_buffer[i].remove(0);
            }
        }
    }

    fn n_observations(&self) -> usize {
        if self.returns_buffer.is_empty() {
            0
        } else {
            self.returns_buffer[0].len()
        }
    }
}

/// Streaming cross-impact estimator with rolling window.
#[pyclass]
pub struct CrossImpactEstimator {
    inner: Mutex<EstimatorInner>,
}

#[pymethods]
impl CrossImpactEstimator {
    /// Create a new estimator for N markets with given lookback window.
    #[new]
    #[pyo3(signature = (n_markets, lookback))]
    fn new(n_markets: usize, lookback: usize) -> PyResult<Self> {
        crate::auth::require_ultra()
            .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

        if n_markets < 2 {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "need at least 2 markets",
            ));
        }
        if lookback < 10 {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "lookback must be at least 10",
            ));
        }

        Ok(Self {
            inner: Mutex::new(EstimatorInner::new(n_markets, lookback)),
        })
    }

    /// Add a single observation (returns and signed volumes for all markets).
    #[pyo3(signature = (returns, signed_volumes))]
    fn add_observation(&self, returns: Vec<f64>, signed_volumes: Vec<f64>) -> PyResult<()> {
        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        if returns.len() != inner.n_markets {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "expected {} returns, got {}", inner.n_markets, returns.len()
            )));
        }
        if signed_volumes.len() != inner.n_markets {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "expected {} volumes, got {}", inner.n_markets, signed_volumes.len()
            )));
        }
        inner.add_observation(&returns, &signed_volumes);
        Ok(())
    }

    /// Estimate the cross-impact matrix from accumulated observations.
    fn estimate(&self) -> PyResult<CrossImpactMatrix> {
        let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        let n_obs = inner.n_observations();
        if n_obs < inner.n_markets + 1 {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "need at least {} observations (got {})", inner.n_markets + 1, n_obs
            )));
        }

        let names: Vec<String> = (0..inner.n_markets).map(|i| format!("market_{}", i)).collect();
        estimate_cross_impact_impl(
            &inner.returns_buffer,
            &inner.volumes_buffer,
            names,
        )
    }
}

// ===========================================================================
// Standalone estimation function
// ===========================================================================

/// Estimate the cross-impact matrix from return and volume matrices.
///
/// Parameters
/// ----------
/// returns_matrix : list[list[float]]
///     Each inner list is one market's return series.
/// volume_matrix : list[list[float]]
///     Each inner list is one market's signed volume series.
/// names : list[str]
///     Market names.
///
/// Returns
/// -------
/// CrossImpactMatrix
#[pyfunction]
#[pyo3(signature = (returns_matrix, volume_matrix, names))]
fn estimate_cross_impact(
    returns_matrix: Vec<Vec<f64>>,
    volume_matrix: Vec<Vec<f64>>,
    names: Vec<String>,
) -> PyResult<CrossImpactMatrix> {
    crate::auth::require_ultra()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

    if returns_matrix.len() < 2 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "need at least 2 markets",
        ));
    }
    if returns_matrix.len() != volume_matrix.len() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "returns_matrix and volume_matrix must have same number of markets",
        ));
    }
    if returns_matrix.len() != names.len() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "names must match number of markets",
        ));
    }
    let n_obs = returns_matrix[0].len();
    if n_obs < returns_matrix.len() + 1 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "need more observations than number of markets",
        ));
    }
    for (i, (r, v)) in returns_matrix.iter().zip(volume_matrix.iter()).enumerate() {
        if r.len() != n_obs || v.len() != n_obs {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "inconsistent length at market {}", i
            )));
        }
        for j in 0..n_obs {
            if !r[j].is_finite() || !v[j].is_finite() {
                return Err(pyo3::exceptions::PyValueError::new_err(format!(
                    "non-finite value at market {}, index {}", i, j
                )));
            }
        }
    }

    estimate_cross_impact_impl(&returns_matrix, &volume_matrix, names)
}

// ===========================================================================
// Module registration
// ===========================================================================

pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<CrossImpactMatrix>()?;
    m.add_class::<CrossImpactEstimator>()?;
    m.add_function(wrap_pyfunction!(estimate_cross_impact, m)?)?;
    Ok(())
}

// ===========================================================================
// Tests
// ===========================================================================

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

    #[test]
    fn test_ols_solve_simple() {
        // y = 2*x, data: (1,2), (2,4), (3,6)
        let x = vec![vec![1.0], vec![2.0], vec![3.0]];
        let y = vec![2.0, 4.0, 6.0];
        let beta = ols_solve(&x, &y);
        assert!((beta[0] - 2.0).abs() < 1e-4);
    }

    #[test]
    fn test_eigen_symmetric_identity() {
        let m = mat_eye(3);
        let (eigs, _vecs) = eigen_symmetric(&m);
        for &e in &eigs {
            assert!((e - 1.0).abs() < 0.1);
        }
    }

    #[test]
    fn test_project_psd() {
        // Already PSD matrix
        let m = vec![
            vec![2.0, 1.0],
            vec![1.0, 2.0],
        ];
        let (psd, eigs, was_psd) = project_psd(&m);
        assert!(was_psd);
        assert!(eigs.iter().all(|&e| e >= -1e-10));
        // Check symmetry
        assert!((psd[0][1] - psd[1][0]).abs() < 1e-6);
    }

    #[test]
    fn test_finite_or_zero_basics() {
        assert_eq!(finite_or_zero(1.0), 1.0);
        assert_eq!(finite_or_zero(f64::NAN), 0.0);
        assert_eq!(finite_or_zero(f64::INFINITY), 0.0);
    }
}
