//! Statistical moment calculations for portfolio analysis.
//!
//! Calculates first four moments (mean, variance, skewness, kurtosis)
//! and their cross-asset equivalents (covariance, co-skewness, co-kurtosis).

use quantumport_shared::{
    AssetReturns, CokurtosisTensor, CoskewnessTensor, CovarianceMatrix, MomentStatistics,
};

/// Moment calculator for portfolio returns.
#[derive(Debug)]
pub struct MomentCalculator {
    /// Annualization factor (252 trading days)
    annualization_factor: f64,
}

impl Default for MomentCalculator {
    fn default() -> Self {
        Self::new()
    }
}

impl MomentCalculator {
    /// Create a new moment calculator.
    #[must_use]
    pub fn new() -> Self {
        Self {
            annualization_factor: 252.0,
        }
    }

    /// Calculate individual asset statistics.
    #[must_use]
    pub fn calculate_asset_stats(&self, returns: &[f64]) -> MomentStatistics {
        let n = returns.len() as f64;
        if n == 0.0 {
            return MomentStatistics {
                mean: 0.0,
                variance: 0.0,
                skewness: 0.0,
                kurtosis: 3.0,
            };
        }

        // Mean (annualized)
        let mean: f64 = returns.iter().sum::<f64>() / n * self.annualization_factor;

        // Variance (annualized)
        let daily_mean = returns.iter().sum::<f64>() / n;
        let variance: f64 = returns
            .iter()
            .map(|r| (r - daily_mean).powi(2))
            .sum::<f64>()
            / (n - 1.0)
            * self.annualization_factor;

        // Standard deviation for normalization
        let std = (returns
            .iter()
            .map(|r| (r - daily_mean).powi(2))
            .sum::<f64>()
            / (n - 1.0))
            .sqrt();

        // Skewness
        let skewness = if std > 0.0 {
            let m3: f64 = returns
                .iter()
                .map(|r| ((r - daily_mean) / std).powi(3))
                .sum::<f64>();
            m3 * n / ((n - 1.0) * (n - 2.0))
        } else {
            0.0
        };

        // Kurtosis (excess)
        let kurtosis = if std > 0.0 {
            let m4: f64 = returns
                .iter()
                .map(|r| ((r - daily_mean) / std).powi(4))
                .sum::<f64>();
            (m4 / n) - 3.0
        } else {
            0.0
        };

        MomentStatistics {
            mean,
            variance,
            skewness,
            kurtosis,
        }
    }

    /// Calculate covariance matrix for multiple assets.
    #[must_use]
    pub fn calculate_covariance(&self, returns: &[AssetReturns]) -> CovarianceMatrix {
        let n = returns.len();
        let symbols: Vec<String> = returns.iter().map(|r| r.symbol.clone()).collect();

        // Calculate means
        let means: Vec<f64> = returns
            .iter()
            .map(|r| r.returns.iter().sum::<f64>() / r.returns.len() as f64)
            .collect();

        // Calculate covariance matrix
        let mut data = vec![0.0; n * n];
        let num_obs = returns[0].returns.len() as f64;

        for i in 0..n {
            for j in 0..n {
                let mut cov = 0.0;
                for t in 0..returns[i].returns.len() {
                    let dev_i = returns[i].returns[t] - means[i];
                    let dev_j = returns[j].returns[t] - means[j];
                    cov += dev_i * dev_j;
                }
                // Annualize
                data[i * n + j] = cov / (num_obs - 1.0) * self.annualization_factor;
            }
        }

        CovarianceMatrix {
            symbols,
            data,
            dimension: n,
        }
    }

    /// Calculate co-skewness tensor.
    #[must_use]
    pub fn calculate_coskewness(&self, returns: &[AssetReturns]) -> CoskewnessTensor {
        let n = returns.len();
        let symbols: Vec<String> = returns.iter().map(|r| r.symbol.clone()).collect();

        // Calculate means and standard deviations
        let stats: Vec<(f64, f64)> = returns
            .iter()
            .map(|r| {
                let mean = r.returns.iter().sum::<f64>() / r.returns.len() as f64;
                let var = r.returns.iter().map(|x| (x - mean).powi(2)).sum::<f64>()
                    / (r.returns.len() - 1) as f64;
                (mean, var.sqrt())
            })
            .collect();

        // Calculate co-skewness tensor
        let mut data = vec![0.0; n * n * n];
        let num_obs = returns[0].returns.len() as f64;

        for i in 0..n {
            for j in 0..n {
                for k in 0..n {
                    let mut coskew = 0.0;
                    for t in 0..returns[i].returns.len() {
                        let z_i = if stats[i].1 > 0.0 {
                            (returns[i].returns[t] - stats[i].0) / stats[i].1
                        } else {
                            0.0
                        };
                        let z_j = if stats[j].1 > 0.0 {
                            (returns[j].returns[t] - stats[j].0) / stats[j].1
                        } else {
                            0.0
                        };
                        let z_k = if stats[k].1 > 0.0 {
                            (returns[k].returns[t] - stats[k].0) / stats[k].1
                        } else {
                            0.0
                        };
                        coskew += z_i * z_j * z_k;
                    }
                    data[i * n * n + j * n + k] = coskew / num_obs;
                }
            }
        }

        CoskewnessTensor {
            symbols,
            data,
            dimension: n,
        }
    }

    /// Calculate co-kurtosis tensor.
    #[must_use]
    pub fn calculate_cokurtosis(&self, returns: &[AssetReturns]) -> CokurtosisTensor {
        let n = returns.len();
        let symbols: Vec<String> = returns.iter().map(|r| r.symbol.clone()).collect();

        // Calculate means and standard deviations
        let stats: Vec<(f64, f64)> = returns
            .iter()
            .map(|r| {
                let mean = r.returns.iter().sum::<f64>() / r.returns.len() as f64;
                let var = r.returns.iter().map(|x| (x - mean).powi(2)).sum::<f64>()
                    / (r.returns.len() - 1) as f64;
                (mean, var.sqrt())
            })
            .collect();

        // Calculate co-kurtosis tensor (simplified - diagonal elements only for demo)
        let mut data = vec![0.0; n * n * n * n];
        let num_obs = returns[0].returns.len() as f64;

        // Calculate diagonal elements (i,i,i,i)
        for i in 0..n {
            let mut cokurt = 0.0;
            for t in 0..returns[i].returns.len() {
                let z_i = if stats[i].1 > 0.0 {
                    (returns[i].returns[t] - stats[i].0) / stats[i].1
                } else {
                    0.0
                };
                cokurt += z_i.powi(4);
            }
            data[i * n * n * n + i * n * n + i * n + i] = cokurt / num_obs - 3.0;
        }

        CokurtosisTensor {
            symbols,
            data,
            dimension: n,
        }
    }

    /// Calculate portfolio moments given weights.
    #[must_use]
    pub fn calculate_portfolio_moments(
        &self,
        weights: &[f64],
        returns: &[AssetReturns],
    ) -> MomentStatistics {
        let n = returns[0].returns.len();

        // Calculate portfolio returns
        let mut portfolio_returns = vec![0.0; n];
        for (i, ret) in returns.iter().enumerate() {
            for (t, r) in ret.returns.iter().enumerate() {
                portfolio_returns[t] += weights[i] * r;
            }
        }

        self.calculate_asset_stats(&portfolio_returns)
    }
}

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

    fn create_test_returns() -> Vec<AssetReturns> {
        vec![
            AssetReturns {
                symbol: "A".to_string(),
                returns: vec![
                    0.01, -0.02, 0.015, 0.005, -0.01, 0.02, -0.005, 0.01, 0.008, -0.012,
                ],
                start_date: "2024-01-01".to_string(),
                end_date: "2024-01-10".to_string(),
            },
            AssetReturns {
                symbol: "B".to_string(),
                returns: vec![
                    0.002, 0.001, 0.003, -0.001, 0.002, 0.001, 0.002, -0.001, 0.003, 0.001,
                ],
                start_date: "2024-01-01".to_string(),
                end_date: "2024-01-10".to_string(),
            },
        ]
    }

    #[test]
    fn test_calculator_creation() {
        let calc = MomentCalculator::new();
        assert!((calc.annualization_factor - 252.0).abs() < 1e-10);
    }

    #[test]
    fn test_asset_stats() {
        let calc = MomentCalculator::new();
        let returns = vec![0.01, -0.02, 0.015, 0.005, -0.01];

        let stats = calc.calculate_asset_stats(&returns);
        assert!(stats.mean.abs() < 1.0); // Reasonable annualized return
        assert!(stats.variance > 0.0);
    }

    #[test]
    fn test_covariance_matrix() {
        let calc = MomentCalculator::new();
        let returns = create_test_returns();

        let cov = calc.calculate_covariance(&returns);
        assert_eq!(cov.dimension, 2);
        assert_eq!(cov.data.len(), 4);

        // Variance should be positive
        assert!(cov.get(0, 0).unwrap() > 0.0);
        assert!(cov.get(1, 1).unwrap() > 0.0);

        // Covariance should be symmetric
        assert!((cov.get(0, 1).unwrap() - cov.get(1, 0).unwrap()).abs() < 1e-10);
    }

    #[test]
    fn test_coskewness() {
        let calc = MomentCalculator::new();
        let returns = create_test_returns();

        let coskew = calc.calculate_coskewness(&returns);
        assert_eq!(coskew.dimension, 2);
        assert!(!coskew.data.is_empty());
    }

    #[test]
    fn test_cokurtosis() {
        let calc = MomentCalculator::new();
        let returns = create_test_returns();

        let cokurt = calc.calculate_cokurtosis(&returns);
        assert_eq!(cokurt.dimension, 2);
        assert!(!cokurt.data.is_empty());
    }

    #[test]
    fn test_portfolio_moments() {
        let calc = MomentCalculator::new();
        let returns = create_test_returns();
        let weights = vec![0.6, 0.4];

        let stats = calc.calculate_portfolio_moments(&weights, &returns);
        assert!(stats.variance > 0.0);
    }

    #[test]
    fn test_empty_returns() {
        let calc = MomentCalculator::new();
        let stats = calc.calculate_asset_stats(&[]);

        assert!((stats.mean).abs() < 1e-10);
        assert!((stats.variance).abs() < 1e-10);
    }
}