//! Vine Copulas for high-dimensional dependence modeling.
//!
//! Implements C-vine and D-vine copula structures for modeling joint
//! distributions of 3+ prediction market outcomes. Uses sequential
//! pair-copula estimation tree by tree:
//!
//! - **C-vine**: star structure, one root variable per tree
//! - **D-vine**: path structure, sequential conditioning
//!
//! Each pair copula is automatically selected from the 5 families
//! in the copula module (Gaussian, Student-t, Clayton, Gumbel, Frank).
//!
//! All math from scratch in Rust -- no external crate dependencies beyond PyO3.
//!
//! Ultra tier required.

use pyo3::prelude::*;

use crate::copula::{
    best_copula_fit, copula_h_function, CopulaFamily,
};

// ===========================================================================
// Constants & helpers
// ===========================================================================

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

#[inline]
fn clamp_unit(x: f64) -> f64 {
    x.clamp(1e-10, 1.0 - 1e-10)
}

// ===========================================================================
// Xoshiro256++ PRNG
// ===========================================================================

struct Xoshiro256 {
    s: [u64; 4],
}

impl Xoshiro256 {
    fn new(seed: u64) -> Self {
        let mut z = seed;
        let mut s = [0u64; 4];
        for si in &mut s {
            z = z.wrapping_add(0x9e3779b97f4a7c15);
            let mut x = z;
            x = (x ^ (x >> 30)).wrapping_mul(0xbf58476d1ce4e5b9);
            x = (x ^ (x >> 27)).wrapping_mul(0x94d049bb133111eb);
            *si = x ^ (x >> 31);
        }
        Self { s }
    }

    fn next_u64(&mut self) -> u64 {
        let result = (self.s[0].wrapping_add(self.s[3]))
            .rotate_left(23)
            .wrapping_add(self.s[0]);
        let t = self.s[1] << 17;
        self.s[2] ^= self.s[0];
        self.s[3] ^= self.s[1];
        self.s[1] ^= self.s[2];
        self.s[0] ^= self.s[3];
        self.s[2] ^= t;
        self.s[3] = self.s[3].rotate_left(45);
        result
    }

    fn next_f64(&mut self) -> f64 {
        (self.next_u64() >> 11) as f64 / (1u64 << 53) as f64
    }

}

// ===========================================================================
// Vine type enum
// ===========================================================================

/// Type of vine copula structure.
#[pyclass(eq, eq_int, frozen, hash)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum VineType {
    /// Canonical vine (star structure, one root per tree).
    CVine = 0,
    /// D-vine (path structure, sequential).
    DVine = 1,
}

#[pymethods]
impl VineType {
    fn __repr__(&self) -> String {
        match self {
            VineType::CVine => "VineType.CVine".to_string(),
            VineType::DVine => "VineType.DVine".to_string(),
        }
    }
}

// ===========================================================================
// Frozen result types
// ===========================================================================

/// A single pair copula in the vine structure.
#[pyclass(frozen)]
#[derive(Debug, Clone)]
pub struct PairCopula {
    /// Copula family for this pair.
    #[pyo3(get)]
    pub family: CopulaFamily,
    /// Estimated dependence parameter.
    #[pyo3(get)]
    pub parameter: f64,
    /// First variable index.
    #[pyo3(get)]
    pub var1: usize,
    /// Second variable index.
    #[pyo3(get)]
    pub var2: usize,
    /// Conditioning variable indices.
    #[pyo3(get)]
    pub conditioning: Vec<usize>,
}

#[pymethods]
impl PairCopula {
    fn __repr__(&self) -> String {
        format!(
            "PairCopula(family={:?}, param={:.4}, var1={}, var2={}, cond={:?})",
            self.family, self.parameter, self.var1, self.var2, self.conditioning
        )
    }
}

/// Result of fitting a vine copula model.
#[pyclass(frozen)]
#[derive(Debug, Clone)]
pub struct VineCopulaResult {
    /// Type of vine (C-vine or D-vine).
    #[pyo3(get)]
    pub vine_type: VineType,
    /// All pair copulas across all trees.
    #[pyo3(get)]
    pub pair_copulas: Vec<PairCopula>,
    /// Number of variables in the vine.
    #[pyo3(get)]
    pub n_variables: usize,
    /// Total log-likelihood of the vine.
    #[pyo3(get)]
    pub log_likelihood: f64,
}

#[pymethods]
impl VineCopulaResult {
    fn __repr__(&self) -> String {
        format!(
            "VineCopulaResult(type={:?}, n_vars={}, n_pairs={}, ll={:.4})",
            self.vine_type,
            self.n_variables,
            self.pair_copulas.len(),
            self.log_likelihood
        )
    }
}

// ===========================================================================
// Rank transform (local copy, does not need auth)
// ===========================================================================

fn rank_transform_vec(data: &[f64]) -> Vec<f64> {
    let n = data.len();
    if n == 0 {
        return vec![];
    }
    let mut indexed: Vec<(f64, usize)> = data.iter().copied().enumerate().map(|(i, v)| (v, i)).collect();
    indexed.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal));

    let mut ranks = vec![0.0_f64; n];
    let mut i = 0;
    while i < n {
        let mut j = i + 1;
        while j < n && (indexed[j].0 - indexed[i].0).abs() < 1e-15 {
            j += 1;
        }
        let avg_rank = (i + j) as f64 / 2.0;
        for item in indexed.iter().take(j).skip(i) {
            ranks[item.1] = avg_rank;
        }
        i = j;
    }

    let denom = (n + 1) as f64;
    ranks.iter().map(|&r| clamp_unit(r / denom)).collect()
}

// ===========================================================================
// Vine fitting internals
// ===========================================================================

/// Compute the h-function (conditional CDF) for a pair copula.
fn h_func(family: CopulaFamily, param: f64, u: f64, v: f64) -> f64 {
    copula_h_function(family, param, u, v)
}

/// Fit a D-vine copula sequentially.
fn fit_dvine(data: &[Vec<f64>]) -> (Vec<PairCopula>, f64) {
    let d = data.len();
    let n = data[0].len();
    let mut pair_copulas = Vec::new();
    let mut total_ll = 0.0;

    // Transform each variable to pseudo-uniforms
    let u: Vec<Vec<f64>> = data.iter().map(|col| rank_transform_vec(col)).collect();

    // Tree 1: fit bivariate copulas for adjacent pairs (j, j+1)
    // Then compute h-functions for conditioning in next tree
    let mut v_current = u.clone(); // current pseudo-observations
    let mut _v_left: Vec<Vec<f64>> = Vec::new();  // h(u_{j+1} | u_j)
    let mut _v_right: Vec<Vec<f64>> = Vec::new(); // h(u_j | u_{j+1})

    for tree in 0..(d - 1) {
        let n_pairs = d - 1 - tree;
        let mut new_v_left = Vec::new();
        let mut new_v_right = Vec::new();

        for pair_idx in 0..n_pairs {
            let u_data = &v_current[pair_idx];
            let v_data = &v_current[pair_idx + 1];

            // Fit best copula
            let fit = best_copula_fit(u_data, v_data);
            total_ll += fit.log_likelihood;

            // Build conditioning set
            let conditioning: Vec<usize> = if tree == 0 {
                vec![]
            } else {
                // For D-vine tree k, conditioning set is {pair_idx+1, ..., pair_idx+tree}
                (pair_idx + 1..pair_idx + 1 + tree).collect()
            };

            // Variable indices
            let var1 = pair_idx;
            let var2 = pair_idx + 1 + tree;

            pair_copulas.push(PairCopula {
                family: fit.family,
                parameter: fit.parameter,
                var1,
                var2,
                conditioning,
            });

            // Compute h-functions for next tree
            if tree < d - 2 {
                let mut h_left = Vec::with_capacity(n);
                let mut h_right = Vec::with_capacity(n);
                for i in 0..n {
                    // h(v|u) = C(v|u)
                    h_left.push(clamp_unit(h_func(fit.family, fit.parameter, u_data[i], v_data[i])));
                    // h(u|v) = C(u|v) -- swap
                    h_right.push(clamp_unit(h_func(fit.family, fit.parameter, v_data[i], u_data[i])));
                }
                new_v_left.push(h_left);
                new_v_right.push(h_right);
            }
        }

        // For next tree: pseudo-observations are h-functions
        // D-vine: adjacent pseudo-obs for next tree are h_right[j] and h_left[j+1]
        if tree < d - 2 {
            let next_n_vars = n_pairs - 1;
            let mut next_v = Vec::with_capacity(next_n_vars + 1);
            // First variable of next tree: h_right[0]
            if !new_v_right.is_empty() {
                next_v.push(new_v_right[0].clone());
            }
            for j in 0..next_n_vars {
                if j + 1 < new_v_left.len() {
                    next_v.push(new_v_left[j + 1].clone());
                } else if j < new_v_right.len() {
                    next_v.push(new_v_right[j].clone());
                }
            }
            // Ensure we have enough variables
            while next_v.len() < next_n_vars + 1 {
                // Pad with existing if needed
                if let Some(last) = next_v.last().cloned() {
                    next_v.push(last);
                } else {
                    break;
                }
            }
            v_current = next_v;
        }

        _v_left = new_v_left;
        _v_right = new_v_right;
    }

    (pair_copulas, total_ll)
}

/// Fit a C-vine copula sequentially.
fn fit_cvine(data: &[Vec<f64>]) -> (Vec<PairCopula>, f64) {
    let d = data.len();
    let n = data[0].len();
    let mut pair_copulas = Vec::new();
    let mut total_ll = 0.0;

    // Transform each variable to pseudo-uniforms
    let u: Vec<Vec<f64>> = data.iter().map(|col| rank_transform_vec(col)).collect();

    // C-vine: in tree k, variable k is the root, connected to all remaining variables
    let mut current_u = u.clone();

    for tree in 0..(d - 1) {
        let _root = tree;
        let n_pairs = d - 1 - tree;
        let mut new_u = Vec::with_capacity(n_pairs);

        // Find the root variable (highest Kendall's tau sum)
        let root_data = &current_u[0]; // First variable in current set is root

        for pair_idx in 0..n_pairs {
            let other_idx = pair_idx + 1;
            if other_idx >= current_u.len() {
                break;
            }
            let other_data = &current_u[other_idx];

            // Fit best copula between root and other
            let fit = best_copula_fit(root_data, other_data);
            total_ll += fit.log_likelihood;

            let conditioning: Vec<usize> = (0..tree).collect();
            let var2 = tree + pair_idx + 1;

            pair_copulas.push(PairCopula {
                family: fit.family,
                parameter: fit.parameter,
                var1: tree,
                var2,
                conditioning,
            });

            // Compute h-function for next tree: h(other | root)
            if tree < d - 2 {
                let mut h_vals = Vec::with_capacity(n);
                for i in 0..n {
                    h_vals.push(clamp_unit(h_func(
                        fit.family,
                        fit.parameter,
                        root_data[i],
                        other_data[i],
                    )));
                }
                new_u.push(h_vals);
            }
        }

        if tree < d - 2 {
            current_u = new_u;
        }
    }

    (pair_copulas, total_ll)
}

// ===========================================================================
// Vine sampling
// ===========================================================================

/// Sample from a D-vine copula.
fn sample_dvine(pair_copulas: &[PairCopula], d: usize, n: usize, rng: &mut Xoshiro256) -> Vec<Vec<f64>> {
    let mut result = vec![vec![0.0; n]; d];

    for sample_idx in 0..n {
        // Generate d independent uniforms
        let w: Vec<f64> = (0..d).map(|_| clamp_unit(rng.next_f64())).collect();

        // v[j][k] stores intermediate conditional values
        let mut v = vec![vec![0.0_f64; d]; d];
        v[0][0] = w[0];
        result[0][sample_idx] = w[0];

        for j in 1..d {
            v[j][0] = w[j];
            // Apply inverse h-functions through the vine trees
            for k in (0..j).rev() {
                // Find the pair copula for tree (j-k-1) if it exists
                let _tree = j - k - 1;
                // Search for the pair copula
                let pc = find_pair_copula_dvine(pair_copulas, d, _tree, k, j);
                if let Some(pc) = pc {
                    // Inverse h-function: find u such that h(u|v_cond) = current value
                    v[j][0] = inverse_h_func(pc.family, pc.parameter, v[k][0], v[j][0]);
                }
            }
            result[j][sample_idx] = v[j][0];

            // Update conditional values for future use
            for k in 0..j {
                let pc = find_pair_copula_dvine(pair_copulas, d, 0, k, j);
                if let Some(pc) = pc {
                    v[j][k + 1] = h_func(pc.family, pc.parameter, v[k][0], v[j][0]);
                }
            }
        }
    }

    result
}

/// Sample from a C-vine copula.
fn sample_cvine(pair_copulas: &[PairCopula], d: usize, n: usize, rng: &mut Xoshiro256) -> Vec<Vec<f64>> {
    let mut result = vec![vec![0.0; n]; d];

    for sample_idx in 0..n {
        let w: Vec<f64> = (0..d).map(|_| clamp_unit(rng.next_f64())).collect();
        let mut v = vec![vec![0.0_f64; d]; d];

        v[0][0] = w[0];
        result[0][sample_idx] = w[0];

        for j in 1..d {
            v[j][0] = w[j];
            for k in (0..j).rev() {
                let pc = find_pair_copula_cvine(pair_copulas, k, j);
                if let Some(pc) = pc {
                    v[j][0] = inverse_h_func(pc.family, pc.parameter, v[k][k], v[j][0]);
                }
            }
            result[j][sample_idx] = v[j][0];

            // Compute conditional values
            for k in 0..j {
                let pc = find_pair_copula_cvine(pair_copulas, k, j);
                if let Some(pc) = pc {
                    v[j][k + 1] = h_func(pc.family, pc.parameter, v[k][k], v[j][0]);
                }
            }
        }
    }

    result
}

/// Find pair copula in D-vine structure.
fn find_pair_copula_dvine(pair_copulas: &[PairCopula], _d: usize, _tree: usize, var1: usize, var2: usize) -> Option<&PairCopula> {
    pair_copulas.iter().find(|pc| {
        (pc.var1 == var1 && pc.var2 == var2) || (pc.var1 == var2 && pc.var2 == var1)
    })
}

/// Find pair copula in C-vine structure.
fn find_pair_copula_cvine(pair_copulas: &[PairCopula], var1: usize, var2: usize) -> Option<&PairCopula> {
    pair_copulas.iter().find(|pc| {
        (pc.var1 == var1 && pc.var2 == var2) || (pc.var1 == var2 && pc.var2 == var1)
    })
}

/// Inverse h-function via bisection: find v such that h(v|u) = target.
fn inverse_h_func(family: CopulaFamily, param: f64, u: f64, target: f64) -> f64 {
    let u = clamp_unit(u);
    let target = clamp_unit(target);
    let mut lo = 1e-10_f64;
    let mut hi = 1.0 - 1e-10_f64;

    for _ in 0..80 {
        let mid = 0.5 * (lo + hi);
        let h = h_func(family, param, u, mid);
        if h < target {
            lo = mid;
        } else {
            hi = mid;
        }
    }
    clamp_unit(0.5 * (lo + hi))
}

// ===========================================================================
// Validation
// ===========================================================================

fn validate_data(data: &[Vec<f64>]) -> PyResult<(usize, usize)> {
    if data.len() < 3 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "vine copula requires at least 3 variables",
        ));
    }
    let n = data[0].len();
    if n < 10 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "need at least 10 observations per variable",
        ));
    }
    for (i, col) in data.iter().enumerate() {
        if col.len() != n {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "variable {} has {} observations, expected {}", i, col.len(), n
            )));
        }
        for (j, &v) in col.iter().enumerate() {
            if !v.is_finite() {
                return Err(pyo3::exceptions::PyValueError::new_err(format!(
                    "non-finite value at variable {}, index {}", i, j
                )));
            }
        }
    }
    Ok((data.len(), n))
}

// ===========================================================================
// Exported pyfunctions
// ===========================================================================

/// Fit a vine copula (C-vine or D-vine) to multivariate data.
///
/// Parameters
/// ----------
/// data : list[list[float]]
///     Each inner list is one variable's observations (all same length).
/// vine_type : VineType
///     CVine or DVine.
///
/// Returns
/// -------
/// VineCopulaResult
#[pyfunction]
#[pyo3(signature = (data, vine_type))]
fn fit_vine(data: Vec<Vec<f64>>, vine_type: VineType) -> PyResult<VineCopulaResult> {
    crate::auth::require_ultra()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

    let (d, _n) = validate_data(&data)?;

    let (pair_copulas, total_ll) = match vine_type {
        VineType::DVine => fit_dvine(&data),
        VineType::CVine => fit_cvine(&data),
    };

    Ok(VineCopulaResult {
        vine_type,
        pair_copulas,
        n_variables: d,
        log_likelihood: finite_or_zero(total_ll),
    })
}

/// Sample from a fitted vine copula.
///
/// Parameters
/// ----------
/// vine : VineCopulaResult
///     Fitted vine copula model.
/// n : int
///     Number of samples to generate.
/// seed : int
///     Random seed for reproducibility.
///
/// Returns
/// -------
/// list[list[float]]
///     Each inner list is one variable's samples.
#[pyfunction]
#[pyo3(signature = (vine, n, seed))]
fn vine_sample(vine: &VineCopulaResult, n: usize, seed: u64) -> PyResult<Vec<Vec<f64>>> {
    crate::auth::require_ultra()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

    if n == 0 {
        return Ok(vec![vec![]; vine.n_variables]);
    }

    let mut rng = Xoshiro256::new(seed);
    let d = vine.n_variables;

    let result = match vine.vine_type {
        VineType::DVine => sample_dvine(&vine.pair_copulas, d, n, &mut rng),
        VineType::CVine => sample_cvine(&vine.pair_copulas, d, n, &mut rng),
    };

    Ok(result)
}

/// Estimate tail risk (joint exceedance probability) from a vine copula.
///
/// Estimates P(all variables exceed their respective thresholds)
/// via Monte Carlo simulation from the fitted vine.
///
/// Parameters
/// ----------
/// vine : VineCopulaResult
///     Fitted vine copula model.
/// threshold : float
///     Uniform threshold in (0, 1). Joint probability that all
///     variables exceed this quantile.
/// n_sim : int
///     Number of Monte Carlo samples.
/// seed : int
///     Random seed.
///
/// Returns
/// -------
/// float
///     Estimated joint exceedance probability.
#[pyfunction]
#[pyo3(signature = (vine, threshold, n_sim, seed))]
fn vine_tail_risk(vine: &VineCopulaResult, threshold: f64, n_sim: usize, seed: u64) -> PyResult<f64> {
    crate::auth::require_ultra()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

    if !threshold.is_finite() || threshold <= 0.0 || threshold >= 1.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "threshold must be in (0, 1)",
        ));
    }
    if n_sim == 0 {
        return Ok(0.0);
    }

    let samples = vine_sample(vine, n_sim, seed)?;
    let d = vine.n_variables;
    let n = if samples.is_empty() || samples[0].is_empty() {
        0
    } else {
        samples[0].len()
    };

    let mut count = 0usize;
    for i in 0..n {
        let all_exceed = (0..d).all(|j| samples[j][i] > threshold);
        if all_exceed {
            count += 1;
        }
    }

    Ok(finite_or_zero(count as f64 / n_sim as f64))
}

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

pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<VineType>()?;
    m.add_class::<PairCopula>()?;
    m.add_class::<VineCopulaResult>()?;
    m.add_function(wrap_pyfunction!(fit_vine, m)?)?;
    m.add_function(wrap_pyfunction!(vine_sample, m)?)?;
    m.add_function(wrap_pyfunction!(vine_tail_risk, m)?)?;
    Ok(())
}

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

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

    #[test]
    fn test_rank_transform_vec() {
        let data = vec![10.0, 20.0, 30.0, 40.0, 50.0];
        let ranks = rank_transform_vec(&data);
        assert_eq!(ranks.len(), 5);
        for &r in &ranks {
            assert!(r > 0.0 && r < 1.0);
        }
        // Should be monotonically increasing
        for i in 1..ranks.len() {
            assert!(ranks[i] > ranks[i - 1]);
        }
    }

    #[test]
    fn test_inverse_h_func_roundtrip() {
        let family = CopulaFamily::Gaussian;
        let param = 0.5;
        let u = 0.4;
        let v = 0.6;
        let h = h_func(family, param, u, v);
        let v_recovered = inverse_h_func(family, param, u, h);
        assert!(
            (v - v_recovered).abs() < 0.05,
            "roundtrip failed: v={}, recovered={}",
            v,
            v_recovered
        );
    }

    #[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);
    }
}
