//! Rust analytics for prediction market backtesting.
//!
//! Provides O(N) computation of:
//! - Calibration curves with Brier score, log-loss, and ECE
//! - Standalone log-loss
//! - Edge decay analysis (edge vs time-to-resolution)

use pyo3::prelude::*;

// ---------------------------------------------------------------------------
// CalibrationResult
// ---------------------------------------------------------------------------

/// Result of calibration curve computation.
#[pyclass]
#[derive(Debug, Clone)]
pub struct CalibrationResult {
    /// Vec of (bin_center, actual_frequency, count) for each bin.
    #[pyo3(get)]
    pub bins: Vec<(f64, f64, u32)>,
    #[pyo3(get)]
    pub brier_score: f64,
    #[pyo3(get)]
    pub log_loss: f64,
    /// Expected Calibration Error.
    #[pyo3(get)]
    pub ece: f64,
}

#[pymethods]
impl CalibrationResult {
    fn __repr__(&self) -> String {
        format!(
            "CalibrationResult(bins={}, brier={:.4}, log_loss={:.4}, ece={:.4})",
            self.bins.len(),
            self.brier_score,
            self.log_loss,
            self.ece
        )
    }
}

// ---------------------------------------------------------------------------
// EdgeDecayResult
// ---------------------------------------------------------------------------

/// Result of edge decay analysis.
#[pyclass]
#[derive(Debug, Clone)]
pub struct EdgeDecayResult {
    /// Vec of (hours_before_resolution, avg_edge) for each bucket.
    #[pyo3(get)]
    pub decay_curve: Vec<(f64, f64)>,
    /// Estimated half-life in hours (time for edge to halve).
    #[pyo3(get)]
    pub half_life_hours: f64,
}

#[pymethods]
impl EdgeDecayResult {
    fn __repr__(&self) -> String {
        format!(
            "EdgeDecayResult(buckets={}, half_life={:.1}h)",
            self.decay_curve.len(),
            self.half_life_hours
        )
    }
}

// ---------------------------------------------------------------------------
// calibration_curve
// ---------------------------------------------------------------------------

/// Bin predictions into n_bins buckets, compute actual frequency per bin.
/// Also computes Brier score, log-loss, and ECE in a single O(N) pass.
///
/// Args:
///     predictions: predicted probabilities (0-1)
///     outcomes: actual outcomes (0.0 or 1.0)
///     n_bins: number of calibration bins
#[pyfunction]
#[pyo3(signature = (predictions, outcomes, n_bins=10))]
pub fn calibration_curve(
    predictions: Vec<f64>,
    outcomes: Vec<f64>,
    n_bins: usize,
) -> PyResult<CalibrationResult> {
    if predictions.len() != outcomes.len() {
        return Err(pyo3::exceptions::PyValueError::new_err(format!(
            "predictions ({}) and outcomes ({}) must have same length",
            predictions.len(),
            outcomes.len()
        )));
    }
    if predictions.is_empty() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "predictions must not be empty",
        ));
    }
    let n_bins = n_bins.max(1);
    let n = predictions.len() as f64;

    // Accumulators for each bin
    let mut bin_pred_sum = vec![0.0f64; n_bins];
    let mut bin_outcome_sum = vec![0.0f64; n_bins];
    let mut bin_count = vec![0u32; n_bins];

    let mut brier_sum = 0.0;
    let mut log_loss_sum = 0.0;

    let eps = 1e-15;

    for (&pred, &outcome) in predictions.iter().zip(outcomes.iter()) {
        let p = pred.clamp(eps, 1.0 - eps);
        let o = outcome.clamp(0.0, 1.0);

        // Brier: (p - o)^2
        brier_sum += (p - o).powi(2);

        // Log-loss: -(o*ln(p) + (1-o)*ln(1-p))
        log_loss_sum += -(o * p.ln() + (1.0 - o) * (1.0 - p).ln());

        // Bin assignment
        let bin_idx = ((p * n_bins as f64).floor() as usize).min(n_bins - 1);
        bin_pred_sum[bin_idx] += p;
        bin_outcome_sum[bin_idx] += o;
        bin_count[bin_idx] += 1;
    }

    let brier_score = brier_sum / n;
    let log_loss = log_loss_sum / n;

    // Build bins and compute ECE
    let mut bins = Vec::with_capacity(n_bins);
    let mut ece = 0.0;

    for i in 0..n_bins {
        let count = bin_count[i];
        let bin_center = (i as f64 + 0.5) / n_bins as f64;

        if count > 0 {
            let avg_pred = bin_pred_sum[i] / count as f64;
            let actual_freq = bin_outcome_sum[i] / count as f64;
            bins.push((bin_center, actual_freq, count));
            // ECE: weighted |avg_pred - actual_freq|
            ece += (count as f64 / n) * (avg_pred - actual_freq).abs();
        } else {
            bins.push((bin_center, 0.0, 0));
        }
    }

    Ok(CalibrationResult {
        bins,
        brier_score,
        log_loss,
        ece,
    })
}

// ---------------------------------------------------------------------------
// log_loss
// ---------------------------------------------------------------------------

/// Compute log-loss: -1/N * sum(y*ln(p) + (1-y)*ln(1-p)).
/// Predictions clamped to [1e-15, 1-1e-15] for numerical stability.
#[pyfunction]
pub fn log_loss(predictions: Vec<f64>, outcomes: Vec<f64>) -> PyResult<f64> {
    if predictions.len() != outcomes.len() {
        return Err(pyo3::exceptions::PyValueError::new_err(format!(
            "predictions ({}) and outcomes ({}) must have same length",
            predictions.len(),
            outcomes.len()
        )));
    }
    if predictions.is_empty() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "predictions must not be empty",
        ));
    }

    let eps = 1e-15;
    let n = predictions.len() as f64;
    let mut sum = 0.0;

    for (&pred, &outcome) in predictions.iter().zip(outcomes.iter()) {
        let p = pred.clamp(eps, 1.0 - eps);
        let o = outcome.clamp(0.0, 1.0);
        sum += -(o * p.ln() + (1.0 - o) * (1.0 - p).ln());
    }

    Ok(sum / n)
}

// ---------------------------------------------------------------------------
// edge_decay
// ---------------------------------------------------------------------------

/// Group trades by hours-to-resolution, compute avg edge per bucket.
/// Edge = outcome - entry_price for each trade.
///
/// Args:
///     entry_prices: price at trade entry
///     outcomes: actual outcomes (0.0 or 1.0)
///     entry_ts: UNIX timestamp of each trade entry
///     resolution_ts: UNIX timestamp of each market resolution
///     n_buckets: number of time buckets
#[pyfunction]
#[pyo3(signature = (entry_prices, outcomes, entry_ts, resolution_ts, n_buckets=20))]
pub fn edge_decay(
    entry_prices: Vec<f64>,
    outcomes: Vec<f64>,
    entry_ts: Vec<f64>,
    resolution_ts: Vec<f64>,
    n_buckets: usize,
) -> PyResult<EdgeDecayResult> {
    let n = entry_prices.len();
    if n != outcomes.len() || n != entry_ts.len() || n != resolution_ts.len() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "all input vectors must have the same length",
        ));
    }
    if n == 0 {
        return Ok(EdgeDecayResult {
            decay_curve: Vec::new(),
            half_life_hours: 0.0,
        });
    }
    let n_buckets = n_buckets.max(1);

    // Compute hours-to-resolution for each trade
    let mut hours_to_res: Vec<f64> = Vec::with_capacity(n);
    let mut edges: Vec<f64> = Vec::with_capacity(n);

    for i in 0..n {
        let hours = (resolution_ts[i] - entry_ts[i]).max(0.0) / 3600.0;
        hours_to_res.push(hours);
        edges.push(outcomes[i] - entry_prices[i]);
    }

    // Find max hours for bucket sizing
    let max_hours = hours_to_res
        .iter()
        .copied()
        .fold(0.0f64, f64::max)
        .max(1.0);
    let bucket_width = max_hours / n_buckets as f64;

    // Accumulate edges per bucket
    let mut bucket_edge_sum = vec![0.0f64; n_buckets];
    let mut bucket_count = vec![0u32; n_buckets];

    for i in 0..n {
        let bucket = ((hours_to_res[i] / bucket_width).floor() as usize).min(n_buckets - 1);
        bucket_edge_sum[bucket] += edges[i];
        bucket_count[bucket] += 1;
    }

    // Build decay curve
    let mut decay_curve = Vec::with_capacity(n_buckets);
    for i in 0..n_buckets {
        let hours_center = (i as f64 + 0.5) * bucket_width;
        let avg_edge = if bucket_count[i] > 0 {
            bucket_edge_sum[i] / bucket_count[i] as f64
        } else {
            0.0
        };
        decay_curve.push((hours_center, avg_edge));
    }

    // Estimate half-life: find where edge drops to half of max edge
    let max_edge = decay_curve
        .iter()
        .map(|(_, e)| e.abs())
        .fold(0.0f64, f64::max);

    let half_life_hours = if max_edge > 0.0 {
        let half_target = max_edge / 2.0;
        // Find first bucket where |edge| drops below half_target
        let mut hl = max_hours; // default: no decay found
        for &(hours, edge) in &decay_curve {
            if edge.abs() <= half_target {
                hl = hours;
                break;
            }
        }
        hl
    } else {
        0.0
    };

    Ok(EdgeDecayResult {
        decay_curve,
        half_life_hours,
    })
}

/// Register analytics types and functions in the Python module.
pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<CalibrationResult>()?;
    m.add_class::<EdgeDecayResult>()?;
    m.add_function(wrap_pyfunction!(calibration_curve, m)?)?;
    m.add_function(wrap_pyfunction!(log_loss, m)?)?;
    m.add_function(wrap_pyfunction!(edge_decay, m)?)?;
    Ok(())
}

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

    #[test]
    fn calibration_perfect() {
        // Perfect predictions: 0.0 for outcome 0, 1.0 for outcome 1
        let preds = vec![0.1, 0.9, 0.1, 0.9];
        let outcomes = vec![0.0, 1.0, 0.0, 1.0];
        let result = calibration_curve(preds, outcomes, 5).unwrap();
        assert!(result.brier_score < 0.1);
        assert!(result.ece < 0.2);
    }

    #[test]
    fn calibration_terrible() {
        // Predictions opposite to outcomes
        let preds = vec![0.9, 0.1, 0.9, 0.1];
        let outcomes = vec![0.0, 1.0, 0.0, 1.0];
        let result = calibration_curve(preds, outcomes, 5).unwrap();
        assert!(result.brier_score > 0.5);
    }

    #[test]
    fn calibration_length_mismatch() {
        let result = calibration_curve(vec![0.5], vec![0.0, 1.0], 5);
        assert!(result.is_err());
    }

    #[test]
    fn calibration_empty() {
        let result = calibration_curve(vec![], vec![], 5);
        assert!(result.is_err());
    }

    #[test]
    fn log_loss_perfect() {
        let preds = vec![0.99, 0.01];
        let outcomes = vec![1.0, 0.0];
        let ll = log_loss(preds, outcomes).unwrap();
        assert!(ll < 0.05, "log_loss={}", ll);
    }

    #[test]
    fn log_loss_terrible() {
        let preds = vec![0.01, 0.99];
        let outcomes = vec![1.0, 0.0];
        let ll = log_loss(preds, outcomes).unwrap();
        assert!(ll > 2.0, "log_loss={}", ll);
    }

    #[test]
    fn log_loss_clamps_extreme() {
        // 0.0 and 1.0 predictions should not cause NaN/Inf
        let preds = vec![0.0, 1.0];
        let outcomes = vec![1.0, 0.0];
        let ll = log_loss(preds, outcomes).unwrap();
        assert!(ll.is_finite());
    }

    #[test]
    fn edge_decay_basic() {
        let prices = vec![0.4, 0.4, 0.4, 0.4];
        let outcomes = vec![1.0, 1.0, 0.0, 0.0];
        let entry_ts = vec![0.0, 3600.0, 7200.0, 10800.0];
        let resolution_ts = vec![14400.0, 14400.0, 14400.0, 14400.0];
        let result = edge_decay(prices, outcomes, entry_ts, resolution_ts, 4).unwrap();
        assert_eq!(result.decay_curve.len(), 4);
        assert!(result.half_life_hours >= 0.0);
    }

    #[test]
    fn edge_decay_empty() {
        let result = edge_decay(vec![], vec![], vec![], vec![], 10).unwrap();
        assert!(result.decay_curve.is_empty());
        assert_eq!(result.half_life_hours, 0.0);
    }

    #[test]
    fn edge_decay_length_mismatch() {
        let result = edge_decay(vec![0.5], vec![], vec![0.0], vec![1.0], 5);
        assert!(result.is_err());
    }

    #[test]
    fn brier_score_known() {
        // Brier = mean((0.7 - 1)^2, (0.3 - 0)^2) = mean(0.09, 0.09) = 0.09
        let preds = vec![0.7, 0.3];
        let outcomes = vec![1.0, 0.0];
        let result = calibration_curve(preds, outcomes, 5).unwrap();
        assert!((result.brier_score - 0.09).abs() < 0.001);
    }
}
