//! Parametric statistical tests
//!
//! Provides classical parametric tests:
//! - One-sample t-test
//! - Paired t-test
//! - Independent two-sample t-test (Welch's)
//! - One-way ANOVA (F-test)

use crate::{Result, StatsError, utils};
use statrs::distribution::{ContinuousCDF, FisherSnedecor, StudentsT};

/// Result of a t-test
#[derive(Debug, Clone)]
pub struct TTestResult {
    /// The t-statistic
    pub statistic: f64,
    /// The p-value (two-tailed)
    pub pvalue: f64,
    /// Degrees of freedom
    pub df: f64,
    /// Mean of first sample (or difference for paired)
    pub mean: f64,
    /// Standard error
    pub se: f64,
    /// 95% confidence interval for the mean
    pub ci_95: (f64, f64),
}

/// Result of an F-test (ANOVA)
#[derive(Debug, Clone)]
pub struct FTestResult {
    /// The F-statistic
    pub statistic: f64,
    /// The p-value
    pub pvalue: f64,
    /// Degrees of freedom (between groups, within groups)
    pub df: (f64, f64),
    /// Sum of squares between groups
    pub ss_between: f64,
    /// Sum of squares within groups
    pub ss_within: f64,
    /// Mean square between groups
    pub ms_between: f64,
    /// Mean square within groups
    pub ms_within: f64,
    /// Eta-squared effect size
    pub eta_squared: f64,
}

/// One-sample t-test
///
/// Tests whether the mean of a sample differs from a hypothesized value.
///
/// # Arguments
/// * `data` - Sample data
/// * `popmean` - Hypothesized population mean
///
/// # Returns
/// TTestResult with t-statistic, p-value, and confidence interval
pub fn ttest_1samp(data: &[f64], popmean: f64) -> Result<TTestResult> {
    if data.len() < 2 {
        return Err(StatsError::InsufficientData {
            needed: 2,
            got: data.len(),
        });
    }

    let n = data.len() as f64;
    let mean = utils::mean(data);
    let se = utils::sem(data);
    let df = n - 1.0;

    if se < 1e-10 {
        // All values are identical
        let statistic = if (mean - popmean).abs() < 1e-10 {
            0.0
        } else {
            (mean - popmean).signum() * f64::INFINITY
        };
        return Ok(TTestResult {
            statistic,
            pvalue: if statistic.is_infinite() { 0.0 } else { 1.0 },
            df,
            mean,
            se,
            ci_95: (mean, mean),
        });
    }

    let t_stat = (mean - popmean) / se;

    // Compute p-value using Student's t distribution
    let t_dist = StudentsT::new(0.0, 1.0, df).map_err(|e| {
        StatsError::ComputationError(format!("Failed to create t-distribution: {}", e))
    })?;

    let pvalue = 2.0 * (1.0 - t_dist.cdf(t_stat.abs()));

    // 95% CI
    let t_crit = t_dist.inverse_cdf(0.975).max(0.0).min(100.0); // Fallback to reasonable bounds
    let ci_95 = (mean - t_crit * se, mean + t_crit * se);

    Ok(TTestResult {
        statistic: t_stat,
        pvalue,
        df,
        mean,
        se,
        ci_95,
    })
}

/// Paired-sample t-test
///
/// Tests whether the mean difference between paired samples differs from zero.
///
/// # Arguments
/// * `a` - First sample
/// * `b` - Second sample (paired with a)
///
/// # Returns
/// TTestResult for the paired differences
pub fn ttest_rel(a: &[f64], b: &[f64]) -> Result<TTestResult> {
    if a.len() != b.len() {
        return Err(StatsError::DimensionMismatch(format!(
            "Arrays must have same length: {} vs {}",
            a.len(),
            b.len()
        )));
    }

    // Compute differences and use one-sample t-test
    let diff: Vec<f64> = a.iter().zip(b.iter()).map(|(x, y)| x - y).collect();
    ttest_1samp(&diff, 0.0)
}

/// Independent two-sample t-test (Welch's t-test)
///
/// Tests whether two independent samples have different means.
/// Uses Welch's approximation which doesn't assume equal variances.
///
/// # Arguments
/// * `a` - First sample
/// * `b` - Second sample
///
/// # Returns
/// TTestResult for the difference in means
pub fn ttest_ind(a: &[f64], b: &[f64]) -> Result<TTestResult> {
    if a.len() < 2 {
        return Err(StatsError::InsufficientData {
            needed: 2,
            got: a.len(),
        });
    }
    if b.len() < 2 {
        return Err(StatsError::InsufficientData {
            needed: 2,
            got: b.len(),
        });
    }

    let n_a = a.len() as f64;
    let n_b = b.len() as f64;
    let mean_a = utils::mean(a);
    let mean_b = utils::mean(b);
    let var_a = utils::variance(a, 1);
    let var_b = utils::variance(b, 1);

    let se = (var_a / n_a + var_b / n_b).sqrt();
    let mean_diff = mean_a - mean_b;

    if se < 1e-10 {
        let statistic = if mean_diff.abs() < 1e-10 {
            0.0
        } else {
            mean_diff.signum() * f64::INFINITY
        };
        return Ok(TTestResult {
            statistic,
            pvalue: if statistic.is_infinite() { 0.0 } else { 1.0 },
            df: n_a + n_b - 2.0,
            mean: mean_diff,
            se,
            ci_95: (mean_diff, mean_diff),
        });
    }

    let t_stat = mean_diff / se;

    // Welch-Satterthwaite degrees of freedom
    let v_a = var_a / n_a;
    let v_b = var_b / n_b;
    let df = (v_a + v_b).powi(2) / (v_a.powi(2) / (n_a - 1.0) + v_b.powi(2) / (n_b - 1.0));

    // Compute p-value
    let t_dist = StudentsT::new(0.0, 1.0, df).map_err(|e| {
        StatsError::ComputationError(format!("Failed to create t-distribution: {}", e))
    })?;

    let pvalue = 2.0 * (1.0 - t_dist.cdf(t_stat.abs()));

    // 95% CI
    let t_crit = t_dist.inverse_cdf(0.975).max(0.0).min(100.0);
    let ci_95 = (mean_diff - t_crit * se, mean_diff + t_crit * se);

    Ok(TTestResult {
        statistic: t_stat,
        pvalue,
        df,
        mean: mean_diff,
        se,
        ci_95,
    })
}

/// One-way ANOVA (F-test)
///
/// Tests whether the means of multiple groups differ.
///
/// # Arguments
/// * `groups` - Vector of groups, each group is a vector of observations
///
/// # Returns
/// FTestResult with F-statistic, p-value, and effect size
pub fn f_oneway(groups: &[&[f64]]) -> Result<FTestResult> {
    if groups.len() < 2 {
        return Err(StatsError::InvalidInput(
            "Need at least 2 groups for ANOVA".to_string(),
        ));
    }

    for group in groups {
        if group.len() < 2 {
            return Err(StatsError::InsufficientData {
                needed: 2,
                got: group.len(),
            });
        }
    }

    let k = groups.len() as f64; // Number of groups
    let n: f64 = groups.iter().map(|g| g.len() as f64).sum(); // Total observations

    // Grand mean
    let all_values: Vec<f64> = groups.iter().flat_map(|g| g.iter().copied()).collect();
    let grand_mean = utils::mean(&all_values);

    // Group means
    let group_means: Vec<f64> = groups.iter().map(|g| utils::mean(g)).collect();

    // Sum of squares between groups
    let ss_between: f64 = groups
        .iter()
        .zip(group_means.iter())
        .map(|(g, &gm)| g.len() as f64 * (gm - grand_mean).powi(2))
        .sum();

    // Sum of squares within groups
    let ss_within: f64 = groups
        .iter()
        .zip(group_means.iter())
        .map(|(g, &gm)| g.iter().map(|&x| (x - gm).powi(2)).sum::<f64>())
        .sum();

    // Degrees of freedom
    let df_between = k - 1.0;
    let df_within = n - k;

    // Mean squares
    let ms_between = ss_between / df_between;
    let ms_within = ss_within / df_within;

    // F-statistic
    let f_stat = if ms_within < 1e-10 {
        if ms_between < 1e-10 {
            0.0
        } else {
            f64::INFINITY
        }
    } else {
        ms_between / ms_within
    };

    // P-value from F-distribution
    let f_dist = FisherSnedecor::new(df_between, df_within).map_err(|e| {
        StatsError::ComputationError(format!("Failed to create F-distribution: {}", e))
    })?;

    let pvalue = if f_stat.is_infinite() {
        0.0
    } else {
        1.0 - f_dist.cdf(f_stat)
    };

    // Effect size (eta-squared)
    let ss_total = ss_between + ss_within;
    let eta_squared = if ss_total < 1e-10 {
        0.0
    } else {
        ss_between / ss_total
    };

    Ok(FTestResult {
        statistic: f_stat,
        pvalue,
        df: (df_between, df_within),
        ss_between,
        ss_within,
        ms_between,
        ms_within,
        eta_squared,
    })
}

/// Vectorized t-test for multiple features
///
/// Performs element-wise t-tests on multi-dimensional data.
///
/// # Arguments
/// * `data` - 2D array [n_observations x n_features]
/// * `popmean` - Hypothesized population mean
///
/// # Returns
/// Tuple of (t_statistics, p_values) for each feature
pub fn ttest_1samp_vectorized(data: &[Vec<f64>], popmean: f64) -> Result<(Vec<f64>, Vec<f64>)> {
    if data.is_empty() {
        return Err(StatsError::InvalidInput("Empty data".to_string()));
    }

    let n_obs = data.len();
    let n_features = data[0].len();

    if n_obs < 2 {
        return Err(StatsError::InsufficientData {
            needed: 2,
            got: n_obs,
        });
    }

    let df = (n_obs - 1) as f64;
    let t_dist = StudentsT::new(0.0, 1.0, df).map_err(|e| {
        StatsError::ComputationError(format!("Failed to create t-distribution: {}", e))
    })?;

    let mut t_stats = Vec::with_capacity(n_features);
    let mut pvalues = Vec::with_capacity(n_features);

    for j in 0..n_features {
        let col: Vec<f64> = data.iter().map(|row| row[j] - popmean).collect();
        let mean = utils::mean(&col);
        let se = utils::sem(&col);

        let t_stat = if se < 1e-10 { 0.0 } else { mean / se };
        let pvalue = 2.0 * (1.0 - t_dist.cdf(t_stat.abs()));

        t_stats.push(t_stat);
        pvalues.push(pvalue);
    }

    Ok((t_stats, pvalues))
}

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

    #[test]
    fn test_ttest_1samp_significant() {
        let data = vec![2.1, 2.3, 1.9, 2.5, 2.2, 2.0, 2.4, 2.1];
        let result = ttest_1samp(&data, 0.0).unwrap();

        assert!(result.statistic > 0.0);
        assert!(result.pvalue < 0.001);
        assert!(result.ci_95.0 > 0.0);
    }

    #[test]
    fn test_ttest_1samp_not_significant() {
        let data = vec![0.1, -0.2, 0.15, -0.1, 0.05, -0.05, 0.1, -0.15];
        let result = ttest_1samp(&data, 0.0).unwrap();

        assert!(result.pvalue > 0.05);
        assert!(result.ci_95.0 < 0.0 && result.ci_95.1 > 0.0);
    }

    #[test]
    fn test_ttest_rel() {
        let a = vec![10.0, 11.0, 12.0, 13.0, 14.0, 15.0];
        let b = vec![8.0, 9.0, 10.0, 11.0, 12.0, 13.0];
        let result = ttest_rel(&a, &b).unwrap();

        assert!(result.statistic > 0.0);
        assert!(result.pvalue < 0.001);
        assert!((result.mean - 2.0).abs() < 1e-10);
    }

    #[test]
    fn test_ttest_ind() {
        let a = vec![10.0, 11.0, 12.0, 10.5, 11.5, 10.8];
        let b = vec![5.0, 6.0, 5.5, 6.5, 5.8, 6.2];
        let result = ttest_ind(&a, &b).unwrap();

        assert!(result.statistic > 0.0);
        assert!(result.pvalue < 0.001);
        assert!(result.ci_95.0 > 0.0);
    }

    #[test]
    fn test_f_oneway_significant() {
        let group1 = vec![10.0, 11.0, 12.0, 10.5, 11.5];
        let group2 = vec![5.0, 6.0, 5.5, 6.5, 5.8];
        let group3 = vec![15.0, 16.0, 15.5, 16.5, 15.8];

        let result = f_oneway(&[&group1, &group2, &group3]).unwrap();

        assert!(result.statistic > 0.0);
        assert!(result.pvalue < 0.001);
        assert!(result.eta_squared > 0.8); // Large effect
    }

    #[test]
    fn test_f_oneway_not_significant() {
        let group1 = vec![10.0, 11.0, 9.0, 10.5, 9.5];
        let group2 = vec![10.2, 10.8, 9.2, 10.3, 9.7];
        let group3 = vec![9.8, 11.2, 9.1, 10.7, 9.6];

        let result = f_oneway(&[&group1, &group2, &group3]).unwrap();

        assert!(result.pvalue > 0.05);
    }

    #[test]
    fn test_ttest_vectorized() {
        let data = vec![
            vec![2.0, 0.1, 3.0],
            vec![2.2, -0.1, 2.8],
            vec![1.9, 0.05, 3.1],
            vec![2.1, -0.05, 2.9],
        ];

        let (t_stats, pvalues) = ttest_1samp_vectorized(&data, 0.0).unwrap();

        assert_eq!(t_stats.len(), 3);
        assert_eq!(pvalues.len(), 3);
        // First and third features should be significant
        assert!(pvalues[0] < 0.05);
        assert!(pvalues[2] < 0.05);
        // Second feature should not be significant
        assert!(pvalues[1] > 0.05);
    }
}