//! Entropy Pooling (Meucci, 2008) for Bayesian scenario analysis.
//!
//! Minimizes relative entropy KL(p || q) subject to linear view constraints,
//! finding the posterior distribution closest to the prior that satisfies
//! the investor's views on expected values.
//!
//! Dual problem solved via Newton optimization with log-sum-exp for
//! numerical stability.
//!
//! All math from scratch in Rust -- no external crate dependencies beyond PyO3.
//!
//! Ultra tier required.

use pyo3::prelude::*;

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

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

/// Log-sum-exp for numerical stability: log(sum(exp(x_i))).
fn log_sum_exp(values: &[f64]) -> f64 {
    if values.is_empty() {
        return f64::NEG_INFINITY;
    }
    let max_val = values.iter().copied().fold(f64::NEG_INFINITY, f64::max);
    if !max_val.is_finite() {
        return max_val;
    }
    let sum: f64 = values.iter().map(|&x| (x - max_val).exp()).sum();
    finite_or_zero(max_val + sum.ln())
}

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

/// Result of entropy pooling optimization.
#[pyclass(frozen)]
#[derive(Debug, Clone)]
pub struct EntropyPoolResult {
    /// Posterior scenario probabilities (sum to 1).
    #[pyo3(get)]
    pub posterior_probs: Vec<f64>,
    /// Prior scenario probabilities (input).
    #[pyo3(get)]
    pub prior_probs: Vec<f64>,
    /// Relative entropy KL(posterior || prior).
    #[pyo3(get)]
    pub relative_entropy: f64,
    /// Effective number of scenarios = exp(-sum p_i ln p_i).
    #[pyo3(get)]
    pub effective_scenarios: f64,
}

#[pymethods]
impl EntropyPoolResult {
    fn __repr__(&self) -> String {
        format!(
            "EntropyPoolResult(n_scenarios={}, rel_entropy={:.4}, eff_scenarios={:.1})",
            self.posterior_probs.len(),
            self.relative_entropy,
            self.effective_scenarios
        )
    }
}

// ===========================================================================
// View constraint
// ===========================================================================

/// A linear view constraint on the scenario set.
///
/// For equality: sum_j coefficients[j] * scenarios[j] = bound
/// For inequality: sum_j coefficients[j] * scenarios[j] <= bound
#[pyclass(frozen)]
#[derive(Debug, Clone)]
pub struct ViewConstraint {
    /// Coefficients for each scenario variable.
    #[pyo3(get)]
    pub coefficients: Vec<f64>,
    /// Right-hand side bound.
    #[pyo3(get)]
    pub bound: f64,
    /// If true, equality constraint; otherwise inequality (<=).
    #[pyo3(get)]
    pub is_equality: bool,
}

#[pymethods]
impl ViewConstraint {
    #[new]
    #[pyo3(signature = (coefficients, bound, is_equality=true))]
    fn new(coefficients: Vec<f64>, bound: f64, is_equality: bool) -> PyResult<Self> {
        if coefficients.is_empty() {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "coefficients must not be empty",
            ));
        }
        if !bound.is_finite() {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "bound must be finite",
            ));
        }
        for (i, &c) in coefficients.iter().enumerate() {
            if !c.is_finite() {
                return Err(pyo3::exceptions::PyValueError::new_err(format!(
                    "non-finite coefficient at index {}", i
                )));
            }
        }
        Ok(Self {
            coefficients,
            bound,
            is_equality,
        })
    }

    fn __repr__(&self) -> String {
        let kind = if self.is_equality { "==" } else { "<=" };
        format!(
            "ViewConstraint(n_coeffs={}, bound={:.4}, {})",
            self.coefficients.len(),
            self.bound,
            kind
        )
    }
}

// ===========================================================================
// Newton solver for the dual problem
// ===========================================================================

/// Compute posterior probabilities given dual variables (lambda, nu).
///
/// For equality views on E[H*x] = h:
///   p_j = q_j * exp(sum_k lambda_k * H_kj) / Z
///
/// where Z is the normalizing constant.
fn compute_posterior(
    log_q: &[f64],
    a_matrix: &[Vec<f64>],
    lambda: &[f64],
) -> Vec<f64> {
    let j = log_q.len();
    let k = lambda.len();

    // log_p_j = log_q_j + sum_k lambda_k * A_{k,j}
    let mut log_p = vec![0.0_f64; j];
    for jj in 0..j {
        let mut s = log_q[jj];
        for kk in 0..k {
            s += lambda[kk] * a_matrix[kk][jj];
        }
        log_p[jj] = s;
    }

    // Normalize via log-sum-exp
    let lse = log_sum_exp(&log_p);
    let mut p = vec![0.0; j];
    for jj in 0..j {
        p[jj] = finite_or_zero((log_p[jj] - lse).exp()).max(1e-30);
    }

    // Ensure sum = 1
    let total: f64 = p.iter().sum();
    if total > 0.0 && total.is_finite() {
        for pj in &mut p {
            *pj /= total;
        }
    }

    p
}

/// Compute gradient and Hessian of the dual objective.
///
/// The dual function is:
///   L(lambda) = log(sum_j q_j exp(A' lambda)_j) - lambda' b
///
/// Gradient: g_k = E_p[A_k] - b_k
/// Hessian: H_{k,l} = E_p[A_k A_l] - E_p[A_k] E_p[A_l]
fn dual_gradient_hessian(
    p: &[f64],
    a_matrix: &[Vec<f64>],
    b: &[f64],
) -> (Vec<f64>, Vec<Vec<f64>>) {
    let k = a_matrix.len();
    let j = p.len();

    // Gradient: g_k = sum_j p_j * A_{k,j} - b_k
    let mut grad = vec![0.0; k];
    for kk in 0..k {
        let mut s = 0.0;
        for jj in 0..j {
            s += p[jj] * a_matrix[kk][jj];
        }
        grad[kk] = finite_or_zero(s - b[kk]);
    }

    // Hessian: H_{k,l} = sum_j p_j * A_{k,j} * A_{l,j} - (sum_j p_j*A_k) * (sum_j p_j*A_l)
    let mut hess = vec![vec![0.0; k]; k];
    let mut means = vec![0.0; k];
    for kk in 0..k {
        for jj in 0..j {
            means[kk] += p[jj] * a_matrix[kk][jj];
        }
    }

    for kk in 0..k {
        for ll in 0..=kk {
            let mut s = 0.0;
            for jj in 0..j {
                s += p[jj] * a_matrix[kk][jj] * a_matrix[ll][jj];
            }
            let val = finite_or_zero(s - means[kk] * means[ll]);
            hess[kk][ll] = val;
            hess[ll][kk] = val;
        }
        // Regularize diagonal
        hess[kk][kk] += 1e-10;
    }

    (grad, hess)
}

/// Solve linear system Hx = g via Cholesky decomposition.
fn cholesky_solve(h: &[Vec<f64>], g: &[f64]) -> Vec<f64> {
    let n = h.len();

    // Cholesky: H = L L^T
    let mut l = vec![vec![0.0; n]; n];
    for i in 0..n {
        for j in 0..=i {
            let mut s = 0.0;
            for kk in 0..j {
                s += l[i][kk] * l[j][kk];
            }
            if i == j {
                let diag = h[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 {
                    (h[i][j] - s) / l[j][j]
                } else {
                    0.0
                };
            }
        }
    }

    // Forward: L z = g
    let mut z = vec![0.0; n];
    for i in 0..n {
        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 {
            (g[i] - s) / l[i][i]
        } else {
            0.0
        };
    }

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

    x
}

/// Newton optimization of the dual problem.
fn newton_optimize(
    log_q: &[f64],
    a_matrix: &[Vec<f64>],
    b: &[f64],
    max_iter: usize,
    tol: f64,
) -> Vec<f64> {
    let k = a_matrix.len();
    let mut lambda = vec![0.0; k];

    for _iter in 0..max_iter {
        let p = compute_posterior(log_q, a_matrix, &lambda);
        let (grad, hess) = dual_gradient_hessian(&p, a_matrix, b);

        // Check convergence
        let grad_norm: f64 = grad.iter().map(|g| g * g).sum::<f64>().sqrt();
        if grad_norm < tol {
            break;
        }

        // Newton step: delta = H^{-1} * g
        let delta = cholesky_solve(&hess, &grad);

        // Line search with backtracking
        let mut step = 1.0;
        for _ in 0..30 {
            let new_lambda: Vec<f64> = lambda
                .iter()
                .zip(delta.iter())
                .map(|(&l, &d)| l - step * d)
                .collect();
            let new_p = compute_posterior(log_q, a_matrix, &new_lambda);
            let (new_grad, _) = dual_gradient_hessian(&new_p, a_matrix, b);
            let new_norm: f64 = new_grad.iter().map(|g| g * g).sum::<f64>().sqrt();
            if new_norm < grad_norm {
                lambda = new_lambda;
                break;
            }
            step *= 0.5;
            if step < 1e-10 {
                // Accept anyway
                lambda = new_lambda;
                break;
            }
        }
    }

    lambda
}

// ===========================================================================
// KL divergence and effective scenarios
// ===========================================================================

fn kl_divergence(p: &[f64], q: &[f64]) -> f64 {
    let mut kl = 0.0;
    for i in 0..p.len() {
        if p[i] > 1e-30 && q[i] > 1e-30 {
            kl += p[i] * (p[i] / q[i]).ln();
        }
    }
    finite_or_zero(kl.max(0.0))
}

fn effective_scenarios(p: &[f64]) -> f64 {
    let mut entropy = 0.0;
    for &pi in p {
        if pi > 1e-30 {
            entropy -= pi * pi.ln();
        }
    }
    finite_or_zero(entropy.exp())
}

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

/// Perform entropy pooling: find posterior probabilities closest to prior
/// that satisfy the given view constraints.
///
/// Parameters
/// ----------
/// prior_probs : list[float]
///     Prior scenario probabilities (must sum to ~1).
/// scenarios : list[list[float]]
///     Each inner list is one scenario (all same length = n_variables).
/// views : list[ViewConstraint]
///     Linear constraints on expected values.
///
/// Returns
/// -------
/// EntropyPoolResult
#[pyfunction]
#[pyo3(signature = (prior_probs, scenarios, views))]
fn entropy_pool(
    prior_probs: Vec<f64>,
    scenarios: Vec<Vec<f64>>,
    views: Vec<ViewConstraint>,
) -> PyResult<EntropyPoolResult> {
    crate::auth::require_ultra()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

    // Validate
    let j = prior_probs.len(); // number of scenarios
    if j < 2 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "need at least 2 scenarios",
        ));
    }
    if scenarios.len() != j {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "scenarios length must match prior_probs length",
        ));
    }
    if views.is_empty() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "need at least one view constraint",
        ));
    }

    let n_vars = scenarios[0].len();
    for (i, s) in scenarios.iter().enumerate() {
        if s.len() != n_vars {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "scenario {} has {} variables, expected {}", i, s.len(), n_vars
            )));
        }
        for (jj, &v) in s.iter().enumerate() {
            if !v.is_finite() {
                return Err(pyo3::exceptions::PyValueError::new_err(format!(
                    "non-finite value in scenario {}, variable {}", i, jj
                )));
            }
        }
    }

    // Validate prior probabilities
    for (i, &p) in prior_probs.iter().enumerate() {
        if !p.is_finite() || p < 0.0 {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "prior_probs[{}] must be non-negative and finite", i
            )));
        }
    }
    let total: f64 = prior_probs.iter().sum();
    if (total - 1.0).abs() > 0.01 {
        return Err(pyo3::exceptions::PyValueError::new_err(format!(
            "prior_probs must sum to ~1 (got {})", total
        )));
    }

    // Validate view coefficients length
    for (i, view) in views.iter().enumerate() {
        if view.coefficients.len() != n_vars {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "view {} has {} coefficients, expected {} (one per variable)",
                i, view.coefficients.len(), n_vars
            )));
        }
    }

    // Build constraint matrix A (k x j): A_{k,j} = sum_v coefficients[v] * scenarios[j][v]
    let k = views.len();
    let mut a_matrix = vec![vec![0.0; j]; k];
    let mut b = vec![0.0; k];

    for (kk, view) in views.iter().enumerate() {
        for jj in 0..j {
            let mut s = 0.0;
            for v in 0..n_vars {
                s += view.coefficients[v] * scenarios[jj][v];
            }
            a_matrix[kk][jj] = finite_or_zero(s);
        }
        b[kk] = view.bound;

        // For inequality constraints, use slack variable approach:
        // we convert E[A_k x] <= b_k to E[A_k x] = b_k (pessimistic)
        // In practice, the Newton solver handles this naturally
    }

    // Log prior
    let log_q: Vec<f64> = prior_probs
        .iter()
        .map(|&p| if p > 1e-30 { p.ln() } else { -30.0 })
        .collect();

    // Solve dual
    let lambda = newton_optimize(&log_q, &a_matrix, &b, 500, 1e-10);

    // Compute posterior
    let posterior = compute_posterior(&log_q, &a_matrix, &lambda);

    // KL divergence
    let kl = kl_divergence(&posterior, &prior_probs);
    let eff = effective_scenarios(&posterior);

    Ok(EntropyPoolResult {
        posterior_probs: posterior,
        prior_probs,
        relative_entropy: finite_or_zero(kl),
        effective_scenarios: finite_or_zero(eff),
    })
}

/// Compute posterior mean from weighted scenarios.
///
/// Parameters
/// ----------
/// probs : list[float]
///     Scenario probabilities.
/// scenarios : list[list[float]]
///     Scenario values (n_scenarios x n_variables).
///
/// Returns
/// -------
/// list[float]
///     Posterior mean for each variable.
#[pyfunction]
#[pyo3(signature = (probs, scenarios))]
fn posterior_mean(probs: Vec<f64>, scenarios: Vec<Vec<f64>>) -> PyResult<Vec<f64>> {
    crate::auth::require_ultra()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

    if probs.len() != scenarios.len() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "probs and scenarios must have same length",
        ));
    }
    if scenarios.is_empty() {
        return Ok(vec![]);
    }

    let n_vars = scenarios[0].len();
    let j = probs.len();
    let mut mean = vec![0.0; n_vars];

    for jj in 0..j {
        for v in 0..n_vars {
            mean[v] += probs[jj] * scenarios[jj][v];
        }
    }

    Ok(mean.into_iter().map(finite_or_zero).collect())
}

/// Compute posterior covariance from weighted scenarios.
///
/// Parameters
/// ----------
/// probs : list[float]
///     Scenario probabilities.
/// scenarios : list[list[float]]
///     Scenario values (n_scenarios x n_variables).
///
/// Returns
/// -------
/// list[list[float]]
///     Posterior covariance matrix (n_variables x n_variables).
#[pyfunction]
#[pyo3(signature = (probs, scenarios))]
fn posterior_covariance(probs: Vec<f64>, scenarios: Vec<Vec<f64>>) -> PyResult<Vec<Vec<f64>>> {
    crate::auth::require_ultra()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

    if probs.len() != scenarios.len() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "probs and scenarios must have same length",
        ));
    }
    if scenarios.is_empty() {
        return Ok(vec![]);
    }

    let n_vars = scenarios[0].len();
    let j = probs.len();

    // Compute mean
    let mut mean = vec![0.0; n_vars];
    for jj in 0..j {
        for v in 0..n_vars {
            mean[v] += probs[jj] * scenarios[jj][v];
        }
    }

    // Compute covariance: Cov = sum_j p_j * (x_j - mu)(x_j - mu)^T
    let mut cov = vec![vec![0.0; n_vars]; n_vars];
    for jj in 0..j {
        for vi in 0..n_vars {
            let di = scenarios[jj][vi] - mean[vi];
            for vj in 0..=vi {
                let dj = scenarios[jj][vj] - mean[vj];
                let val = probs[jj] * di * dj;
                cov[vi][vj] += val;
                if vi != vj {
                    cov[vj][vi] += val;
                }
            }
        }
    }

    // Ensure finite
    for row in &mut cov {
        for val in row.iter_mut() {
            *val = finite_or_zero(*val);
        }
    }

    Ok(cov)
}

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

pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<EntropyPoolResult>()?;
    m.add_class::<ViewConstraint>()?;
    m.add_function(wrap_pyfunction!(entropy_pool, m)?)?;
    m.add_function(wrap_pyfunction!(posterior_mean, m)?)?;
    m.add_function(wrap_pyfunction!(posterior_covariance, m)?)?;
    Ok(())
}

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

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

    #[test]
    fn test_log_sum_exp() {
        let vals = vec![1.0, 2.0, 3.0];
        let lse = log_sum_exp(&vals);
        let expected = (1.0_f64.exp() + 2.0_f64.exp() + 3.0_f64.exp()).ln();
        assert!((lse - expected).abs() < 1e-10);
    }

    #[test]
    fn test_log_sum_exp_large() {
        // Test numerical stability with large values
        let vals = vec![1000.0, 1001.0, 1002.0];
        let lse = log_sum_exp(&vals);
        assert!(lse.is_finite());
        assert!(lse > 1000.0);
    }

    #[test]
    fn test_kl_divergence_same() {
        let p = vec![0.25, 0.25, 0.25, 0.25];
        let kl = kl_divergence(&p, &p);
        assert!(kl.abs() < 1e-10);
    }

    #[test]
    fn test_effective_scenarios_uniform() {
        let p = vec![0.25, 0.25, 0.25, 0.25];
        let eff = effective_scenarios(&p);
        assert!((eff - 4.0).abs() < 1e-6);
    }

    #[test]
    fn test_effective_scenarios_concentrated() {
        let p = vec![0.97, 0.01, 0.01, 0.01];
        let eff = effective_scenarios(&p);
        assert!(eff < 2.0);
    }

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

    #[test]
    fn test_cholesky_solve_simple() {
        // Solve [[2, 1], [1, 2]] x = [3, 3] => x = [1, 1]
        let h = vec![vec![2.0, 1.0], vec![1.0, 2.0]];
        let g = vec![3.0, 3.0];
        let x = cholesky_solve(&h, &g);
        assert!((x[0] - 1.0).abs() < 1e-6);
        assert!((x[1] - 1.0).abs() < 1e-6);
    }
}
