//! Robust Portfolio Optimization under parameter uncertainty.
//!
//! Implements worst-case mean-variance optimization where expected returns
//! are uncertain within an ellipsoidal uncertainty set:
//!
//!   max_w { min_mu { w'mu - lambda * w'Sigma*w } }
//!   s.t.  ||mu - mu_hat|| <= epsilon
//!         sum(w) = 1, min_weight <= w_i <= max_weight
//!
//! The inner minimization has a closed form:
//!   worst_case_mu = mu_hat - epsilon * Sigma*w / ||Sigma*w||
//!
//! Outer optimization via projected gradient descent with box + simplex
//! constraints.
//!
//! All math from scratch in Rust -- no external crate dependencies beyond PyO3.
//!
//! Pro tier required.

use pyo3::prelude::*;

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

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

// ===========================================================================
// Linear algebra helpers
// ===========================================================================

/// Matrix-vector multiplication: A (n x n) * v (n) -> w (n).
fn mat_vec_mul(a: &[Vec<f64>], v: &[f64]) -> Vec<f64> {
    let n = a.len();
    let mut w = vec![0.0; n];
    for i in 0..n {
        let mut s = 0.0;
        for j in 0..v.len() {
            s += a[i][j] * v[j];
        }
        w[i] = finite_or_zero(s);
    }
    w
}

fn vec_dot(a: &[f64], b: &[f64]) -> f64 {
    let mut s = 0.0;
    for i in 0..a.len() {
        s += a[i] * b[i];
    }
    finite_or_zero(s)
}

fn vec_norm(a: &[f64]) -> f64 {
    vec_dot(a, a).sqrt()
}

fn vec_add(a: &[f64], b: &[f64]) -> Vec<f64> {
    a.iter().zip(b.iter()).map(|(&x, &y)| x + y).collect()
}

#[allow(dead_code)]
fn vec_sub(a: &[f64], b: &[f64]) -> Vec<f64> {
    a.iter().zip(b.iter()).map(|(&x, &y)| x - y).collect()
}

fn vec_scale(a: &[f64], s: f64) -> Vec<f64> {
    a.iter().map(|&x| finite_or_zero(x * s)).collect()
}

// ===========================================================================
// Projection operators
// ===========================================================================

/// Project weights onto the feasible set:
/// - Each weight in [min_weight, max_weight]
/// - Weights sum to 1.0
fn project_weights(w: &[f64], min_weight: f64, max_weight: f64) -> Vec<f64> {
    let n = w.len();

    // First, clamp to box constraints
    let mut projected: Vec<f64> = w.iter().map(|&wi| wi.clamp(min_weight, max_weight)).collect();

    // Then, project onto simplex (sum = 1) while respecting box constraints
    // Iterative projection (Dykstra-like)
    for _ in 0..100 {
        let total: f64 = projected.iter().sum();
        if (total - 1.0).abs() < 1e-12 {
            break;
        }

        // Distribute excess/deficit equally, respecting bounds
        let excess = total - 1.0;
        let n_free = projected
            .iter()
            .enumerate()
            .filter(|(_, &wi)| {
                if excess > 0.0 {
                    wi > min_weight + 1e-15
                } else {
                    wi < max_weight - 1e-15
                }
            })
            .count();

        if n_free == 0 {
            break;
        }

        let adj = excess / n_free as f64;
        for i in 0..n {
            if excess > 0.0 && projected[i] > min_weight + 1e-15 {
                projected[i] = (projected[i] - adj).clamp(min_weight, max_weight);
            } else if excess < 0.0 && projected[i] < max_weight - 1e-15 {
                projected[i] = (projected[i] - adj).clamp(min_weight, max_weight);
            }
        }
    }

    // Final normalization as safety net
    let total: f64 = projected.iter().sum();
    if total > 0.0 && total.is_finite() {
        for p in &mut projected {
            *p /= total;
            *p = p.clamp(min_weight, max_weight);
        }
    }

    projected
}

// ===========================================================================
// Worst-case return computation
// ===========================================================================

/// Compute worst-case expected return under ellipsoidal uncertainty.
///
/// worst_case = w' * mu_hat - epsilon * ||D * w||
///
/// where D = diag(return_uncertainty).
fn worst_case_return_impl(
    weights: &[f64],
    expected_returns: &[f64],
    return_uncertainty: &[f64],
    budget: f64,
) -> f64 {
    let nominal = vec_dot(weights, expected_returns);

    // Compute ||D * w|| where D = diag(return_uncertainty)
    let dw: Vec<f64> = weights
        .iter()
        .zip(return_uncertainty.iter())
        .map(|(&wi, &di)| wi * di)
        .collect();
    let dw_norm = vec_norm(&dw);

    finite_or_zero(nominal - budget * dw_norm)
}

// ===========================================================================
// Robust optimization via projected gradient descent
// ===========================================================================

/// Objective function: worst-case mean-variance
///   f(w) = w'mu_worst - lambda * w'Sigma*w
/// where mu_worst = mu_hat - epsilon * D*w / ||D*w||
fn robust_objective(
    w: &[f64],
    mu: &[f64],
    sigma: &[Vec<f64>],
    return_uncertainty: &[f64],
    budget: f64,
    risk_aversion: f64,
) -> f64 {
    let wc_ret = worst_case_return_impl(w, mu, return_uncertainty, budget);
    let sigma_w = mat_vec_mul(sigma, w);
    let risk = vec_dot(w, &sigma_w);
    finite_or_zero(wc_ret - risk_aversion * risk)
}

/// Gradient of the robust objective.
///
/// d/dw [w'mu_hat - epsilon*||D*w|| - lambda*w'Sigma*w]
/// = mu_hat - epsilon * D^2 * w / ||D*w|| - 2*lambda*Sigma*w
fn robust_gradient(
    w: &[f64],
    mu: &[f64],
    sigma: &[Vec<f64>],
    return_uncertainty: &[f64],
    budget: f64,
    risk_aversion: f64,
) -> Vec<f64> {
    let n = w.len();

    // d/dw nominal = mu_hat
    let mut grad = mu.to_vec();

    // d/dw (-epsilon * ||D*w||) = -epsilon * D^2 * w / ||D*w||
    let dw: Vec<f64> = w
        .iter()
        .zip(return_uncertainty.iter())
        .map(|(&wi, &di)| wi * di)
        .collect();
    let dw_norm = vec_norm(&dw);
    if dw_norm > 1e-15 {
        for i in 0..n {
            grad[i] -= budget * return_uncertainty[i] * return_uncertainty[i] * w[i] / dw_norm;
        }
    }

    // d/dw (-lambda * w'Sigma*w) = -2*lambda*Sigma*w
    let sigma_w = mat_vec_mul(sigma, w);
    for i in 0..n {
        grad[i] -= 2.0 * risk_aversion * sigma_w[i];
    }

    grad.iter().map(|&g| finite_or_zero(g)).collect()
}

/// Run projected gradient descent for robust portfolio optimization.
fn optimize_robust(
    mu: &[f64],
    sigma: &[Vec<f64>],
    return_uncertainty: &[f64],
    budget: f64,
    max_weight: f64,
    min_weight: f64,
    risk_aversion: f64,
) -> Vec<f64> {
    let n = mu.len();
    let max_iter = 2000;
    let tol = 1e-10;

    // Initialize with equal weights
    let mut w = vec![1.0 / n as f64; n];
    w = project_weights(&w, min_weight, max_weight);

    let mut best_w = w.clone();
    let mut best_obj = robust_objective(&w, mu, sigma, return_uncertainty, budget, risk_aversion);

    let mut step_size = 0.01;

    for iter in 0..max_iter {
        let grad = robust_gradient(&w, mu, sigma, return_uncertainty, budget, risk_aversion);
        let grad_norm = vec_norm(&grad);

        if grad_norm < tol {
            break;
        }

        // Gradient ascent (we are maximizing)
        let w_new_raw = vec_add(&w, &vec_scale(&grad, step_size));
        let w_new = project_weights(&w_new_raw, min_weight, max_weight);

        let new_obj = robust_objective(&w_new, mu, sigma, return_uncertainty, budget, risk_aversion);

        if new_obj > best_obj {
            best_obj = new_obj;
            best_w = w_new.clone();
            w = w_new;
            // Increase step size slightly
            step_size *= 1.05;
            step_size = step_size.min(0.1);
        } else {
            // Decrease step size, keep old w (don't move to worse point)
            step_size *= 0.5;
            if step_size < 1e-15 {
                break;
            }
        }

        // Adaptive step size decay
        if iter > 0 && iter % 200 == 0 {
            step_size *= 0.5;
        }
    }

    best_w
}

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

/// Result of robust portfolio optimization.
#[pyclass(frozen)]
#[derive(Debug, Clone)]
pub struct RobustPortfolioResult {
    /// Optimal portfolio weights (sum to 1).
    #[pyo3(get)]
    pub weights: Vec<f64>,
    /// Worst-case expected return.
    #[pyo3(get)]
    pub worst_case_return: f64,
    /// Nominal (point estimate) expected return.
    #[pyo3(get)]
    pub nominal_return: f64,
    /// Uncertainty budget used.
    #[pyo3(get)]
    pub uncertainty_budget: f64,
}

#[pymethods]
impl RobustPortfolioResult {
    fn __repr__(&self) -> String {
        let n = self.weights.len();
        let top3: Vec<String> = self
            .weights
            .iter()
            .take(3)
            .map(|w| format!("{:.4}", w))
            .collect();
        format!(
            "RobustPortfolioResult(n={}, weights=[{}{}], wc_ret={:.6}, nom_ret={:.6})",
            n,
            top3.join(", "),
            if n > 3 { ", ..." } else { "" },
            self.worst_case_return,
            self.nominal_return
        )
    }
}

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

fn validate_inputs(
    mu: &[f64],
    sigma: &[Vec<f64>],
    uncertainty: &[f64],
) -> PyResult<usize> {
    let n = mu.len();
    if n < 2 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "need at least 2 assets",
        ));
    }
    if sigma.len() != n {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "covariance matrix rows must match number of assets",
        ));
    }
    for (i, row) in sigma.iter().enumerate() {
        if row.len() != n {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "covariance matrix row {} has wrong length", i
            )));
        }
        for (j, &v) in row.iter().enumerate() {
            if !v.is_finite() {
                return Err(pyo3::exceptions::PyValueError::new_err(format!(
                    "non-finite value in covariance at ({}, {})", i, j
                )));
            }
        }
    }
    if uncertainty.len() != n {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "return_uncertainty must match number of assets",
        ));
    }
    for (i, &v) in mu.iter().enumerate() {
        if !v.is_finite() {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "non-finite expected return at index {}", i
            )));
        }
    }
    for (i, &v) in uncertainty.iter().enumerate() {
        if !v.is_finite() || v < 0.0 {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "return_uncertainty[{}] must be non-negative and finite", i
            )));
        }
    }
    Ok(n)
}

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

/// Solve the robust mean-variance optimization problem.
///
/// Parameters
/// ----------
/// expected_returns : list[float]
///     Point estimates of expected returns.
/// covariance : list[list[float]]
///     Covariance matrix (N x N).
/// return_uncertainty : list[float]
///     Uncertainty radius for each asset's expected return.
/// uncertainty_budget : float
///     Overall budget for the uncertainty set (epsilon).
/// max_weight : float
///     Maximum weight per asset.
/// min_weight : float
///     Minimum weight per asset (can be 0 for long-only).
///
/// Returns
/// -------
/// RobustPortfolioResult
#[pyfunction]
#[pyo3(signature = (expected_returns, covariance, return_uncertainty, uncertainty_budget, max_weight=1.0, min_weight=0.0))]
fn robust_optimize(
    expected_returns: Vec<f64>,
    covariance: Vec<Vec<f64>>,
    return_uncertainty: Vec<f64>,
    uncertainty_budget: f64,
    max_weight: f64,
    min_weight: f64,
) -> PyResult<RobustPortfolioResult> {
    crate::auth::require_pro()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

    let n = validate_inputs(&expected_returns, &covariance, &return_uncertainty)?;

    if !uncertainty_budget.is_finite() || uncertainty_budget < 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "uncertainty_budget must be non-negative and finite",
        ));
    }
    if !max_weight.is_finite() || !min_weight.is_finite() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "weight bounds must be finite",
        ));
    }
    if min_weight > max_weight {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "min_weight must be <= max_weight",
        ));
    }
    if min_weight * (n as f64) > 1.0 + 1e-10 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "min_weight * n_assets > 1: infeasible",
        ));
    }
    if max_weight * (n as f64) < (1.0 - 1e-10) {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "max_weight * n_assets < 1: infeasible",
        ));
    }

    // Risk aversion: scale based on return magnitudes
    let avg_return: f64 = expected_returns.iter().map(|r| r.abs()).sum::<f64>() / n as f64;
    let risk_aversion = if avg_return > 1e-10 { 0.5 / avg_return } else { 1.0 };

    let weights = optimize_robust(
        &expected_returns,
        &covariance,
        &return_uncertainty,
        uncertainty_budget,
        max_weight,
        min_weight,
        risk_aversion,
    );

    let wc = worst_case_return_impl(&weights, &expected_returns, &return_uncertainty, uncertainty_budget);
    let nominal = vec_dot(&weights, &expected_returns);

    Ok(RobustPortfolioResult {
        weights,
        worst_case_return: finite_or_zero(wc),
        nominal_return: finite_or_zero(nominal),
        uncertainty_budget,
    })
}

/// Compute the worst-case expected return for given weights.
///
/// Parameters
/// ----------
/// weights : list[float]
/// expected_returns : list[float]
/// return_uncertainty : list[float]
/// budget : float
///
/// Returns
/// -------
/// float
#[pyfunction]
#[pyo3(signature = (weights, expected_returns, return_uncertainty, budget))]
fn worst_case_return(
    weights: Vec<f64>,
    expected_returns: Vec<f64>,
    return_uncertainty: Vec<f64>,
    budget: f64,
) -> PyResult<f64> {
    crate::auth::require_pro()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

    if weights.len() != expected_returns.len() || weights.len() != return_uncertainty.len() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "all vectors must have the same length",
        ));
    }
    Ok(finite_or_zero(worst_case_return_impl(
        &weights,
        &expected_returns,
        &return_uncertainty,
        budget,
    )))
}

/// Compute the robust efficient frontier.
///
/// Parameters
/// ----------
/// expected_returns : list[float]
/// covariance : list[list[float]]
/// return_uncertainty : list[float]
/// n_points : int
///     Number of points on the frontier.
///
/// Returns
/// -------
/// list[tuple[float, float]]
///     (worst_case_return, portfolio_risk) pairs.
#[pyfunction]
#[pyo3(signature = (expected_returns, covariance, return_uncertainty, n_points=20))]
fn robust_efficient_frontier(
    expected_returns: Vec<f64>,
    covariance: Vec<Vec<f64>>,
    return_uncertainty: Vec<f64>,
    n_points: usize,
) -> PyResult<Vec<(f64, f64)>> {
    crate::auth::require_pro()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

    let n = validate_inputs(&expected_returns, &covariance, &return_uncertainty)?;

    if n_points < 2 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "n_points must be at least 2",
        ));
    }

    let mut frontier = Vec::with_capacity(n_points);

    // Vary uncertainty budget from 0 (no robustness) to max
    let max_budget = return_uncertainty.iter().copied().fold(0.0_f64, f64::max) * 3.0;
    let max_budget = max_budget.max(0.5);

    for i in 0..n_points {
        let budget = max_budget * i as f64 / (n_points - 1) as f64;

        let avg_return: f64 = expected_returns.iter().map(|r| r.abs()).sum::<f64>() / n as f64;
        let risk_aversion = if avg_return > 1e-10 { 0.5 / avg_return } else { 1.0 };

        let weights = optimize_robust(
            &expected_returns,
            &covariance,
            &return_uncertainty,
            budget,
            1.0,
            0.0,
            risk_aversion,
        );

        let wc = worst_case_return_impl(&weights, &expected_returns, &return_uncertainty, budget);
        let sigma_w = mat_vec_mul(&covariance, &weights);
        let risk = vec_dot(&weights, &sigma_w).sqrt();

        frontier.push((finite_or_zero(wc), finite_or_zero(risk)));
    }

    Ok(frontier)
}

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

pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<RobustPortfolioResult>()?;
    m.add_function(wrap_pyfunction!(robust_optimize, m)?)?;
    m.add_function(wrap_pyfunction!(worst_case_return, m)?)?;
    m.add_function(wrap_pyfunction!(robust_efficient_frontier, m)?)?;
    Ok(())
}

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

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

    #[test]
    fn test_project_weights_equal() {
        let w = vec![0.5, 0.5];
        let proj = project_weights(&w, 0.0, 1.0);
        let sum: f64 = proj.iter().sum();
        assert!((sum - 1.0).abs() < 1e-6);
    }

    #[test]
    fn test_project_weights_box() {
        let w = vec![2.0, -1.0, 0.5];
        let proj = project_weights(&w, 0.0, 0.5);
        assert!(proj.iter().all(|&wi| wi >= -1e-10 && wi <= 0.5 + 1e-10));
    }

    #[test]
    fn test_worst_case_return_no_uncertainty() {
        let w = vec![0.5, 0.5];
        let mu = vec![0.1, 0.2];
        let unc = vec![0.0, 0.0];
        let wc = worst_case_return_impl(&w, &mu, &unc, 1.0);
        // No uncertainty => wc = nominal = 0.5*0.1 + 0.5*0.2 = 0.15
        assert!((wc - 0.15).abs() < 1e-10);
    }

    #[test]
    fn test_worst_case_return_with_uncertainty() {
        let w = vec![0.5, 0.5];
        let mu = vec![0.1, 0.2];
        let unc = vec![0.05, 0.05];
        let wc = worst_case_return_impl(&w, &mu, &unc, 1.0);
        // Should be less than nominal
        let nominal = 0.15;
        assert!(wc < nominal);
    }

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

    #[test]
    fn test_mat_vec_mul() {
        let a = vec![vec![1.0, 2.0], vec![3.0, 4.0]];
        let v = vec![1.0, 1.0];
        let w = mat_vec_mul(&a, &v);
        assert!((w[0] - 3.0).abs() < 1e-10);
        assert!((w[1] - 7.0).abs() < 1e-10);
    }
}
