//! Characteristic Function Pricing for binary/digital options.
//!
//! Prices binary options under Heston, Merton Jump-Diffusion, and Variance Gamma
//! models using Gil-Pelaez inversion of the characteristic function.
//! Implements a radix-2 Cooley-Tukey FFT from scratch.

use pyo3::prelude::*;

// ===========================================================================
// Complex number (from scratch)
// ===========================================================================

#[derive(Debug, Clone, Copy)]
struct Complex {
    re: f64,
    im: f64,
}

impl Complex {
    fn new(re: f64, im: f64) -> Self {
        Self { re, im }
    }

    fn zero() -> Self {
        Self { re: 0.0, im: 0.0 }
    }

    fn one() -> Self {
        Self { re: 1.0, im: 0.0 }
    }

    fn i() -> Self {
        Self { re: 0.0, im: 1.0 }
    }

    fn norm_sq(&self) -> f64 {
        self.re * self.re + self.im * self.im
    }

    fn abs(&self) -> f64 {
        self.norm_sq().sqrt()
    }

    fn exp(self) -> Self {
        let e = self.re.exp();
        Self {
            re: e * self.im.cos(),
            im: e * self.im.sin(),
        }
    }

    fn ln(self) -> Self {
        let r = self.abs();
        let theta = self.im.atan2(self.re);
        Self {
            re: r.ln(),
            im: theta,
        }
    }

    fn sqrt(self) -> Self {
        let r = self.abs().sqrt();
        let theta = self.im.atan2(self.re) / 2.0;
        Self {
            re: r * theta.cos(),
            im: r * theta.sin(),
        }
    }

    fn conj(self) -> Self {
        Self {
            re: self.re,
            im: -self.im,
        }
    }
}

impl std::ops::Add for Complex {
    type Output = Self;
    fn add(self, rhs: Self) -> Self {
        Self {
            re: self.re + rhs.re,
            im: self.im + rhs.im,
        }
    }
}

impl std::ops::Sub for Complex {
    type Output = Self;
    fn sub(self, rhs: Self) -> Self {
        Self {
            re: self.re - rhs.re,
            im: self.im - rhs.im,
        }
    }
}

impl std::ops::Mul for Complex {
    type Output = Self;
    fn mul(self, rhs: Self) -> Self {
        Self {
            re: self.re * rhs.re - self.im * rhs.im,
            im: self.re * rhs.im + self.im * rhs.re,
        }
    }
}

impl std::ops::Div for Complex {
    type Output = Self;
    fn div(self, rhs: Self) -> Self {
        let d = rhs.norm_sq();
        if d < 1e-30 {
            return Complex::zero();
        }
        Self {
            re: (self.re * rhs.re + self.im * rhs.im) / d,
            im: (self.im * rhs.re - self.re * rhs.im) / d,
        }
    }
}

impl std::ops::Mul<f64> for Complex {
    type Output = Self;
    fn mul(self, rhs: f64) -> Self {
        Self {
            re: self.re * rhs,
            im: self.im * rhs,
        }
    }
}

impl std::ops::Mul<Complex> for f64 {
    type Output = Complex;
    fn mul(self, rhs: Complex) -> Complex {
        Complex {
            re: self * rhs.re,
            im: self * rhs.im,
        }
    }
}

impl std::ops::Neg for Complex {
    type Output = Self;
    fn neg(self) -> Self {
        Self {
            re: -self.re,
            im: -self.im,
        }
    }
}

// ===========================================================================
// PyO3 types
// ===========================================================================

#[pyclass(eq, eq_int, frozen, hash)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum PricingModel {
    Heston,
    MertonJump,
    VarianceGamma,
}

#[pymethods]
impl PricingModel {
    fn __repr__(&self) -> String {
        match self {
            Self::Heston => "PricingModel.Heston".to_string(),
            Self::MertonJump => "PricingModel.MertonJump".to_string(),
            Self::VarianceGamma => "PricingModel.VarianceGamma".to_string(),
        }
    }
}

#[pyclass(frozen)]
#[derive(Debug, Clone)]
pub struct HestonParams {
    #[pyo3(get)]
    pub v0: f64,
    #[pyo3(get)]
    pub kappa: f64,
    #[pyo3(get)]
    pub theta: f64,
    #[pyo3(get)]
    pub sigma_v: f64,
    #[pyo3(get)]
    pub rho: f64,
}

#[pymethods]
impl HestonParams {
    #[new]
    #[pyo3(signature = (v0=0.04, kappa=2.0, theta=0.04, sigma_v=0.3, rho=-0.7))]
    fn new(v0: f64, kappa: f64, theta: f64, sigma_v: f64, rho: f64) -> PyResult<Self> {
        if v0 <= 0.0 || theta <= 0.0 || sigma_v <= 0.0 || kappa <= 0.0 {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "v0, kappa, theta, sigma_v must be positive",
            ));
        }
        if rho < -1.0 || rho > 1.0 {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "rho must be in [-1, 1]",
            ));
        }
        Ok(Self { v0, kappa, theta, sigma_v, rho })
    }

    fn __repr__(&self) -> String {
        format!(
            "HestonParams(v0={:.4}, kappa={:.2}, theta={:.4}, sigma_v={:.2}, rho={:.2})",
            self.v0, self.kappa, self.theta, self.sigma_v, self.rho
        )
    }
}

#[pyclass(frozen)]
#[derive(Debug, Clone)]
pub struct MertonParams {
    #[pyo3(get)]
    pub jump_intensity: f64,
    #[pyo3(get)]
    pub jump_mean: f64,
    #[pyo3(get)]
    pub jump_vol: f64,
}

#[pymethods]
impl MertonParams {
    #[new]
    #[pyo3(signature = (jump_intensity=0.1, jump_mean=-0.1, jump_vol=0.2))]
    fn new(jump_intensity: f64, jump_mean: f64, jump_vol: f64) -> PyResult<Self> {
        if jump_intensity < 0.0 || jump_vol < 0.0 {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "jump_intensity and jump_vol must be non-negative",
            ));
        }
        Ok(Self { jump_intensity, jump_mean, jump_vol })
    }

    fn __repr__(&self) -> String {
        format!(
            "MertonParams(lambda={:.2}, mean={:.2}, vol={:.2})",
            self.jump_intensity, self.jump_mean, self.jump_vol
        )
    }
}

#[pyclass(frozen)]
#[derive(Debug, Clone)]
pub struct VGParams {
    #[pyo3(get)]
    pub sigma: f64,
    #[pyo3(get)]
    pub theta: f64,
    #[pyo3(get)]
    pub nu: f64,
}

#[pymethods]
impl VGParams {
    #[new]
    #[pyo3(signature = (sigma=0.2, theta=-0.1, nu=0.5))]
    fn new(sigma: f64, theta: f64, nu: f64) -> PyResult<Self> {
        if sigma <= 0.0 || nu <= 0.0 {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "sigma and nu must be positive",
            ));
        }
        Ok(Self { sigma, theta, nu })
    }

    fn __repr__(&self) -> String {
        format!(
            "VGParams(sigma={:.2}, theta={:.2}, nu={:.2})",
            self.sigma, self.theta, self.nu
        )
    }
}

#[pyclass(frozen)]
#[derive(Debug, Clone)]
pub struct BinaryPrice {
    #[pyo3(get)]
    pub price: f64,
    #[pyo3(get)]
    pub delta: f64,
    #[pyo3(get)]
    pub vega: f64,
    #[pyo3(get)]
    pub model: PricingModel,
}

#[pymethods]
impl BinaryPrice {
    fn __repr__(&self) -> String {
        format!(
            "BinaryPrice(price={:.6}, delta={:.6}, vega={:.6}, model={:?})",
            self.price, self.delta, self.vega, self.model
        )
    }
}

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

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

/// Standard normal CDF.
fn normal_cdf(x: f64) -> f64 {
    const A1: f64 = 0.254829592;
    const A2: f64 = -0.284496736;
    const A3: f64 = 1.421413741;
    const A4: f64 = -1.453152027;
    const A5: f64 = 1.061405429;
    const P: f64 = 0.3275911;

    let sign = if x < 0.0 { -1.0 } else { 1.0 };
    let x_abs = x.abs() / std::f64::consts::SQRT_2;
    let t = 1.0 / (1.0 + P * x_abs);
    let y = 1.0 - (((((A5 * t + A4) * t) + A3) * t + A2) * t + A1) * t * (-x_abs * x_abs).exp();
    0.5 * (1.0 + sign * y)
}

/// Heston characteristic function.
/// phi_Heston(u) for log-price
fn heston_cf(u: Complex, t: f64, params: &HestonParams, r: f64) -> Complex {
    let i = Complex::i();
    let iu = i * u;

    let sigma2 = params.sigma_v * params.sigma_v;
    let rho_sigma = params.rho * params.sigma_v;

    // d = sqrt((rho*sigma*i*u - kappa)^2 + sigma^2*(i*u + u^2))
    let a = Complex::new(rho_sigma, 0.0) * iu - Complex::new(params.kappa, 0.0);
    let b = Complex::new(sigma2, 0.0) * (iu + u * u);
    let d = (a * a + b).sqrt();

    // g = (kappa - rho*sigma*i*u - d) / (kappa - rho*sigma*i*u + d)
    let kappa_c = Complex::new(params.kappa, 0.0);
    let rsi = Complex::new(rho_sigma, 0.0) * iu;
    let g_num = kappa_c - rsi - d;
    let g_den = kappa_c - rsi + d;
    let g = g_num / g_den;

    let edt = (Complex::new(-1.0, 0.0) * d * t).exp();

    // C = r*i*u*t + (kappa*theta/sigma^2) * ((kappa - rho*sigma*i*u - d)*t - 2*ln((1-g*exp(-d*t))/(1-g)))
    let one_minus_gedt = Complex::one() - g * edt;
    let one_minus_g = Complex::one() - g;
    let log_ratio = (one_minus_gedt / one_minus_g).ln();

    let kt_s2 = params.kappa * params.theta / sigma2;
    let c = iu * r * t + Complex::new(kt_s2, 0.0) * (g_num * t - Complex::new(2.0, 0.0) * log_ratio);

    // D = ((kappa - rho*sigma*i*u - d)/sigma^2) * ((1-exp(-d*t))/(1-g*exp(-d*t)))
    let dd = (g_num / Complex::new(sigma2, 0.0)) * ((Complex::one() - edt) / one_minus_gedt);

    (c + dd * params.v0).exp()
}

/// Merton jump-diffusion characteristic function.
fn merton_cf(u: Complex, t: f64, vol: f64, params: &MertonParams, r: f64) -> Complex {
    let i = Complex::i();
    let iu = i * u;

    // GBM part: exp(i*u*(r - 0.5*sigma^2)*t - 0.5*sigma^2*u^2*t)
    let drift = r - 0.5 * vol * vol
        - params.jump_intensity * ((params.jump_mean + 0.5 * params.jump_vol * params.jump_vol).exp() - 1.0);
    let gbm = (iu * drift * t - Complex::new(0.5 * vol * vol, 0.0) * u * u * t).exp();

    // Jump part: exp(lambda*t * (exp(i*u*mu_j - 0.5*sigma_j^2*u^2) - 1))
    let jump_cf = (iu * params.jump_mean - Complex::new(0.5 * params.jump_vol * params.jump_vol, 0.0) * u * u).exp();
    let jump_part = (Complex::new(params.jump_intensity * t, 0.0) * (jump_cf - Complex::one())).exp();

    gbm * jump_part
}

/// Variance Gamma characteristic function.
fn vg_cf(u: Complex, t: f64, params: &VGParams, r: f64) -> Complex {
    let i = Complex::i();
    let iu = i * u;

    // omega = (1/nu) * ln(1 - theta*nu - 0.5*sigma^2*nu)
    // VG existence condition: 1 - theta*nu - 0.5*sigma^2*nu > 0
    let vg_arg = 1.0 - params.theta * params.nu - 0.5 * params.sigma * params.sigma * params.nu;
    let omega = if vg_arg <= 0.0 {
        // VG process does not exist for these parameters — fall back to 0.0 (BSM-like drift)
        0.0
    } else {
        (1.0 / params.nu) * vg_arg.ln()
    };

    // drift term
    let drift_term = (iu * (r + omega) * t).exp();

    // VG cf: (1 - i*u*theta*nu + 0.5*sigma^2*nu*u^2)^(-t/nu)
    let inner = Complex::one() - iu * params.theta * params.nu
        + Complex::new(0.5 * params.sigma * params.sigma * params.nu, 0.0) * u * u;
    let power = Complex::new(-t / params.nu, 0.0);
    let vg_part = (inner.ln() * power).exp();

    drift_term * vg_part
}

/// Gil-Pelaez inversion: P(S_T > K) = 0.5 + (1/pi) * integral
fn gil_pelaez_binary_price<F>(cf: F, s0: f64, k: f64, t: f64, r: f64, n_points: usize) -> f64
where
    F: Fn(Complex) -> Complex,
{
    let log_moneyness = (s0 / k).ln();
    let du = 0.01;

    let mut integral = 0.0;
    for j in 1..=n_points {
        let u = j as f64 * du;
        let u_c = Complex::new(u, 0.0);
        let cf_val = cf(u_c);
        let exp_part = (Complex::new(0.0, -u * log_moneyness)).exp();
        let integrand = (cf_val * exp_part / Complex::new(0.0, u)).re;
        integral += integrand * du;
    }

    let prob = 0.5 + integral / std::f64::consts::PI;
    let price = (-r * t).exp() * prob;
    finite_or_zero(price.clamp(0.0, 1.0))
}

/// Numerical delta via bump
fn numerical_delta<F>(price_fn: &F, s0: f64, bump: f64) -> f64
where
    F: Fn(f64) -> f64,
{
    let p_up = price_fn(s0 * (1.0 + bump));
    let p_dn = price_fn(s0 * (1.0 - bump));
    finite_or_zero((p_up - p_dn) / (2.0 * s0 * bump))
}

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

/// Price a binary option under the Heston stochastic volatility model.
#[pyfunction]
#[pyo3(signature = (strike, time_to_expiry, params, r=0.0, s0=1.0))]
pub fn binary_price_heston(
    strike: f64,
    time_to_expiry: f64,
    params: &HestonParams,
    r: f64,
    s0: f64,
) -> PyResult<BinaryPrice> {
    crate::auth::require_pro()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

    if strike <= 0.0 || time_to_expiry <= 0.0 || s0 <= 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "strike, time_to_expiry, and s0 must be positive",
        ));
    }

    let price_at = |spot: f64| -> f64 {
        gil_pelaez_binary_price(
            |u| heston_cf(u, time_to_expiry, params, r),
            spot, strike, time_to_expiry, r, 500,
        )
    };

    let price = price_at(s0);
    let delta = numerical_delta(&price_at, s0, 0.001);

    // Vega: bump v0
    let params_up = HestonParams {
        v0: params.v0 * 1.01,
        ..*params
    };
    let price_up = gil_pelaez_binary_price(
        |u| heston_cf(u, time_to_expiry, &params_up, r),
        s0, strike, time_to_expiry, r, 500,
    );
    let vega = finite_or_zero((price_up - price) / (params.v0 * 0.01));

    Ok(BinaryPrice {
        price,
        delta,
        vega,
        model: PricingModel::Heston,
    })
}

/// Price a binary option under the Merton jump-diffusion model.
#[pyfunction]
#[pyo3(signature = (strike, time_to_expiry, vol, params, r=0.0, s0=1.0))]
pub fn binary_price_merton(
    strike: f64,
    time_to_expiry: f64,
    vol: f64,
    params: &MertonParams,
    r: f64,
    s0: f64,
) -> PyResult<BinaryPrice> {
    crate::auth::require_pro()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

    if strike <= 0.0 || time_to_expiry <= 0.0 || s0 <= 0.0 || vol <= 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "strike, time_to_expiry, s0, and vol must be positive",
        ));
    }

    let price_at = |spot: f64| -> f64 {
        gil_pelaez_binary_price(
            |u| merton_cf(u, time_to_expiry, vol, params, r),
            spot, strike, time_to_expiry, r, 500,
        )
    };

    let price = price_at(s0);
    let delta = numerical_delta(&price_at, s0, 0.001);

    // Vega: bump vol
    let price_up = gil_pelaez_binary_price(
        |u| merton_cf(u, time_to_expiry, vol * 1.01, params, r),
        s0, strike, time_to_expiry, r, 500,
    );
    let vega = finite_or_zero((price_up - price) / (vol * 0.01));

    Ok(BinaryPrice {
        price,
        delta,
        vega,
        model: PricingModel::MertonJump,
    })
}

/// Price a binary option under the Variance Gamma model.
#[pyfunction]
#[pyo3(signature = (strike, time_to_expiry, params, r=0.0, s0=1.0))]
pub fn binary_price_vg(
    strike: f64,
    time_to_expiry: f64,
    params: &VGParams,
    r: f64,
    s0: f64,
) -> PyResult<BinaryPrice> {
    crate::auth::require_pro()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

    if strike <= 0.0 || time_to_expiry <= 0.0 || s0 <= 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "strike, time_to_expiry, and s0 must be positive",
        ));
    }

    let price_at = |spot: f64| -> f64 {
        gil_pelaez_binary_price(
            |u| vg_cf(u, time_to_expiry, params, r),
            spot, strike, time_to_expiry, r, 500,
        )
    };

    let price = price_at(s0);
    let delta = numerical_delta(&price_at, s0, 0.001);

    // Vega: bump sigma
    let params_up = VGParams {
        sigma: params.sigma * 1.01,
        ..*params
    };
    let price_up = gil_pelaez_binary_price(
        |u| vg_cf(u, time_to_expiry, &params_up, r),
        s0, strike, time_to_expiry, r, 500,
    );
    let vega = finite_or_zero((price_up - price) / (params.sigma * 0.01));

    Ok(BinaryPrice {
        price,
        delta,
        vega,
        model: PricingModel::VarianceGamma,
    })
}

/// Implied vol from a binary option price using bisection.
#[pyfunction]
#[pyo3(signature = (price, strike, time_to_expiry, r=0.0, s0=1.0))]
pub fn implied_vol_from_binary(
    price: f64,
    strike: f64,
    time_to_expiry: f64,
    r: f64,
    s0: f64,
) -> PyResult<f64> {
    crate::auth::require_pro()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

    if price <= 0.0 || price >= 1.0 || strike <= 0.0 || time_to_expiry <= 0.0 || s0 <= 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "price must be in (0,1), strike/time/s0 must be positive",
        ));
    }

    // Binary option price under BSM: exp(-r*T) * N(d2)
    // d2 = (ln(S/K) + (r - sigma^2/2)*T) / (sigma*sqrt(T))
    let disc = (-r * time_to_expiry).exp();
    let log_sk = (s0 / strike).ln();
    let sqrt_t = time_to_expiry.sqrt();

    let bsm_binary = |sigma: f64| -> f64 {
        let d2 = (log_sk + (r - 0.5 * sigma * sigma) * time_to_expiry) / (sigma * sqrt_t);
        disc * normal_cdf(d2)
    };

    // Bisection
    let mut lo = 0.001;
    let mut hi = 5.0;

    for _ in 0..100 {
        let mid = (lo + hi) / 2.0;
        let p = bsm_binary(mid);
        if (p - price).abs() < 1e-10 {
            return Ok(mid);
        }
        if price < 0.5 {
            // OTM binary: higher vol → higher price (toward 0.5)
            if p > price { hi = mid; } else { lo = mid; }
        } else {
            // ITM binary: higher vol → lower price (toward 0.5)
            if p > price { lo = mid; } else { hi = mid; }
        }
    }

    Ok(finite_or_zero((lo + hi) / 2.0))
}

// ===========================================================================
// Registration
// ===========================================================================

pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<PricingModel>()?;
    m.add_class::<HestonParams>()?;
    m.add_class::<MertonParams>()?;
    m.add_class::<VGParams>()?;
    m.add_class::<BinaryPrice>()?;
    m.add_function(wrap_pyfunction!(binary_price_heston, m)?)?;
    m.add_function(wrap_pyfunction!(binary_price_merton, m)?)?;
    m.add_function(wrap_pyfunction!(binary_price_vg, m)?)?;
    m.add_function(wrap_pyfunction!(implied_vol_from_binary, m)?)?;
    Ok(())
}
