//! Volatility Signature Plot and Two-Scale Realized Volatility.
//!
//! Computes realized volatility at different sampling frequencies to produce
//! a signature plot. Implements TSRV (Zhang, Mykland, Aït-Sahalia 2005) for
//! noise-robust volatility estimation, and optimal sampling frequency detection.

use pyo3::prelude::*;

// ===========================================================================
// Result type
// ===========================================================================

#[pyclass(frozen)]
#[derive(Debug, Clone)]
pub struct SignaturePlot {
    /// Sampling frequencies (in ticks).
    #[pyo3(get)]
    pub frequencies: Vec<usize>,
    /// Realized volatility at each frequency.
    #[pyo3(get)]
    pub realized_vols: Vec<f64>,
    /// Optimal sampling frequency (minimum variance).
    #[pyo3(get)]
    pub optimal_frequency: usize,
    /// Two-scale realized volatility estimate.
    #[pyo3(get)]
    pub two_scale_rv: f64,
    /// Estimated microstructure noise variance.
    #[pyo3(get)]
    pub noise_variance: f64,
}

#[pymethods]
impl SignaturePlot {
    fn __repr__(&self) -> String {
        format!(
            "SignaturePlot(n_points={}, optimal_freq={}, tsrv={:.6}, noise_var={:.8})",
            self.frequencies.len(),
            self.optimal_frequency,
            self.two_scale_rv,
            self.noise_variance
        )
    }
}

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

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

/// Compute realized variance at a given sampling frequency.
fn realized_var_at_freq(prices: &[f64], freq: usize) -> f64 {
    if prices.len() < freq + 1 || freq == 0 {
        return 0.0;
    }

    let mut rv = 0.0;
    let mut i = 0;
    while i + freq < prices.len() {
        if prices[i] > 0.0 {
            let ret = (prices[i + freq] / prices[i]).ln();
            rv += ret * ret;
        }
        i += freq;
    }
    finite_or_zero(rv)
}

/// Compute realized variance using subsampled (averaged) estimator.
fn realized_var_subsampled(prices: &[f64], freq: usize) -> f64 {
    if prices.len() < freq + 1 || freq == 0 {
        return 0.0;
    }

    let mut total_rv = 0.0;
    let mut count = 0;

    for offset in 0..freq {
        let mut rv = 0.0;
        let mut i = offset;
        while i + freq < prices.len() {
            if prices[i] > 0.0 {
                let ret = (prices[i + freq] / prices[i]).ln();
                rv += ret * ret;
            }
            i += freq;
        }
        total_rv += rv;
        count += 1;
    }

    if count > 0 {
        finite_or_zero(total_rv / count as f64)
    } else {
        0.0
    }
}

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

/// Compute the volatility signature plot.
#[pyfunction]
#[pyo3(signature = (prices, max_freq=None, n_points=20))]
pub fn volatility_signature(
    prices: Vec<f64>,
    max_freq: Option<usize>,
    n_points: usize,
) -> PyResult<SignaturePlot> {
    crate::auth::require_pro()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

    if prices.len() < 10 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "need at least 10 prices for signature plot",
        ));
    }

    let n = prices.len();
    let max_freq = max_freq.unwrap_or(n / 4).max(2).min(n / 2);
    let n_points = n_points.max(3).min(max_freq);

    // Generate log-spaced frequencies
    let log_min = 1.0_f64.ln();
    let log_max = (max_freq as f64).ln();
    let mut frequencies = Vec::with_capacity(n_points);
    for i in 0..n_points {
        let t = i as f64 / (n_points - 1) as f64;
        let freq = (log_min + t * (log_max - log_min)).exp().round() as usize;
        let freq = freq.max(1).min(max_freq);
        if frequencies.is_empty() || *frequencies.last().unwrap() != freq {
            frequencies.push(freq);
        }
    }

    // Compute realized vol at each frequency
    let mut realized_vols = Vec::with_capacity(frequencies.len());
    let mut min_vol = f64::MAX;
    let mut optimal_freq = 1;

    for &freq in &frequencies {
        let rv = realized_var_subsampled(&prices, freq);
        let vol = rv.sqrt();
        realized_vols.push(finite_or_zero(vol));

        if vol > 0.0 && vol < min_vol {
            min_vol = vol;
            optimal_freq = freq;
        }
    }

    // TSRV and noise variance
    let noise_var = noise_variance_estimate_impl(&prices);
    let tsrv = two_scale_rv_impl(&prices, optimal_freq.max(2));

    Ok(SignaturePlot {
        frequencies,
        realized_vols,
        optimal_frequency: optimal_freq,
        two_scale_rv: finite_or_zero(tsrv),
        noise_variance: finite_or_zero(noise_var),
    })
}

/// Two-Scale Realized Volatility (TSRV).
/// Corrects for microstructure noise using fast and slow scale estimators.
#[pyfunction]
#[pyo3(signature = (prices, slow_freq=5))]
pub fn two_scale_realized_vol(prices: Vec<f64>, slow_freq: usize) -> PyResult<f64> {
    crate::auth::require_pro()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

    if prices.len() < 10 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "need at least 10 prices",
        ));
    }

    Ok(finite_or_zero(two_scale_rv_impl(&prices, slow_freq.max(2))))
}

fn two_scale_rv_impl(prices: &[f64], slow_freq: usize) -> f64 {
    let n = prices.len();
    if n < slow_freq + 1 {
        return 0.0;
    }

    // Fast scale: tick-by-tick RV
    let rv_fast = realized_var_at_freq(prices, 1);

    // Slow scale: subsampled RV at slow_freq
    let rv_slow = realized_var_subsampled(prices, slow_freq);

    // TSRV = RV_slow - (n_bar / n) * RV_fast
    // where n_bar = n / slow_freq (average number of returns per subsample)
    let n_bar = (n as f64 - 1.0) / slow_freq as f64;
    let n_fast = (n - 1) as f64;

    let tsrv = rv_slow - (n_bar / n_fast) * rv_fast;
    tsrv.max(0.0).sqrt()
}

/// Estimate microstructure noise variance from tick-by-tick returns.
/// Uses first-order autocovariance: Var(noise) = -Cov(r_t, r_{t-1}).
#[pyfunction]
pub fn noise_variance_estimate(prices: Vec<f64>) -> f64 {
    crate::auth::require_pro()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))
        .ok();
    finite_or_zero(noise_variance_estimate_impl(&prices))
}

fn noise_variance_estimate_impl(prices: &[f64]) -> f64 {
    if prices.len() < 3 {
        return 0.0;
    }

    // Compute log returns
    let mut returns = Vec::with_capacity(prices.len() - 1);
    for i in 1..prices.len() {
        if prices[i - 1] > 0.0 && prices[i] > 0.0 {
            returns.push((prices[i] / prices[i - 1]).ln());
        }
    }

    if returns.len() < 2 {
        return 0.0;
    }

    // First-order autocovariance
    let n = returns.len();
    let mean_r: f64 = returns.iter().sum::<f64>() / n as f64;
    let mut autocov = 0.0;
    for i in 1..n {
        autocov += (returns[i] - mean_r) * (returns[i - 1] - mean_r);
    }
    autocov /= (n - 1) as f64;

    // Noise variance = -autocovariance (under noise model)
    (-autocov).max(0.0)
}

/// Find optimal sampling frequency that minimizes variance.
#[pyfunction]
pub fn optimal_sampling_frequency(prices: Vec<f64>) -> usize {
    crate::auth::require_pro()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))
        .ok();

    if prices.len() < 10 {
        return 1;
    }

    let n = prices.len();
    let noise_var = noise_variance_estimate_impl(&prices);
    if noise_var <= 1e-15 {
        return 1;
    }

    // Optimal frequency ~ (n * noise_var / IQ)^(1/3)
    // where IQ is integrated quarticity, approximated from data
    let rv1 = realized_var_at_freq(&prices, 1);
    let iq_approx = rv1 * rv1 / (n as f64).max(1.0);

    if iq_approx <= 1e-15 {
        return 1;
    }

    let opt = (n as f64 * noise_var / iq_approx).powf(1.0 / 3.0);
    let opt = opt.round() as usize;
    opt.max(1).min(n / 4)
}

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

pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<SignaturePlot>()?;
    m.add_function(wrap_pyfunction!(volatility_signature, m)?)?;
    m.add_function(wrap_pyfunction!(two_scale_realized_vol, m)?)?;
    m.add_function(wrap_pyfunction!(noise_variance_estimate, m)?)?;
    m.add_function(wrap_pyfunction!(optimal_sampling_frequency, m)?)?;
    Ok(())
}
