//! Synthetic index construction for prediction markets.
//!
//! Build custom indices from prediction market prices with various weighting methods.
//! Track index levels, compute rebalance actions, and detect arbitrage signals.

use pyo3::prelude::*;

/// Weighting method for index construction.
#[pyclass(eq, eq_int, frozen, hash)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum WeightingMethod {
    Equal = 0,
    InverseVol = 1,
    Liquidity = 2,
    Custom = 3,
}

#[pymethods]
impl WeightingMethod {
    fn __repr__(&self) -> String {
        match self {
            Self::Equal => "WeightingMethod.Equal".to_string(),
            Self::InverseVol => "WeightingMethod.InverseVol".to_string(),
            Self::Liquidity => "WeightingMethod.Liquidity".to_string(),
            Self::Custom => "WeightingMethod.Custom".to_string(),
        }
    }
    fn __str__(&self) -> String {
        self.__repr__()
    }
}

/// A single component weight in the index.
#[pyclass(frozen, get_all)]
#[derive(Debug, Clone)]
pub struct IndexWeight {
    pub market_id: String,
    pub weight: f64,
}

#[pymethods]
impl IndexWeight {
    #[new]
    fn new(market_id: String, weight: f64) -> Self {
        Self { market_id, weight }
    }

    fn __repr__(&self) -> String {
        format!("IndexWeight(market='{}', weight={:.4})", self.market_id, self.weight)
    }
}

/// Computed index level.
#[pyclass(frozen, get_all)]
#[derive(Debug, Clone)]
pub struct IndexLevel {
    pub level: f64,
    pub timestamp: f64,
    pub num_components: usize,
    pub total_weight: f64,
}

#[pymethods]
impl IndexLevel {
    fn __repr__(&self) -> String {
        format!(
            "IndexLevel(level={:.4}, components={}, total_weight={:.4})",
            self.level, self.num_components, self.total_weight
        )
    }
}

/// Rebalance action for a single component.
#[pyclass(frozen, get_all)]
#[derive(Debug, Clone)]
pub struct IndexRebalanceAction {
    pub market_id: String,
    pub current_weight: f64,
    pub target_weight: f64,
    pub delta: f64,
}

#[pymethods]
impl IndexRebalanceAction {
    fn __repr__(&self) -> String {
        format!(
            "IndexRebalanceAction(market='{}', delta={:+.4})",
            self.market_id, self.delta
        )
    }
}

/// Arbitrage signal between index level and component prices.
#[pyclass(frozen, get_all)]
#[derive(Debug, Clone)]
pub struct IndexArbSignal {
    pub index_level: f64,
    pub fair_level: f64,
    pub deviation: f64,
    pub signal_strength: f64,
}

#[pymethods]
impl IndexArbSignal {
    fn __repr__(&self) -> String {
        format!(
            "IndexArbSignal(index={:.4}, fair={:.4}, dev={:+.4}, strength={:.4})",
            self.index_level, self.fair_level, self.deviation, self.signal_strength
        )
    }
}

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

/// Compute index level as weighted sum of prices.
///
/// prices[i] corresponds to weights[i].market_id.
/// Returns IndexLevel with the computed level.
#[pyfunction]
#[pyo3(signature = (prices, weights, base_level=100.0))]
pub fn compute_index_level(
    prices: Vec<f64>,
    weights: Vec<f64>,
    base_level: f64,
) -> IndexLevel {
    let n = prices.len().min(weights.len());
    if n == 0 {
        return IndexLevel {
            level: base_level,
            timestamp: now_secs(),
            num_components: 0,
            total_weight: 0.0,
        };
    }

    let total_weight: f64 = weights.iter().take(n).sum();
    if total_weight <= 0.0 {
        return IndexLevel {
            level: base_level,
            timestamp: now_secs(),
            num_components: n,
            total_weight: 0.0,
        };
    }

    let weighted_sum: f64 = prices
        .iter()
        .zip(weights.iter())
        .take(n)
        .map(|(p, w)| p * w)
        .sum();

    let level = finite_or_zero(base_level * weighted_sum / total_weight);

    IndexLevel {
        level,
        timestamp: now_secs(),
        num_components: n,
        total_weight: finite_or_zero(total_weight),
    }
}

/// Compute equal weights for N components.
#[pyfunction]
pub fn compute_weights_equal(n: usize) -> Vec<f64> {
    if n == 0 {
        return Vec::new();
    }
    vec![1.0 / n as f64; n]
}

/// Compute inverse-volatility weights.
///
/// Components with lower volatility get higher weight.
#[pyfunction]
pub fn compute_weights_inverse_vol(volatilities: Vec<f64>) -> Vec<f64> {
    if volatilities.is_empty() {
        return Vec::new();
    }

    let inv_vols: Vec<f64> = volatilities
        .iter()
        .map(|v| {
            if *v > 0.0 && v.is_finite() {
                1.0 / v
            } else {
                0.0
            }
        })
        .collect();

    let total: f64 = inv_vols.iter().sum();
    if total <= 0.0 {
        return compute_weights_equal(volatilities.len());
    }

    inv_vols.iter().map(|iv| finite_or_zero(iv / total)).collect()
}

/// Compute liquidity-weighted weights.
///
/// Components with higher volume get higher weight.
#[pyfunction]
pub fn compute_weights_liquidity(volumes: Vec<f64>) -> Vec<f64> {
    if volumes.is_empty() {
        return Vec::new();
    }

    let total: f64 = volumes.iter().filter(|v| v.is_finite() && **v > 0.0).sum();
    if total <= 0.0 {
        return compute_weights_equal(volumes.len());
    }

    volumes
        .iter()
        .map(|v| {
            if *v > 0.0 && v.is_finite() {
                finite_or_zero(v / total)
            } else {
                0.0
            }
        })
        .collect()
}

/// Compute rebalance actions to move from current to target weights.
#[pyfunction]
pub fn rebalance_weights(
    market_ids: Vec<String>,
    current_weights: Vec<f64>,
    target_weights: Vec<f64>,
) -> Vec<IndexRebalanceAction> {
    let n = market_ids
        .len()
        .min(current_weights.len())
        .min(target_weights.len());

    (0..n)
        .map(|i| {
            let delta = finite_or_zero(target_weights[i] - current_weights[i]);
            IndexRebalanceAction {
                market_id: market_ids[i].clone(),
                current_weight: current_weights[i],
                target_weight: target_weights[i],
                delta,
            }
        })
        .collect()
}

/// Compute tracking error between portfolio and index returns.
///
/// Tracking error = std(portfolio_returns - index_returns)
#[pyfunction]
pub fn index_tracking_error(
    portfolio_returns: Vec<f64>,
    index_returns: Vec<f64>,
) -> f64 {
    let n = portfolio_returns.len().min(index_returns.len());
    if n < 2 {
        return 0.0;
    }

    let diffs: Vec<f64> = portfolio_returns
        .iter()
        .zip(index_returns.iter())
        .take(n)
        .map(|(p, i)| p - i)
        .collect();

    let mean: f64 = diffs.iter().sum::<f64>() / n as f64;
    let var: f64 = diffs.iter().map(|d| (d - mean).powi(2)).sum::<f64>() / (n - 1) as f64;

    finite_or_zero(var.sqrt())
}

/// Detect arbitrage signal between index and components.
///
/// Compares a tracked index level to the fair level computed from component prices.
#[pyfunction]
#[pyo3(signature = (index_level, prices, weights, threshold=0.01))]
pub fn index_arb_signal(
    index_level: f64,
    prices: Vec<f64>,
    weights: Vec<f64>,
    threshold: f64,
) -> IndexArbSignal {
    let n = prices.len().min(weights.len());
    if n == 0 {
        return IndexArbSignal {
            index_level,
            fair_level: index_level,
            deviation: 0.0,
            signal_strength: 0.0,
        };
    }

    let total_weight: f64 = weights.iter().take(n).sum();
    let fair_level = if total_weight > 0.0 {
        let weighted_sum: f64 = prices
            .iter()
            .zip(weights.iter())
            .take(n)
            .map(|(p, w)| p * w)
            .sum();
        finite_or_zero(weighted_sum / total_weight)
    } else {
        index_level
    };

    let deviation = finite_or_zero(index_level - fair_level);
    let signal_strength = if fair_level.abs() > 1e-10 {
        let raw = (deviation / fair_level).abs();
        if raw > threshold {
            finite_or_zero(raw / threshold).min(10.0)
        } else {
            0.0
        }
    } else {
        0.0
    };

    IndexArbSignal {
        index_level,
        fair_level,
        deviation,
        signal_strength,
    }
}

fn now_secs() -> f64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs_f64()
}

/// Register all index functions in the Python module.
pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<WeightingMethod>()?;
    m.add_class::<IndexWeight>()?;
    m.add_class::<IndexLevel>()?;
    m.add_class::<IndexRebalanceAction>()?;
    m.add_class::<IndexArbSignal>()?;
    m.add_function(wrap_pyfunction!(compute_index_level, m)?)?;
    m.add_function(wrap_pyfunction!(compute_weights_equal, m)?)?;
    m.add_function(wrap_pyfunction!(compute_weights_inverse_vol, m)?)?;
    m.add_function(wrap_pyfunction!(compute_weights_liquidity, m)?)?;
    m.add_function(wrap_pyfunction!(rebalance_weights, m)?)?;
    m.add_function(wrap_pyfunction!(index_tracking_error, m)?)?;
    m.add_function(wrap_pyfunction!(index_arb_signal, m)?)?;
    Ok(())
}

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

    #[test]
    fn test_compute_index_level_basic() {
        let prices = vec![0.5, 0.3, 0.8];
        let weights = vec![1.0, 1.0, 1.0];
        let level = compute_index_level(prices, weights, 100.0);
        // weighted avg = (0.5+0.3+0.8)/3 = 0.5333
        // level = 100 * 0.5333 = 53.33
        assert!((level.level - 53.333).abs() < 0.1);
        assert_eq!(level.num_components, 3);
    }

    #[test]
    fn test_compute_index_level_empty() {
        let level = compute_index_level(vec![], vec![], 100.0);
        assert!((level.level - 100.0).abs() < 1e-6);
        assert_eq!(level.num_components, 0);
    }

    #[test]
    fn test_compute_index_level_weighted() {
        let prices = vec![0.5, 0.8];
        let weights = vec![2.0, 1.0];
        let level = compute_index_level(prices, weights, 100.0);
        // weighted avg = (0.5*2 + 0.8*1)/3 = 1.8/3 = 0.6
        assert!((level.level - 60.0).abs() < 0.1);
    }

    #[test]
    fn test_compute_weights_equal() {
        let w = compute_weights_equal(4);
        assert_eq!(w.len(), 4);
        for wi in &w {
            assert!((wi - 0.25).abs() < 1e-10);
        }
    }

    #[test]
    fn test_compute_weights_equal_zero() {
        assert!(compute_weights_equal(0).is_empty());
    }

    #[test]
    fn test_compute_weights_inverse_vol() {
        let w = compute_weights_inverse_vol(vec![0.1, 0.2, 0.4]);
        // inv: 10, 5, 2.5 -> total=17.5
        assert!((w[0] - 10.0 / 17.5).abs() < 1e-6);
        assert!((w[1] - 5.0 / 17.5).abs() < 1e-6);
        assert!((w[2] - 2.5 / 17.5).abs() < 1e-6);
    }

    #[test]
    fn test_compute_weights_inverse_vol_zero_vol() {
        let w = compute_weights_inverse_vol(vec![0.0, 0.0]);
        // Falls back to equal weights
        assert!((w[0] - 0.5).abs() < 1e-6);
    }

    #[test]
    fn test_compute_weights_liquidity() {
        let w = compute_weights_liquidity(vec![1000.0, 3000.0]);
        assert!((w[0] - 0.25).abs() < 1e-6);
        assert!((w[1] - 0.75).abs() < 1e-6);
    }

    #[test]
    fn test_rebalance_weights() {
        let actions = rebalance_weights(
            vec!["A".into(), "B".into()],
            vec![0.6, 0.4],
            vec![0.5, 0.5],
        );
        assert_eq!(actions.len(), 2);
        assert!((actions[0].delta - (-0.1)).abs() < 1e-10);
        assert!((actions[1].delta - 0.1).abs() < 1e-10);
    }

    #[test]
    fn test_index_tracking_error() {
        let te = index_tracking_error(
            vec![0.01, 0.02, -0.01, 0.03],
            vec![0.01, 0.02, -0.01, 0.03],
        );
        assert!(te < 1e-10); // identical returns
    }

    #[test]
    fn test_index_tracking_error_different() {
        let te = index_tracking_error(
            vec![0.01, 0.05, -0.02, 0.04],
            vec![0.02, 0.01, -0.01, 0.01],
        );
        assert!(te > 0.0);
    }

    #[test]
    fn test_index_tracking_error_short() {
        assert!((index_tracking_error(vec![0.01], vec![0.02]) - 0.0).abs() < 1e-10);
    }

    #[test]
    fn test_index_arb_signal_no_deviation() {
        let sig = index_arb_signal(0.5, vec![0.5, 0.5], vec![1.0, 1.0], 0.01);
        assert!((sig.deviation).abs() < 1e-10);
        assert!((sig.signal_strength).abs() < 1e-10);
    }

    #[test]
    fn test_index_arb_signal_with_deviation() {
        let sig = index_arb_signal(0.55, vec![0.5, 0.5], vec![1.0, 1.0], 0.01);
        assert!((sig.fair_level - 0.5).abs() < 1e-10);
        assert!((sig.deviation - 0.05).abs() < 1e-10);
        assert!(sig.signal_strength > 0.0);
    }

    #[test]
    fn test_index_arb_signal_empty() {
        let sig = index_arb_signal(0.5, vec![], vec![], 0.01);
        assert!((sig.deviation).abs() < 1e-10);
    }
}
