//! Optimal Stopping via Longstaff-Schwartz (LSM) algorithm.
//!
//! Implements American-style optimal exercise/exit decisions using
//! backward induction with polynomial basis regression:
//!
//! 1. Generate paths (or accept user-provided paths)
//! 2. At each time step (backward), regress continuation value on
//!    basis functions of current price
//! 3. Compare intrinsic value vs. estimated continuation value
//! 4. Exercise if intrinsic > continuation
//!
//! Includes GBM path generation with Xoshiro256++ PRNG and Box-Muller.
//!
//! All math from scratch in Rust -- no external crate dependencies beyond PyO3.
//!
//! Pro tier required.

use pyo3::prelude::*;

// ===========================================================================
// Constants
// ===========================================================================

const PI: f64 = std::f64::consts::PI;

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

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

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

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

impl Xoshiro256 {
    fn new(seed: u64) -> Self {
        // SplitMix64 to expand seed
        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
    }

    /// Box-Muller transform: return pair of independent standard normals.
    fn next_normal_pair(&mut self) -> (f64, f64) {
        loop {
            let u1 = self.next_f64();
            let u2 = self.next_f64();
            if u1 > 1e-30 {
                let r = (-2.0 * u1.ln()).sqrt();
                let theta = 2.0 * PI * u2;
                return (r * theta.cos(), r * theta.sin());
            }
        }
    }
}

// ===========================================================================
// Polynomial basis functions (Laguerre-like)
// ===========================================================================

/// Evaluate polynomial basis functions for LSM regression.
/// Uses simple polynomial basis: {1, x, x^2, ..., x^{n_basis-1}}
/// normalized by a scaling factor for numerical stability.
#[allow(dead_code)]
fn evaluate_basis(price: f64, n_basis: usize, scale: f64) -> Vec<f64> {
    let x = if scale > 1e-15 { price / scale } else { price };
    let mut basis = Vec::with_capacity(n_basis);
    let mut xpow = 1.0;
    for _ in 0..n_basis {
        basis.push(finite_or_zero(xpow));
        xpow *= x;
        // Prevent overflow for high powers
        if !xpow.is_finite() {
            xpow = 0.0;
        }
    }
    basis
}

/// Laguerre polynomial basis (alternative):
/// L_0(x) = 1
/// L_1(x) = 1 - x
/// L_2(x) = 1 - 2x + x^2/2
/// Weighted by exp(-x/2) for stability.
fn evaluate_laguerre_basis(price: f64, n_basis: usize, scale: f64) -> Vec<f64> {
    let x = if scale > 1e-15 { price / scale } else { price };
    let x = x.max(0.0); // Laguerre defined for x >= 0
    let w = (-x / 2.0).exp();

    let mut basis = Vec::with_capacity(n_basis);
    if n_basis >= 1 {
        basis.push(finite_or_zero(w));
    }
    if n_basis >= 2 {
        basis.push(finite_or_zero(w * (1.0 - x)));
    }
    if n_basis >= 3 {
        basis.push(finite_or_zero(w * (1.0 - 2.0 * x + 0.5 * x * x)));
    }
    // Higher orders: simple polynomial
    let mut xpow = x * x * x;
    for _ in 3..n_basis {
        basis.push(finite_or_zero(w * xpow));
        xpow *= x;
        if !xpow.is_finite() {
            xpow = 0.0;
        }
    }
    basis
}

// ===========================================================================
// Least-squares regression
// ===========================================================================

/// Ordinary least squares: solve X^T X beta = X^T y.
/// X is (n x p), y is (n,). Returns beta (p,).
fn ols_solve(x: &[Vec<f64>], y: &[f64]) -> Vec<f64> {
    let n = x.len();
    let p = if n > 0 { x[0].len() } else { return vec![]; };

    // X^T X (p x p)
    let mut xtx = vec![vec![0.0; p]; p];
    for i in 0..n {
        for j in 0..p {
            for k in 0..=j {
                let val = x[i][j] * x[i][k];
                xtx[j][k] += val;
                if j != k {
                    xtx[k][j] += val;
                }
            }
        }
    }

    // Regularize diagonal
    for j in 0..p {
        xtx[j][j] += 1e-8;
    }

    // X^T y (p,)
    let mut xty = vec![0.0; p];
    for i in 0..n {
        for j in 0..p {
            xty[j] += x[i][j] * y[i];
        }
    }

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

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

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

    beta
}

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

/// Result of Longstaff-Schwartz optimal stopping.
#[pyclass(frozen)]
#[derive(Debug, Clone)]
pub struct StopPolicy {
    /// Exercise boundary at each time step (intrinsic value threshold).
    #[pyo3(get)]
    pub exercise_boundary: Vec<f64>,
    /// Expected payoff (discounted) under the optimal stopping policy.
    #[pyo3(get)]
    pub expected_payoff: f64,
    /// Expected optimal stopping time (as fraction of total horizon).
    #[pyo3(get)]
    pub optimal_stop_time: f64,
}

#[pymethods]
impl StopPolicy {
    fn __repr__(&self) -> String {
        format!(
            "StopPolicy(n_steps={}, payoff={:.6}, stop_time={:.4})",
            self.exercise_boundary.len(),
            self.expected_payoff,
            self.optimal_stop_time
        )
    }
}

// ===========================================================================
// Longstaff-Schwartz algorithm
// ===========================================================================

/// Run the Longstaff-Schwartz algorithm for optimal stopping.
///
/// paths: Vec<Vec<f64>>, paths[i][t] = price of path i at time t
/// strike: the reference level (for put-like payoff: max(strike - S, 0))
/// n_basis: number of basis functions for regression
fn lsm_impl(paths: &[Vec<f64>], strike: f64, n_basis: usize) -> StopPolicy {
    let n_paths = paths.len();
    let n_steps = if n_paths > 0 { paths[0].len() } else { return empty_policy(); };
    if n_steps < 2 {
        return empty_policy();
    }

    // Intrinsic value: max(strike - S, 0) for put; max(S - strike, 0) for call
    // We use put-like payoff (useful for stop-loss / exit decisions)
    let intrinsic = |price: f64| -> f64 { (strike - price).max(0.0) };

    // Discount factor (assume 0 risk-free rate for prediction markets)
    let discount = 1.0;

    // Scaling factor for basis functions
    let avg_price: f64 = paths.iter().map(|p| p[0]).sum::<f64>() / n_paths as f64;
    let scale = avg_price.max(1.0);

    // Cash flow matrix: cf[i] = discounted payoff for path i
    let mut cashflow = vec![0.0_f64; n_paths];
    let mut stop_time = vec![n_steps as f64 - 1.0; n_paths]; // when each path exercises

    // Initialize with terminal payoff
    for i in 0..n_paths {
        cashflow[i] = intrinsic(paths[i][n_steps - 1]);
    }

    // Exercise boundary at each step
    let mut exercise_boundary = vec![0.0_f64; n_steps];
    exercise_boundary[n_steps - 1] = strike;

    // Backward induction
    for t in (1..n_steps - 1).rev() {
        // Collect in-the-money paths
        let mut itm_indices = Vec::new();
        let mut itm_x = Vec::new();
        let mut itm_y = Vec::new();

        for i in 0..n_paths {
            let iv = intrinsic(paths[i][t]);
            if iv > 0.0 {
                itm_indices.push(i);
                itm_x.push(evaluate_laguerre_basis(paths[i][t], n_basis, scale));
                itm_y.push(cashflow[i] * discount); // continuation value (discounted)
            }
        }

        if itm_indices.len() < n_basis + 1 {
            // Not enough in-the-money paths for regression
            exercise_boundary[t] = strike;
            continue;
        }

        // Regression: continuation value = f(basis(S_t))
        let beta = ols_solve(&itm_x, &itm_y);

        // Decision: exercise if intrinsic > continuation estimate
        let mut boundary_sum = 0.0;
        let mut boundary_count = 0usize;

        for (idx_in_itm, &path_idx) in itm_indices.iter().enumerate() {
            let basis = &itm_x[idx_in_itm];
            let continuation_est: f64 = beta
                .iter()
                .zip(basis.iter())
                .map(|(&b, &x)| b * x)
                .sum();
            let continuation_est = finite_or_zero(continuation_est);

            let iv = intrinsic(paths[path_idx][t]);

            if iv > continuation_est {
                // Exercise at time t
                cashflow[path_idx] = iv;
                stop_time[path_idx] = t as f64;
                boundary_sum += paths[path_idx][t];
                boundary_count += 1;
            }
            // Otherwise, keep the future cashflow (already set)
        }

        // Exercise boundary: average price at which exercise happens
        if boundary_count > 0 {
            exercise_boundary[t] = finite_or_zero(boundary_sum / boundary_count as f64);
        } else {
            exercise_boundary[t] = strike;
        }
    }

    // Also check exercise at t=0
    exercise_boundary[0] = strike;

    // Expected payoff
    let total_payoff: f64 = cashflow.iter().sum();
    let expected_payoff = finite_or_zero(total_payoff / n_paths as f64);

    // Expected stop time (as fraction of horizon)
    let total_stop: f64 = stop_time.iter().sum();
    let avg_stop = finite_or_zero(total_stop / n_paths as f64 / (n_steps - 1) as f64);

    StopPolicy {
        exercise_boundary,
        expected_payoff,
        optimal_stop_time: avg_stop,
    }
}

fn empty_policy() -> StopPolicy {
    StopPolicy {
        exercise_boundary: vec![],
        expected_payoff: 0.0,
        optimal_stop_time: 0.0,
    }
}

// ===========================================================================
// GBM path generation
// ===========================================================================

/// Generate Geometric Brownian Motion paths.
///
/// S_{t+1} = S_t * exp((mu - sigma^2/2)*dt + sigma*sqrt(dt)*Z)
fn generate_gbm_impl(
    s0: f64,
    mu: f64,
    sigma: f64,
    dt: f64,
    n_steps: usize,
    n_paths: usize,
    seed: u64,
) -> Vec<Vec<f64>> {
    let mut rng = Xoshiro256::new(seed);
    let drift = (mu - 0.5 * sigma * sigma) * dt;
    let vol = sigma * dt.sqrt();

    let mut paths = Vec::with_capacity(n_paths);

    for _ in 0..n_paths {
        let mut path = Vec::with_capacity(n_steps);
        let mut s = s0;
        path.push(s);

        for _ in 1..n_steps {
            let (z, _) = rng.next_normal_pair();
            s *= (drift + vol * z).exp();
            s = finite_or_zero(s).max(1e-20);
            path.push(s);
        }

        paths.push(path);
    }

    paths
}

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

/// Run Longstaff-Schwartz optimal stopping on price paths.
///
/// Determines when to optimally exit a position by comparing
/// intrinsic value (distance from strike) against estimated
/// continuation value from polynomial regression.
///
/// Parameters
/// ----------
/// paths : list[list[float]]
///     Price paths (n_paths x n_steps).
/// strike : float
///     Strike/reference price for the exit decision.
/// n_basis : int
///     Number of polynomial basis functions (default 4).
///
/// Returns
/// -------
/// StopPolicy
#[pyfunction]
#[pyo3(signature = (paths, strike, n_basis=4))]
fn longstaff_schwartz(paths: Vec<Vec<f64>>, strike: f64, n_basis: usize) -> PyResult<StopPolicy> {
    crate::auth::require_pro()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

    if paths.is_empty() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "paths must not be empty",
        ));
    }
    let n_steps = paths[0].len();
    if n_steps < 3 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "paths must have at least 3 time steps",
        ));
    }
    if paths.len() < 10 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "need at least 10 paths",
        ));
    }
    for (i, path) in paths.iter().enumerate() {
        if path.len() != n_steps {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "path {} has {} steps, expected {}", i, path.len(), n_steps
            )));
        }
        for (t, &v) in path.iter().enumerate() {
            if !v.is_finite() || v <= 0.0 {
                return Err(pyo3::exceptions::PyValueError::new_err(format!(
                    "path {}[{}] must be positive and finite (got {})", i, t, v
                )));
            }
        }
    }
    if !strike.is_finite() || strike <= 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "strike must be positive and finite",
        ));
    }
    if n_basis < 2 || n_basis > 20 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "n_basis must be in [2, 20]",
        ));
    }

    Ok(lsm_impl(&paths, strike, n_basis))
}

/// Check if the current price is above the exercise boundary.
///
/// Parameters
/// ----------
/// price : float
///     Current price.
/// boundary : list[float]
///     Exercise boundary from StopPolicy.
/// time_step : int
///     Current time step index.
///
/// Returns
/// -------
/// bool
///     True if position should be exited (price below boundary for put-like).
#[pyfunction]
#[pyo3(signature = (price, boundary, time_step))]
fn should_exit(price: f64, boundary: Vec<f64>, time_step: usize) -> bool {
    if time_step >= boundary.len() {
        return true; // Past the horizon, exit
    }
    // Exit if price is below the exercise boundary (put-like: exit on drawdown)
    price <= boundary[time_step]
}

/// Generate Geometric Brownian Motion price paths.
///
/// S_{t+1} = S_t * exp((mu - sigma^2/2)*dt + sigma*sqrt(dt)*Z)
///
/// Parameters
/// ----------
/// s0 : float
///     Initial price.
/// mu : float
///     Drift (annualized).
/// sigma : float
///     Volatility (annualized).
/// dt : float
///     Time step size (in years, e.g. 1/252 for daily).
/// n_steps : int
///     Number of time steps per path.
/// n_paths : int
///     Number of paths to generate.
/// seed : int
///     Random seed.
///
/// Returns
/// -------
/// list[list[float]]
///     Price paths (n_paths x n_steps).
#[pyfunction]
#[pyo3(signature = (s0, mu, sigma, dt, n_steps, n_paths, seed))]
fn generate_gbm_paths(
    s0: f64,
    mu: f64,
    sigma: f64,
    dt: f64,
    n_steps: usize,
    n_paths: usize,
    seed: u64,
) -> PyResult<Vec<Vec<f64>>> {
    crate::auth::require_pro()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

    if !s0.is_finite() || s0 <= 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "s0 must be positive and finite",
        ));
    }
    if !mu.is_finite() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "mu must be finite",
        ));
    }
    if !sigma.is_finite() || sigma <= 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "sigma must be positive and finite",
        ));
    }
    if !dt.is_finite() || dt <= 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "dt must be positive and finite",
        ));
    }
    if n_steps < 2 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "n_steps must be at least 2",
        ));
    }
    if n_paths == 0 {
        return Ok(vec![]);
    }
    if n_paths > 1_000_000 || n_steps > 100_000 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "n_paths * n_steps too large (max 1M paths, 100K steps)",
        ));
    }

    Ok(generate_gbm_impl(s0, mu, sigma, dt, n_steps, n_paths, seed))
}

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

pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<StopPolicy>()?;
    m.add_function(wrap_pyfunction!(longstaff_schwartz, m)?)?;
    m.add_function(wrap_pyfunction!(should_exit, m)?)?;
    m.add_function(wrap_pyfunction!(generate_gbm_paths, m)?)?;
    Ok(())
}

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

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

    #[test]
    fn test_xoshiro_deterministic() {
        let mut rng1 = Xoshiro256::new(42);
        let mut rng2 = Xoshiro256::new(42);
        for _ in 0..100 {
            assert_eq!(rng1.next_u64(), rng2.next_u64());
        }
    }

    #[test]
    fn test_xoshiro_different_seeds() {
        let mut rng1 = Xoshiro256::new(42);
        let mut rng2 = Xoshiro256::new(43);
        let mut all_same = true;
        for _ in 0..10 {
            if rng1.next_u64() != rng2.next_u64() {
                all_same = false;
                break;
            }
        }
        assert!(!all_same);
    }

    #[test]
    fn test_gbm_basic() {
        let paths = generate_gbm_impl(100.0, 0.05, 0.2, 1.0 / 252.0, 50, 100, 42);
        assert_eq!(paths.len(), 100);
        assert_eq!(paths[0].len(), 50);
        // All paths start at s0
        for p in &paths {
            assert!((p[0] - 100.0).abs() < 1e-10);
        }
        // All values positive
        for p in &paths {
            for &v in p {
                assert!(v > 0.0);
            }
        }
    }

    #[test]
    fn test_evaluate_basis() {
        let basis = evaluate_basis(2.0, 4, 1.0);
        assert_eq!(basis.len(), 4);
        assert!((basis[0] - 1.0).abs() < 1e-10); // 1
        assert!((basis[1] - 2.0).abs() < 1e-10); // x
        assert!((basis[2] - 4.0).abs() < 1e-10); // x^2
        assert!((basis[3] - 8.0).abs() < 1e-10); // x^3
    }

    #[test]
    fn test_lsm_basic() {
        // Simple test: paths clearly in/out of the money
        let paths = generate_gbm_impl(100.0, -0.1, 0.3, 1.0 / 252.0, 20, 500, 42);
        let policy = lsm_impl(&paths, 100.0, 4);
        assert_eq!(policy.exercise_boundary.len(), 20);
        assert!(policy.expected_payoff >= 0.0);
        assert!(policy.optimal_stop_time >= 0.0 && policy.optimal_stop_time <= 1.0);
    }

    #[test]
    fn test_should_exit_boundary() {
        let boundary = vec![90.0, 85.0, 80.0, 75.0];
        assert!(!should_exit(95.0, boundary.clone(), 0)); // above boundary
        assert!(should_exit(85.0, boundary.clone(), 1));  // at boundary
        assert!(should_exit(70.0, boundary.clone(), 2));  // below boundary
        assert!(should_exit(100.0, vec![], 0));           // past horizon
    }

    #[test]
    fn test_ols_solve_simple() {
        // y = 2*x + 1, data: (0,1), (1,3), (2,5)
        let x = vec![vec![1.0, 0.0], vec![1.0, 1.0], vec![1.0, 2.0]];
        let y = vec![1.0, 3.0, 5.0];
        let beta = ols_solve(&x, &y);
        assert!((beta[0] - 1.0).abs() < 0.1); // intercept
        assert!((beta[1] - 2.0).abs() < 0.1); // slope
    }

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