"""Arrow-Debreu state-contingent securities analysis.

Core theoretical contribution: tests for market completeness,
risk-neutral density extraction, moment computation, and arbitrage
violation detection using Polymarket state-price data.
"""

from __future__ import annotations

import pandas as pd
import numpy as np
from dataclasses import dataclass


# ---------------------------------------------------------------------------
# Data classes
# ---------------------------------------------------------------------------

@dataclass
class CompletenessTestResult:
    """Result container for the Arrow-Debreu completeness test."""
    daily_sums: pd.Series          # sum of state prices each day
    mean_sum: float                # should be close to 1.0
    std_sum: float
    mean_deviation: float          # mean(sum - 1.0)
    max_deviation: float
    all_positive: bool             # no negative state prices
    discount_factor: float         # implied 1/(1+r) from mean sum


# ---------------------------------------------------------------------------
# Completeness
# ---------------------------------------------------------------------------

def test_market_completeness(state_prices: pd.DataFrame) -> CompletenessTestResult:
    """Test Arrow-Debreu completeness conditions.

    Parameters
    ----------
    state_prices : pd.DataFrame
        DataFrame where columns = states, values = prices.
        In a complete market the sum of state prices equals 1/(1+r),
        which is approximately 1.0 for short horizons.

    Returns
    -------
    CompletenessTestResult
    """
    daily_sums = state_prices.sum(axis=1)
    deviations = daily_sums - 1.0

    return CompletenessTestResult(
        daily_sums=daily_sums,
        mean_sum=float(daily_sums.mean()),
        std_sum=float(daily_sums.std()),
        mean_deviation=float(deviations.abs().mean()),
        max_deviation=float(deviations.abs().max()),
        all_positive=bool((state_prices >= 0).all().all()),
        discount_factor=float(daily_sums.mean()),  # 1/(1+r) ≈ mean sum
    )


# ---------------------------------------------------------------------------
# Risk-neutral density
# ---------------------------------------------------------------------------

def extract_risk_neutral_density(
    state_prices: pd.DataFrame,
    state_values: list[float],
) -> pd.DataFrame:
    """Extract risk-neutral probability density from state prices.

    Normalises each row so prices sum to 1.0; the resulting values are
    interpreted as risk-neutral probabilities.

    Parameters
    ----------
    state_prices : pd.DataFrame
        Columns = states, values = prices.
    state_values : list[float]
        Numerical value associated with each state (same order as columns).

    Returns
    -------
    pd.DataFrame
        Same shape as *state_prices* with values summing to 1.0 per row.
    """
    if len(state_values) != state_prices.shape[1]:
        raise ValueError(
            f"state_values length ({len(state_values)}) must match "
            f"number of columns ({state_prices.shape[1]})"
        )

    row_sums = state_prices.sum(axis=1)
    # Avoid division by zero
    row_sums = row_sums.replace(0, np.nan)
    rn_density = state_prices.div(row_sums, axis=0)
    return rn_density


# ---------------------------------------------------------------------------
# Risk-neutral moments
# ---------------------------------------------------------------------------

def compute_risk_neutral_moments(
    rn_density: pd.DataFrame,
    state_values: list[float],
) -> pd.DataFrame:
    """Compute time series of risk-neutral moments.

    Parameters
    ----------
    rn_density : pd.DataFrame
        Risk-neutral probabilities (rows sum to 1).
    state_values : list[float]
        Numerical values for each state (e.g., basis-point changes
        [-50, -25, 0, 25, 50]).

    Returns
    -------
    pd.DataFrame
        Columns ``[mean, variance, skewness, kurtosis]`` indexed by date.
    """
    vals = np.array(state_values, dtype=float)
    probs = rn_density.values  # (T, K)

    # Mean
    mu = probs @ vals  # (T,)

    # Central moments
    deviations = vals[np.newaxis, :] - mu[:, np.newaxis]  # (T, K)
    var = (probs * deviations ** 2).sum(axis=1)
    sigma = np.sqrt(np.maximum(var, 0))

    # Protect against zero variance
    sigma_safe = np.where(sigma > 0, sigma, np.nan)
    skew = (probs * deviations ** 3).sum(axis=1) / sigma_safe ** 3
    kurt = (probs * deviations ** 4).sum(axis=1) / sigma_safe ** 4

    return pd.DataFrame(
        {"mean": mu, "variance": var, "skewness": skew, "kurtosis": kurt},
        index=rn_density.index,
    )


# ---------------------------------------------------------------------------
# Arbitrage violations
# ---------------------------------------------------------------------------

def detect_arbitrage_violations(
    prices: pd.DataFrame,
    market_type: str = "binary",
) -> pd.DataFrame:
    """Detect no-arbitrage violations.

    For binary markets: p(Yes) + p(No) should equal 1.
    For state-space groups: sum of state prices should be approximately 1.

    Parameters
    ----------
    prices : pd.DataFrame
        Price data (columns = states or Yes/No).
    market_type : str
        ``"binary"`` or ``"state_space"``.

    Returns
    -------
    pd.DataFrame
        Columns ``[date, violation_size, violation_pct]``.
    """
    row_sums = prices.sum(axis=1)
    deviation = (row_sums - 1.0).abs()

    result = pd.DataFrame({
        "date": prices.index,
        "violation_size": deviation.values,
        "violation_pct": (deviation.values * 100),
    })
    result = result.set_index("date")
    return result


# ---------------------------------------------------------------------------
# State-price dynamics
# ---------------------------------------------------------------------------

def state_price_dynamics(state_prices: pd.DataFrame) -> dict:
    """Analyse dynamics of state prices over time.

    Returns
    -------
    dict
        ``correlation_matrix`` : cross-correlations between states.
        ``persistence``        : AR(1) coefficients for each state price.
        ``mean_reversion_speed``: half-life of deviations from mean.
    """
    # Cross-correlation matrix
    corr_matrix = state_prices.corr()

    # AR(1) coefficients: regress p_t on p_{t-1} for each column
    persistence = {}
    half_lives = {}
    for col in state_prices.columns:
        series = state_prices[col].dropna()
        if len(series) < 3:
            persistence[col] = np.nan
            half_lives[col] = np.nan
            continue

        y = series.iloc[1:].values
        x = series.iloc[:-1].values
        # OLS: phi = cov(y,x) / var(x)
        x_dm = x - x.mean()
        y_dm = y - y.mean()
        var_x = (x_dm ** 2).sum()
        if var_x == 0:
            persistence[col] = np.nan
            half_lives[col] = np.nan
            continue

        phi = (x_dm * y_dm).sum() / var_x
        persistence[col] = float(phi)

        # Half-life: t_{1/2} = -ln(2) / ln(|phi|), defined when |phi| < 1
        if 0 < abs(phi) < 1:
            half_lives[col] = float(-np.log(2) / np.log(abs(phi)))
        else:
            half_lives[col] = np.nan

    return {
        "correlation_matrix": corr_matrix,
        "persistence": pd.Series(persistence),
        "mean_reversion_speed": pd.Series(half_lives),
    }