//! Statistical analysis module

use pyo3::prelude::*;
use numpy::{PyArray1, PyReadonlyArray1, IntoPyArray};
use rtx_neuro_stats::{
    ttest_1samp, ttest_ind, ttest_rel,
    fdr_correction, bonferroni_correction,
    permutation_test_ind, Tail,
};

/// Create the stats submodule
pub fn create_module(py: Python<'_>) -> PyResult<Bound<'_, PyModule>> {
    let m = PyModule::new(py, "stats")?;

    // Basic tests
    m.add_function(wrap_pyfunction!(py_ttest_1samp, &m)?)?;
    m.add_function(wrap_pyfunction!(py_ttest_ind, &m)?)?;
    m.add_function(wrap_pyfunction!(py_ttest_rel, &m)?)?;

    // Permutation tests
    m.add_function(wrap_pyfunction!(py_permutation_t_test, &m)?)?;

    // Multiple comparison correction
    m.add_function(wrap_pyfunction!(py_fdr_correction, &m)?)?;
    m.add_function(wrap_pyfunction!(py_bonferroni_correction, &m)?)?;

    Ok(m)
}

fn stats_err_to_py(e: rtx_neuro_stats::StatsError) -> PyErr {
    PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string())
}

/// One-sample t-test
///
/// # Arguments
/// * `data` - Sample data [n_samples]
/// * `popmean` - Expected population mean (default: 0.0)
///
/// # Returns
/// Tuple of (t-statistic, p-value)
#[pyfunction]
#[pyo3(name = "ttest_1samp", signature = (data, popmean=0.0))]
fn py_ttest_1samp(
    data: PyReadonlyArray1<f64>,
    popmean: f64,
) -> PyResult<(f64, f64)> {
    let input = data.as_slice()?.to_vec();

    let result = ttest_1samp(&input, popmean)
        .map_err(stats_err_to_py)?;

    Ok((result.statistic, result.pvalue))
}

/// Independent samples t-test (Welch's)
///
/// # Arguments
/// * `a` - First sample [n_samples_a]
/// * `b` - Second sample [n_samples_b]
///
/// # Returns
/// Tuple of (t-statistic, p-value)
#[pyfunction]
#[pyo3(name = "ttest_ind")]
fn py_ttest_ind(
    a: PyReadonlyArray1<f64>,
    b: PyReadonlyArray1<f64>,
) -> PyResult<(f64, f64)> {
    let a_vec = a.as_slice()?.to_vec();
    let b_vec = b.as_slice()?.to_vec();

    let result = ttest_ind(&a_vec, &b_vec)
        .map_err(stats_err_to_py)?;

    Ok((result.statistic, result.pvalue))
}

/// Paired samples t-test
///
/// # Arguments
/// * `a` - First sample [n_samples]
/// * `b` - Second sample [n_samples]
///
/// # Returns
/// Tuple of (t-statistic, p-value)
#[pyfunction]
#[pyo3(name = "ttest_rel")]
fn py_ttest_rel(
    a: PyReadonlyArray1<f64>,
    b: PyReadonlyArray1<f64>,
) -> PyResult<(f64, f64)> {
    let a_vec = a.as_slice()?.to_vec();
    let b_vec = b.as_slice()?.to_vec();

    let result = ttest_rel(&a_vec, &b_vec)
        .map_err(stats_err_to_py)?;

    Ok((result.statistic, result.pvalue))
}

/// Permutation t-test
///
/// Non-parametric test using permutation of labels.
///
/// # Arguments
/// * `a` - First group [n_samples_a]
/// * `b` - Second group [n_samples_b]
/// * `n_permutations` - Number of permutations (default: 1000)
/// * `seed` - Random seed (default: None)
///
/// # Returns
/// Tuple of (observed statistic, p-value)
#[pyfunction]
#[pyo3(name = "permutation_t_test", signature = (a, b, n_permutations=1000, seed=None))]
fn py_permutation_t_test(
    a: PyReadonlyArray1<f64>,
    b: PyReadonlyArray1<f64>,
    n_permutations: usize,
    seed: Option<u64>,
) -> PyResult<(f64, f64)> {
    let a_vec = a.as_slice()?.to_vec();
    let b_vec = b.as_slice()?.to_vec();

    let result = permutation_test_ind(&a_vec, &b_vec, n_permutations, Tail::TwoSided, seed)
        .map_err(stats_err_to_py)?;

    Ok((result.statistic, result.pvalue))
}

/// False Discovery Rate (FDR) correction
///
/// Benjamini-Hochberg procedure for controlling FDR.
///
/// # Arguments
/// * `p_values` - Array of p-values
/// * `alpha` - Significance level (default: 0.05)
///
/// # Returns
/// Tuple of (rejected, corrected_p_values)
#[pyfunction]
#[pyo3(name = "fdr_correction", signature = (p_values, alpha=0.05))]
fn py_fdr_correction<'py>(
    py: Python<'py>,
    p_values: PyReadonlyArray1<f64>,
    alpha: f64,
) -> PyResult<(Bound<'py, PyArray1<bool>>, Bound<'py, PyArray1<f64>>)> {
    let p_vec = p_values.as_slice()?.to_vec();

    let (rejected, corrected) = fdr_correction(&p_vec, alpha);

    Ok((
        rejected.into_pyarray(py),
        corrected.into_pyarray(py),
    ))
}

/// Bonferroni correction
///
/// # Arguments
/// * `p_values` - Array of p-values
/// * `alpha` - Significance level (default: 0.05)
///
/// # Returns
/// Tuple of (rejected, corrected_p_values)
#[pyfunction]
#[pyo3(name = "bonferroni_correction", signature = (p_values, alpha=0.05))]
fn py_bonferroni_correction<'py>(
    py: Python<'py>,
    p_values: PyReadonlyArray1<f64>,
    alpha: f64,
) -> PyResult<(Bound<'py, PyArray1<bool>>, Bound<'py, PyArray1<f64>>)> {
    let p_vec = p_values.as_slice()?.to_vec();

    let (rejected, corrected) = bonferroni_correction(&p_vec, alpha);

    Ok((
        rejected.into_pyarray(py),
        corrected.into_pyarray(py),
    ))
}