//! Almgren-Chriss (2000) Optimal Liquidation.
//!
//! Implements the classic optimal execution framework from "Optimal Execution
//! of Portfolio Transactions" (Almgren & Chriss, 2000).
//!
//! The model balances execution risk (price variance from holding inventory)
//! against market impact costs (temporary and permanent price impact from
//! trading). The solution yields a deterministic trading schedule that
//! minimizes a mean-variance objective.
//!
//! Key results:
//! - Optimal schedule: n_j = X * sinh(kappa*(T-t_j)) / sinh(kappa*T)
//! - kappa controls the liquidation speed (balances risk vs. impact)
//! - Efficient frontier: set of (expected cost, variance) trade-offs
//!
//! All math from scratch in Rust -- no external crate dependencies beyond PyO3.

use pyo3::prelude::*;

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

/// Optimal liquidation trajectory from the Almgren-Chriss model.
#[pyclass(frozen)]
#[derive(Debug, Clone)]
pub struct ACTrajectory {
    /// Remaining shares at each time step (n_steps + 1 entries, from X to 0).
    #[pyo3(get)]
    pub schedule: Vec<f64>,
    /// Expected execution cost (implementation shortfall).
    #[pyo3(get)]
    pub expected_cost: f64,
    /// Variance of execution cost.
    #[pyo3(get)]
    pub variance: f64,
    /// Efficient frontier points: Vec of (expected_cost, variance).
    #[pyo3(get)]
    pub efficient_frontier: Vec<(f64, f64)>,
}

#[pymethods]
impl ACTrajectory {
    fn __repr__(&self) -> String {
        format!(
            "ACTrajectory(steps={}, expected_cost={:.6}, variance={:.6}, frontier_points={})",
            self.schedule.len().saturating_sub(1),
            self.expected_cost,
            self.variance,
            self.efficient_frontier.len()
        )
    }
}

// ===========================================================================
// Private helpers
// ===========================================================================

/// Clamp a value to finite, returning 0.0 for NaN/Inf.
#[inline]
fn finite_or_zero(v: f64) -> f64 {
    if v.is_finite() { v } else { 0.0 }
}

/// Compute the AC schedule for a given risk aversion, returning
/// (schedule, expected_cost, variance).
fn compute_ac_schedule(
    total_shares: f64,
    n_steps: usize,
    volatility: f64,
    permanent_impact: f64,
    temporary_impact: f64,
    risk_aversion: f64,
) -> (Vec<f64>, f64, f64) {
    let x = total_shares;
    let n = n_steps;
    let tau = 1.0 / n as f64; // normalized time step (T=1)
    let sigma = volatility;
    let gamma = permanent_impact;
    let eta = temporary_impact;

    // kappa = arccosh(1 + (risk_aversion * sigma^2 * tau^2) / (2 * eta))
    // adjusted for permanent impact
    let kappa_arg = 1.0 + (risk_aversion * sigma * sigma * tau * tau) / (2.0 * eta);
    let kappa_arg = kappa_arg.max(1.0); // ensure >= 1 for arccosh

    // arccosh(x) = ln(x + sqrt(x^2 - 1))
    let kappa = (kappa_arg + (kappa_arg * kappa_arg - 1.0).max(0.0).sqrt()).ln();
    let kappa = if kappa.is_finite() && kappa > 0.0 {
        kappa
    } else {
        1e-10 // fallback for degenerate case
    };

    // Optimal schedule: n_j = X * sinh(kappa * (T - t_j)) / sinh(kappa * T)
    // where T = n * tau, t_j = j * tau
    let kappa_t = kappa * n as f64; // kappa * T (T = n*tau, but kappa absorbs tau scaling)
    let sinh_kt = kappa_t.sinh();

    let mut schedule = Vec::with_capacity(n + 1);
    for j in 0..=n {
        let t_j = j as f64; // j * tau, but kappa already normalized
        let remaining_time = n as f64 - t_j;
        let n_j = if sinh_kt.abs() > 1e-30 {
            x * (kappa * remaining_time).sinh() / sinh_kt
        } else {
            // Linear fallback when kappa ~ 0 (no risk aversion)
            x * (1.0 - j as f64 / n as f64)
        };
        schedule.push(finite_or_zero(n_j));
    }

    // Trade list: x_j = n_{j-1} - n_j (shares sold in period j)
    let mut trades = Vec::with_capacity(n);
    for j in 1..=n {
        trades.push(schedule[j - 1] - schedule[j]);
    }

    // Expected cost (implementation shortfall)
    // E[cost] = 0.5 * gamma * X^2 + eta * sum(x_j^2 / tau)
    // The permanent impact cost is independent of schedule
    let permanent_cost = 0.5 * gamma * x * x;
    let temporary_cost: f64 = trades.iter().map(|xj| eta * xj * xj / tau).sum();
    let expected_cost = permanent_cost + temporary_cost;

    // Variance of cost
    // Var[cost] = sigma^2 * tau * sum(n_j^2 for j=1..N)
    let variance: f64 = sigma * sigma * tau * schedule[1..].iter().map(|nj| nj * nj).sum::<f64>();

    (
        schedule,
        finite_or_zero(expected_cost),
        finite_or_zero(variance),
    )
}

// ===========================================================================
// Public functions
// ===========================================================================

/// Compute the Almgren-Chriss kappa parameter.
///
/// kappa controls the speed of liquidation:
/// - Large kappa: aggressive (TWAP-like), minimizes variance but higher impact
/// - Small kappa: patient, minimizes impact but higher variance
///
/// kappa = arccosh(1 + risk_aversion * volatility^2 / (2 * temporary_impact * n^2))
///
/// Args:
///     volatility: Asset price volatility (positive).
///     permanent_impact: Permanent market impact coefficient (non-negative).
///     temporary_impact: Temporary market impact coefficient (positive).
///     risk_aversion: Risk aversion parameter (positive).
///
/// Returns:
///     kappa parameter (positive scalar).
#[pyfunction]
pub fn ac_kappa(
    volatility: f64,
    permanent_impact: f64,
    temporary_impact: f64,
    risk_aversion: f64,
) -> PyResult<f64> {
    crate::auth::require_ultra()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

    if !volatility.is_finite() || volatility <= 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "volatility must be a positive finite number",
        ));
    }
    if !permanent_impact.is_finite() || permanent_impact < 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "permanent_impact must be a non-negative finite number",
        ));
    }
    if !temporary_impact.is_finite() || temporary_impact <= 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "temporary_impact must be a positive finite number",
        ));
    }
    if !risk_aversion.is_finite() || risk_aversion <= 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "risk_aversion must be a positive finite number",
        ));
    }

    // kappa = arccosh(1 + risk_aversion * sigma^2 / (2 * eta))
    let kappa_arg = 1.0 + risk_aversion * volatility * volatility / (2.0 * temporary_impact);
    // arccosh(x) = ln(x + sqrt(x^2 - 1))
    let kappa = (kappa_arg + (kappa_arg * kappa_arg - 1.0).max(0.0).sqrt()).ln();

    Ok(finite_or_zero(kappa.max(0.0)))
}

/// Compute the Almgren-Chriss optimal liquidation schedule.
///
/// Solves the AC2000 mean-variance optimization to find the trading schedule
/// that minimizes expected cost + risk_aversion * variance.
///
/// The optimal schedule is:
///     n_j = total_shares * sinh(kappa * (T - t_j)) / sinh(kappa * T)
///
/// Args:
///     total_shares: Total number of shares to liquidate (positive).
///     n_steps: Number of discrete trading periods.
///     volatility: Asset price volatility (positive).
///     permanent_impact: Permanent market impact coefficient (non-negative).
///     temporary_impact: Temporary market impact coefficient (positive).
///     risk_aversion: Risk aversion parameter (positive).
///
/// Returns:
///     ACTrajectory with optimal schedule, expected cost, variance, and frontier.
#[pyfunction]
#[pyo3(signature = (total_shares, n_steps, volatility, permanent_impact, temporary_impact, risk_aversion))]
pub fn ac_optimal_schedule(
    total_shares: f64,
    n_steps: usize,
    volatility: f64,
    permanent_impact: f64,
    temporary_impact: f64,
    risk_aversion: f64,
) -> PyResult<ACTrajectory> {
    crate::auth::require_ultra()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

    if !total_shares.is_finite() || total_shares <= 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "total_shares must be a positive finite number",
        ));
    }
    if n_steps == 0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "n_steps must be at least 1",
        ));
    }
    if !volatility.is_finite() || volatility <= 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "volatility must be a positive finite number",
        ));
    }
    if !permanent_impact.is_finite() || permanent_impact < 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "permanent_impact must be a non-negative finite number",
        ));
    }
    if !temporary_impact.is_finite() || temporary_impact <= 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "temporary_impact must be a positive finite number",
        ));
    }
    if !risk_aversion.is_finite() || risk_aversion <= 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "risk_aversion must be a positive finite number",
        ));
    }

    let n_steps = n_steps.min(100_000);

    let (schedule, expected_cost, variance) = compute_ac_schedule(
        total_shares,
        n_steps,
        volatility,
        permanent_impact,
        temporary_impact,
        risk_aversion,
    );

    // Generate a small efficient frontier around the given risk aversion
    let frontier = compute_efficient_frontier(
        total_shares,
        n_steps,
        volatility,
        permanent_impact,
        temporary_impact,
        10,
    );

    Ok(ACTrajectory {
        schedule,
        expected_cost,
        variance,
        efficient_frontier: frontier,
    })
}

/// Compute the efficient frontier for AC optimal execution.
///
/// Traces out the set of (expected_cost, variance) pairs achievable by
/// varying the risk aversion parameter from very patient to very aggressive.
///
/// Args:
///     total_shares: Total shares to liquidate.
///     n_steps: Number of trading periods.
///     volatility: Asset volatility.
///     permanent_impact: Permanent impact coefficient.
///     temporary_impact: Temporary impact coefficient.
///     n_points: Number of points on the frontier.
///
/// Returns:
///     List of (expected_cost, variance) tuples.
#[pyfunction]
#[pyo3(signature = (total_shares, n_steps, volatility, permanent_impact, temporary_impact, n_points))]
pub fn ac_efficient_frontier(
    total_shares: f64,
    n_steps: usize,
    volatility: f64,
    permanent_impact: f64,
    temporary_impact: f64,
    n_points: usize,
) -> PyResult<Vec<(f64, f64)>> {
    crate::auth::require_ultra()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

    if !total_shares.is_finite() || total_shares <= 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "total_shares must be a positive finite number",
        ));
    }
    if n_steps == 0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "n_steps must be at least 1",
        ));
    }
    if !volatility.is_finite() || volatility <= 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "volatility must be a positive finite number",
        ));
    }
    if !permanent_impact.is_finite() || permanent_impact < 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "permanent_impact must be a non-negative finite number",
        ));
    }
    if !temporary_impact.is_finite() || temporary_impact <= 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "temporary_impact must be a positive finite number",
        ));
    }
    if n_points == 0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "n_points must be at least 1",
        ));
    }

    let n_points = n_points.min(1000);
    let n_steps = n_steps.min(100_000);

    Ok(compute_efficient_frontier(
        total_shares,
        n_steps,
        volatility,
        permanent_impact,
        temporary_impact,
        n_points,
    ))
}

/// Internal efficient frontier computation.
fn compute_efficient_frontier(
    total_shares: f64,
    n_steps: usize,
    volatility: f64,
    permanent_impact: f64,
    temporary_impact: f64,
    n_points: usize,
) -> Vec<(f64, f64)> {
    let mut frontier = Vec::with_capacity(n_points);

    // Sweep risk aversion on a log scale from very patient to very aggressive
    let log_min = -4.0f64; // lambda = 0.0001 (very patient)
    let log_max = 4.0f64;  // lambda = 10000 (very aggressive)

    for i in 0..n_points {
        let frac = if n_points > 1 {
            i as f64 / (n_points - 1) as f64
        } else {
            0.5
        };
        let log_lambda = log_min + frac * (log_max - log_min);
        let lambda = (10.0f64).powf(log_lambda);

        let (_, cost, var) = compute_ac_schedule(
            total_shares,
            n_steps,
            volatility,
            permanent_impact,
            temporary_impact,
            lambda,
        );

        frontier.push((finite_or_zero(cost), finite_or_zero(var)));
    }

    // Sort by expected cost ascending
    frontier.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal));

    frontier
}

/// Register Almgren-Chriss types and functions in the Python module.
pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<ACTrajectory>()?;
    m.add_function(wrap_pyfunction!(ac_kappa, m)?)?;
    m.add_function(wrap_pyfunction!(ac_optimal_schedule, m)?)?;
    m.add_function(wrap_pyfunction!(ac_efficient_frontier, m)?)?;
    Ok(())
}

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

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

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

    #[test]
    fn test_ac_kappa_basic() {
        let k = ac_kappa(0.3, 0.01, 0.05, 1.0).unwrap();
        assert!(k > 0.0);
        assert!(k.is_finite());
    }

    #[test]
    fn test_ac_kappa_higher_risk_aversion() {
        let k1 = ac_kappa(0.3, 0.01, 0.05, 0.1).unwrap();
        let k2 = ac_kappa(0.3, 0.01, 0.05, 10.0).unwrap();
        assert!(k2 > k1, "higher risk aversion should increase kappa");
    }

    #[test]
    fn test_ac_kappa_invalid_inputs() {
        assert!(ac_kappa(-0.3, 0.01, 0.05, 1.0).is_err());
        assert!(ac_kappa(0.3, -0.01, 0.05, 1.0).is_err());
        assert!(ac_kappa(0.3, 0.01, -0.05, 1.0).is_err());
        assert!(ac_kappa(0.3, 0.01, 0.05, -1.0).is_err());
        assert!(ac_kappa(0.3, 0.01, 0.0, 1.0).is_err());
        assert!(ac_kappa(f64::NAN, 0.01, 0.05, 1.0).is_err());
    }

    #[test]
    fn test_ac_schedule_basic() {
        let result = ac_optimal_schedule(10000.0, 10, 0.3, 0.01, 0.05, 1.0).unwrap();
        assert_eq!(result.schedule.len(), 11); // n_steps + 1
        // First element should be total_shares
        assert!((result.schedule[0] - 10000.0).abs() < 1e-6);
        // Last element should be 0 (fully liquidated)
        assert!(result.schedule[10].abs() < 1e-6, "last={}", result.schedule[10]);
        assert!(result.expected_cost > 0.0);
        assert!(result.variance >= 0.0);
    }

    #[test]
    fn test_ac_schedule_monotonically_decreasing() {
        let result = ac_optimal_schedule(10000.0, 20, 0.3, 0.01, 0.05, 1.0).unwrap();
        for i in 1..result.schedule.len() {
            assert!(
                result.schedule[i] <= result.schedule[i - 1] + 1e-10,
                "schedule should be monotonically decreasing at step {}",
                i
            );
        }
    }

    #[test]
    fn test_ac_schedule_high_risk_aversion_is_fast() {
        // High risk aversion: should liquidate faster (more at beginning)
        let slow = ac_optimal_schedule(10000.0, 10, 0.3, 0.01, 0.05, 0.01).unwrap();
        let fast = ac_optimal_schedule(10000.0, 10, 0.3, 0.01, 0.05, 100.0).unwrap();
        // After step 1, fast should have less remaining than slow
        assert!(
            fast.schedule[1] < slow.schedule[1],
            "fast={}, slow={}",
            fast.schedule[1],
            slow.schedule[1]
        );
    }

    #[test]
    fn test_ac_schedule_invalid_inputs() {
        assert!(ac_optimal_schedule(-100.0, 10, 0.3, 0.01, 0.05, 1.0).is_err());
        assert!(ac_optimal_schedule(100.0, 0, 0.3, 0.01, 0.05, 1.0).is_err());
        assert!(ac_optimal_schedule(100.0, 10, -0.3, 0.01, 0.05, 1.0).is_err());
        assert!(ac_optimal_schedule(100.0, 10, 0.3, -0.01, 0.05, 1.0).is_err());
        assert!(ac_optimal_schedule(100.0, 10, 0.3, 0.01, 0.0, 1.0).is_err());
        assert!(ac_optimal_schedule(100.0, 10, 0.3, 0.01, 0.05, 0.0).is_err());
    }

    #[test]
    fn test_ac_efficient_frontier_basic() {
        let frontier = ac_efficient_frontier(10000.0, 10, 0.3, 0.01, 0.05, 10).unwrap();
        assert_eq!(frontier.len(), 10);
        for &(cost, var) in &frontier {
            assert!(cost.is_finite(), "cost not finite");
            assert!(var.is_finite(), "var not finite");
            assert!(cost >= 0.0);
            assert!(var >= 0.0);
        }
    }

    #[test]
    fn test_ac_efficient_frontier_sorted_by_cost() {
        let frontier = ac_efficient_frontier(10000.0, 10, 0.3, 0.01, 0.05, 20).unwrap();
        for i in 1..frontier.len() {
            assert!(
                frontier[i].0 >= frontier[i - 1].0 - 1e-10,
                "frontier should be sorted by cost ascending"
            );
        }
    }

    #[test]
    fn test_ac_efficient_frontier_invalid() {
        assert!(ac_efficient_frontier(0.0, 10, 0.3, 0.01, 0.05, 5).is_err());
        assert!(ac_efficient_frontier(100.0, 0, 0.3, 0.01, 0.05, 5).is_err());
        assert!(ac_efficient_frontier(100.0, 10, 0.3, 0.01, 0.05, 0).is_err());
    }

    #[test]
    fn test_ac_schedule_small_position() {
        // Edge case: very small position
        let result = ac_optimal_schedule(1.0, 5, 0.3, 0.01, 0.05, 1.0).unwrap();
        assert_eq!(result.schedule.len(), 6);
        assert!((result.schedule[0] - 1.0).abs() < 1e-10);
    }

    #[test]
    fn test_ac_cost_vs_variance_tradeoff() {
        // Low risk aversion: lower cost, higher variance
        let low = ac_optimal_schedule(10000.0, 10, 0.3, 0.01, 0.05, 0.001).unwrap();
        // High risk aversion: higher cost, lower variance
        let high = ac_optimal_schedule(10000.0, 10, 0.3, 0.01, 0.05, 100.0).unwrap();
        assert!(
            high.expected_cost >= low.expected_cost - 1e-6,
            "higher risk aversion should have higher or equal cost"
        );
        assert!(
            high.variance <= low.variance + 1e-6,
            "higher risk aversion should have lower or equal variance"
        );
    }
}
