"""Plot-ready data extraction for backtest results.

Processes raw trading data into frozen dataclasses that any charting
library (matplotlib, plotly, etc.) can consume directly.  No external
dependencies - stdlib only.
"""

from __future__ import annotations

import math
from collections import defaultdict
from dataclasses import dataclass
from datetime import datetime, timezone


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _safe_float(val: float) -> float:
    """Ensure a float value is finite. Replace NaN/Inf with 0.0."""
    if math.isfinite(val):
        return val
    return 0.0


# ---------------------------------------------------------------------------
# Output dataclasses (all frozen)
# ---------------------------------------------------------------------------

@dataclass(frozen=True)
class BandData:
    """Bollinger-style band envelope for time-series plots."""

    timestamps: tuple[float, ...]
    middle: tuple[float, ...]
    upper: tuple[float, ...]
    lower: tuple[float, ...]
    bandwidth: tuple[float, ...]


@dataclass(frozen=True)
class UnderwaterData:
    """Per-point drawdown depth for area charts."""

    timestamps: tuple[float, ...]
    drawdown_pct: tuple[float, ...]
    peak: tuple[float, ...]


@dataclass(frozen=True)
class HistogramData:
    """Distribution data for bar charts."""

    bin_edges: tuple[float, ...]
    bin_centers: tuple[float, ...]
    counts: tuple[int, ...]
    frequencies: tuple[float, ...]
    mean: float
    median: float
    std: float
    skew: float
    kurtosis: float


@dataclass(frozen=True)
class HeatmapData:
    """2D matrix for heatmap rendering."""

    values: tuple[tuple[float, ...], ...]
    row_labels: tuple[str, ...]
    col_labels: tuple[str, ...]
    title: str


@dataclass(frozen=True)
class CurveData:
    """Single named time series."""

    timestamps: tuple[float, ...]
    values: tuple[float, ...]
    label: str


@dataclass(frozen=True)
class NormalizedEquityData:
    """Equity curve normalised to 1.0 with optional bands."""

    timestamps: tuple[float, ...]
    equity: tuple[float, ...]
    cumulative_pnl: tuple[float, ...]
    bands: BandData | None


@dataclass(frozen=True)
class TradeScatterPoint:
    """Single trade marker for scatter plots."""

    timestamp: float
    price: float
    size: float
    side: str
    market_id: str
    pnl: float


@dataclass(frozen=True)
class TradeScatterData:
    """Grouped trade scatter data."""

    points: tuple[TradeScatterPoint, ...]
    buys: tuple[TradeScatterPoint, ...]
    sells: tuple[TradeScatterPoint, ...]


@dataclass(frozen=True)
class CalibrationPlotData:
    """Predicted-vs-actual calibration chart data."""

    bin_centers: tuple[float, ...]
    actual_freq: tuple[float, ...]
    counts: tuple[int, ...]
    perfect_line: tuple[float, ...]
    ece: float
    brier_score: float


@dataclass(frozen=True)
class RollingStatData:
    """Rolling risk-adjusted metrics."""

    sharpe: CurveData
    sortino: CurveData


@dataclass(frozen=True)
class PlotBundle:
    """One-call extraction of all plot data from a BacktestResult."""

    equity: NormalizedEquityData
    underwater: UnderwaterData
    return_hist: HistogramData
    pnl_hist: HistogramData | None
    monthly_heatmap: HeatmapData | None
    rolling: RollingStatData | None
    trades: TradeScatterData
    calibration: CalibrationPlotData | None


# ---------------------------------------------------------------------------
# Public functions
# ---------------------------------------------------------------------------

def bollinger_bands(
    series: list[tuple[float, float]],
    window: int = 20,
    num_std: float = 2.0,
) -> BandData:
    """Compute Bollinger-style bands over a time series.

    Args:
        series: List of (timestamp, value) tuples.
        window: Rolling window size.
        num_std: Number of standard deviations for bands.

    Returns:
        BandData with middle, upper, lower, and bandwidth.
    """
    if not series or window < 1:
        return BandData((), (), (), (), ())

    timestamps: list[float] = []
    middle: list[float] = []
    upper: list[float] = []
    lower: list[float] = []
    bandwidth: list[float] = []

    values = [v for _, v in series]
    ts_list = [t for t, _ in series]

    for i in range(window - 1, len(values)):
        w = values[i - window + 1: i + 1]
        ts = ts_list[i]
        mean = sum(w) / len(w)
        var = sum((x - mean) ** 2 for x in w) / len(w)
        std = math.sqrt(var) if var > 0 else 0.0

        timestamps.append(ts)
        middle.append(_safe_float(mean))
        upper.append(_safe_float(mean + num_std * std))
        lower.append(_safe_float(mean - num_std * std))
        bw = _safe_float((2.0 * num_std * std) / mean) if mean != 0 else 0.0
        bandwidth.append(bw)

    return BandData(
        timestamps=tuple(timestamps),
        middle=tuple(middle),
        upper=tuple(upper),
        lower=tuple(lower),
        bandwidth=tuple(bandwidth),
    )


def underwater_curve(
    equity_curve: list[tuple[float, float]],
) -> UnderwaterData:
    """Compute per-point drawdown depth from an equity curve.

    Args:
        equity_curve: List of (timestamp, equity) tuples.

    Returns:
        UnderwaterData with drawdown percentages and peak tracking.
    """
    if not equity_curve:
        return UnderwaterData((), (), ())

    timestamps: list[float] = []
    drawdown_pct: list[float] = []
    peaks: list[float] = []
    peak = equity_curve[0][1]

    for ts, equity in equity_curve:
        if equity > peak:
            peak = equity
        dd_pct = _safe_float(((peak - equity) / peak) * 100.0) if peak > 0 else 0.0
        timestamps.append(ts)
        drawdown_pct.append(dd_pct)
        peaks.append(peak)

    return UnderwaterData(
        timestamps=tuple(timestamps),
        drawdown_pct=tuple(drawdown_pct),
        peak=tuple(peaks),
    )


def histogram(
    values: list[float],
    n_bins: int = 30,
) -> HistogramData:
    """Build histogram data from a list of values.

    Args:
        values: Raw data points.
        n_bins: Number of bins.

    Returns:
        HistogramData with edges, centers, counts, frequencies, and stats.
    """
    # Filter out non-finite values
    clean = [v for v in values if math.isfinite(v)]
    if not clean:
        return HistogramData((), (), (), (), 0.0, 0.0, 0.0, 0.0, 0.0)

    n_bins = max(1, n_bins)
    lo = min(clean)
    hi = max(clean)

    # Handle single-value case
    if lo == hi:
        return HistogramData(
            bin_edges=(lo, hi),
            bin_centers=(lo,),
            counts=(len(clean),),
            frequencies=(1.0,),
            mean=lo,
            median=lo,
            std=0.0,
            skew=0.0,
            kurtosis=0.0,
        )

    step = (hi - lo) / n_bins
    edges = [lo + i * step for i in range(n_bins + 1)]
    counts_list = [0] * n_bins

    for v in clean:
        idx = int((v - lo) / step)
        idx = min(idx, n_bins - 1)  # clamp last value into final bin
        counts_list[idx] += 1

    total = len(clean)
    frequencies = [c / total for c in counts_list]
    centers = [(edges[i] + edges[i + 1]) / 2.0 for i in range(n_bins)]

    # Statistics
    mean = sum(clean) / total
    sorted_clean = sorted(clean)
    if total % 2 == 1:
        median = sorted_clean[total // 2]
    else:
        median = (sorted_clean[total // 2 - 1] + sorted_clean[total // 2]) / 2.0

    variance = sum((x - mean) ** 2 for x in clean) / total if total > 0 else 0.0
    std = math.sqrt(variance) if variance > 0 else 0.0

    # Skewness and kurtosis (population)
    if std > 0 and total > 0:
        skew = _safe_float(
            sum(((x - mean) / std) ** 3 for x in clean) / total
        )
        kurtosis = _safe_float(
            sum(((x - mean) / std) ** 4 for x in clean) / total - 3.0
        )
    else:
        skew = 0.0
        kurtosis = 0.0

    return HistogramData(
        bin_edges=tuple(edges),
        bin_centers=tuple(centers),
        counts=tuple(counts_list),
        frequencies=tuple(frequencies),
        mean=_safe_float(mean),
        median=_safe_float(median),
        std=_safe_float(std),
        skew=skew,
        kurtosis=kurtosis,
    )


def return_distribution(
    equity_curve: list[tuple[float, float]],
    n_bins: int = 30,
) -> HistogramData:
    """Extract period returns from an equity curve and build a histogram.

    Args:
        equity_curve: List of (timestamp, equity) tuples.
        n_bins: Number of histogram bins.

    Returns:
        HistogramData of period-over-period returns.
    """
    if len(equity_curve) < 2:
        return HistogramData((), (), (), (), 0.0, 0.0, 0.0, 0.0, 0.0)

    returns: list[float] = []
    for i in range(1, len(equity_curve)):
        prev = equity_curve[i - 1][1]
        curr = equity_curve[i][1]
        if prev > 0:
            returns.append((curr - prev) / prev)

    return histogram(returns, n_bins)


def pnl_distribution(
    trades: list,
    n_bins: int = 30,
) -> HistogramData:
    """Build a PnL histogram from FIFO-matched round-trip trades.

    Args:
        trades: List of Fill objects from the engine.
        n_bins: Number of histogram bins.

    Returns:
        HistogramData of per-trade PnL values.
    """
    from horizon.analytics import _analyze_trades

    analysis = _analyze_trades(trades)
    pnls = [rt.pnl for rt in analysis.round_trips]
    return histogram(pnls, n_bins)


def monthly_returns_heatmap(
    equity_curve: list[tuple[float, float]],
) -> HeatmapData:
    """Build a year x month heatmap of returns from an equity curve.

    Args:
        equity_curve: List of (timestamp, equity) tuples.

    Returns:
        HeatmapData with years as rows and months as columns.
    """
    from horizon.analytics import _compute_monthly_returns

    monthly = _compute_monthly_returns(equity_curve)
    if not monthly:
        return HeatmapData((), (), (), "Monthly Returns (%)")

    # Parse YYYY-MM keys into (year, month) pairs
    year_month: dict[int, dict[int, float]] = defaultdict(dict)
    for key, val in monthly.items():
        parts = key.split("-")
        if len(parts) == 2:
            year, month = int(parts[0]), int(parts[1])
            year_month[year][month] = val

    years = sorted(year_month.keys())
    months = list(range(1, 13))
    month_labels = (
        "Jan", "Feb", "Mar", "Apr", "May", "Jun",
        "Jul", "Aug", "Sep", "Oct", "Nov", "Dec",
    )

    rows: list[tuple[float, ...]] = []
    for y in years:
        row = tuple(year_month[y].get(m, 0.0) for m in months)
        rows.append(row)

    return HeatmapData(
        values=tuple(rows),
        row_labels=tuple(str(y) for y in years),
        col_labels=month_labels,
        title="Monthly Returns (%)",
    )


def correlation_heatmap(
    series_dict: dict[str, list[float]],
) -> HeatmapData:
    """Build a correlation matrix heatmap from named series.

    Args:
        series_dict: Mapping of series name -> list of values.
            All series must have the same length.

    Returns:
        HeatmapData with correlation coefficients.
    """
    names = sorted(series_dict.keys())
    if len(names) < 2:
        return HeatmapData((), (), (), "Correlation Matrix")

    n = len(names)
    # Validate equal lengths
    length = len(series_dict[names[0]])
    for name in names:
        if len(series_dict[name]) != length:
            return HeatmapData((), (), (), "Correlation Matrix")

    if length < 2:
        return HeatmapData((), (), (), "Correlation Matrix")

    # Compute means and stds
    means = {name: sum(series_dict[name]) / length for name in names}
    stds = {}
    for name in names:
        var = sum((x - means[name]) ** 2 for x in series_dict[name]) / length
        stds[name] = math.sqrt(var) if var > 0 else 0.0

    rows: list[tuple[float, ...]] = []
    for i in range(n):
        row: list[float] = []
        for j in range(n):
            if i == j:
                row.append(1.0)
            elif stds[names[i]] == 0 or stds[names[j]] == 0:
                row.append(0.0)
            else:
                cov = sum(
                    (series_dict[names[i]][k] - means[names[i]])
                    * (series_dict[names[j]][k] - means[names[j]])
                    for k in range(length)
                ) / length
                corr = _safe_float(cov / (stds[names[i]] * stds[names[j]]))
                row.append(corr)
        rows.append(tuple(row))

    labels = tuple(names)
    return HeatmapData(
        values=tuple(rows),
        row_labels=labels,
        col_labels=labels,
        title="Correlation Matrix",
    )


def rolling_stats(
    equity_curve: list[tuple[float, float]],
    window: int = 30,
) -> RollingStatData:
    """Compute rolling Sharpe and Sortino from an equity curve.

    Args:
        equity_curve: List of (timestamp, equity) tuples.
        window: Rolling window size.

    Returns:
        RollingStatData with sharpe and sortino CurveData.
    """
    from horizon.analytics import _rolling_sharpe_sortino

    sharpe_raw, sortino_raw = _rolling_sharpe_sortino(equity_curve, window)

    sharpe = CurveData(
        timestamps=tuple(ts for ts, _ in sharpe_raw),
        values=tuple(v for _, v in sharpe_raw),
        label="Rolling Sharpe",
    )
    sortino = CurveData(
        timestamps=tuple(ts for ts, _ in sortino_raw),
        values=tuple(v for _, v in sortino_raw),
        label="Rolling Sortino",
    )
    return RollingStatData(sharpe=sharpe, sortino=sortino)


def normalized_equity(
    equity_curve: list[tuple[float, float]],
    initial_capital: float = 0.0,
    include_bands: bool = False,
    band_window: int = 20,
    band_std: float = 2.0,
) -> NormalizedEquityData:
    """Normalize an equity curve to start at 1.0 with optional bands.

    Args:
        equity_curve: List of (timestamp, equity) tuples.
        initial_capital: Starting capital (uses first equity point if 0).
        include_bands: Whether to compute Bollinger bands on normalized equity.
        band_window: Bollinger band window.
        band_std: Bollinger band standard deviations.

    Returns:
        NormalizedEquityData with normalized equity and cumulative PnL.
    """
    if not equity_curve:
        return NormalizedEquityData((), (), (), None)

    base = initial_capital if initial_capital > 0 else equity_curve[0][1]
    if base <= 0:
        base = 1.0

    timestamps: list[float] = []
    equity: list[float] = []
    cum_pnl: list[float] = []

    for ts, eq in equity_curve:
        timestamps.append(ts)
        norm = _safe_float(eq / base)
        equity.append(norm)
        cum_pnl.append(_safe_float(eq - base))

    bands = None
    if include_bands:
        norm_series = list(zip(timestamps, equity))
        bands = bollinger_bands(norm_series, band_window, band_std)

    return NormalizedEquityData(
        timestamps=tuple(timestamps),
        equity=tuple(equity),
        cumulative_pnl=tuple(cum_pnl),
        bands=bands,
    )


def trade_scatter(
    trades: list,
) -> TradeScatterData:
    """Extract trade scatter plot data from fills via FIFO matching.

    Args:
        trades: List of Fill objects from the engine.

    Returns:
        TradeScatterData with all points split into buys and sells.
    """
    from horizon.analytics import _analyze_trades

    analysis = _analyze_trades(trades)

    all_points: list[TradeScatterPoint] = []
    buy_points: list[TradeScatterPoint] = []
    sell_points: list[TradeScatterPoint] = []

    for rt in analysis.round_trips:
        # Entry point
        entry_side = "buy" if rt.direction == "long" else "sell"
        entry = TradeScatterPoint(
            timestamp=rt.entry_timestamp,
            price=rt.entry_price,
            size=rt.size,
            side=entry_side,
            market_id=rt.market_id,
            pnl=0.0,  # PnL attributed to exit
        )
        all_points.append(entry)
        if entry_side == "buy":
            buy_points.append(entry)
        else:
            sell_points.append(entry)

        # Exit point
        exit_side = "sell" if rt.direction == "long" else "buy"
        exit_pt = TradeScatterPoint(
            timestamp=rt.exit_timestamp,
            price=rt.exit_price,
            size=rt.size,
            side=exit_side,
            market_id=rt.market_id,
            pnl=rt.pnl,
        )
        all_points.append(exit_pt)
        if exit_side == "buy":
            buy_points.append(exit_pt)
        else:
            sell_points.append(exit_pt)

    return TradeScatterData(
        points=tuple(all_points),
        buys=tuple(buy_points),
        sells=tuple(sell_points),
    )


def calibration_plot(
    trades: list,
    outcomes: dict[str, float],
    n_bins: int = 10,
) -> CalibrationPlotData:
    """Compute calibration plot data from trades and outcomes.

    Args:
        trades: List of Fill objects from the engine.
        outcomes: Dict of market_id -> outcome (0.0 or 1.0).
        n_bins: Number of calibration bins.

    Returns:
        CalibrationPlotData with bin centers, actual frequencies, and metrics.
    """
    from horizon.analytics import _extract_predictions_outcomes

    predictions, outcome_vals = _extract_predictions_outcomes(trades, outcomes)

    if not predictions:
        return CalibrationPlotData((), (), (), (), 0.0, 0.0)

    n_bins = max(1, n_bins)

    # Try Rust calibration first
    try:
        from horizon._horizon import calibration_curve as _rust_calibration

        cal = _rust_calibration(predictions, outcome_vals, n_bins)
        bin_centers_list: list[float] = []
        actual_freq_list: list[float] = []
        counts_list: list[int] = []

        for b in cal.bins:
            bin_centers_list.append(b.mean_predicted)
            actual_freq_list.append(b.actual_frequency)
            counts_list.append(b.count)

        perfect = list(bin_centers_list)  # y=x line

        return CalibrationPlotData(
            bin_centers=tuple(bin_centers_list),
            actual_freq=tuple(actual_freq_list),
            counts=tuple(counts_list),
            perfect_line=tuple(perfect),
            ece=cal.ece,
            brier_score=_safe_float(
                sum((p - o) ** 2 for p, o in zip(predictions, outcome_vals))
                / len(predictions)
            ),
        )
    except (ImportError, Exception):
        pass

    # Pure-Python fallback
    step = 1.0 / n_bins
    bins: list[list[tuple[float, float]]] = [[] for _ in range(n_bins)]

    for pred, out in zip(predictions, outcome_vals):
        idx = int(pred / step)
        idx = min(idx, n_bins - 1)
        bins[idx].append((pred, out))

    bin_centers: list[float] = []
    actual_freq: list[float] = []
    counts: list[int] = []

    for i in range(n_bins):
        if bins[i]:
            mean_pred = sum(p for p, _ in bins[i]) / len(bins[i])
            mean_out = sum(o for _, o in bins[i]) / len(bins[i])
            bin_centers.append(_safe_float(mean_pred))
            actual_freq.append(_safe_float(mean_out))
            counts.append(len(bins[i]))
        else:
            center = (i + 0.5) * step
            bin_centers.append(center)
            actual_freq.append(0.0)
            counts.append(0)

    perfect_line = list(bin_centers)
    total = len(predictions)
    ece = sum(
        (c / total) * abs(a - p)
        for c, a, p in zip(counts, actual_freq, bin_centers)
        if c > 0
    ) if total > 0 else 0.0

    brier = sum((p - o) ** 2 for p, o in zip(predictions, outcome_vals)) / total

    return CalibrationPlotData(
        bin_centers=tuple(bin_centers),
        actual_freq=tuple(actual_freq),
        counts=tuple(counts),
        perfect_line=tuple(perfect_line),
        ece=_safe_float(ece),
        brier_score=_safe_float(brier),
    )


def from_backtest(
    result,
    n_bins: int = 30,
    rolling_window: int = 30,
    include_bands: bool = False,
    band_window: int = 20,
    band_std: float = 2.0,
    calibration_bins: int = 10,
) -> PlotBundle:
    """Extract all plot data from a BacktestResult in one call.

    Args:
        result: A BacktestResult instance.
        n_bins: Number of bins for histograms.
        rolling_window: Window size for rolling stats.
        include_bands: Whether to include Bollinger bands on equity.
        band_window: Bollinger band window.
        band_std: Bollinger band standard deviations.
        calibration_bins: Number of calibration bins.

    Returns:
        PlotBundle with all plot data.
    """
    equity_curve = getattr(result, "equity_curve", [])
    trades_list = getattr(result, "trades", [])
    initial_cap = getattr(result, "initial_capital", 0.0)
    outcomes = getattr(result, "outcomes", None)

    # Equity
    eq = normalized_equity(
        equity_curve, initial_cap, include_bands, band_window, band_std,
    )

    # Underwater
    uw = underwater_curve(equity_curve)

    # Return distribution
    ret_hist = return_distribution(equity_curve, n_bins)

    # PnL distribution
    pnl_hist = pnl_distribution(trades_list, n_bins) if trades_list else None

    # Monthly heatmap
    monthly = monthly_returns_heatmap(equity_curve) if len(equity_curve) >= 2 else None

    # Rolling stats
    roll = rolling_stats(equity_curve, rolling_window) if len(equity_curve) >= rolling_window + 1 else None

    # Trade scatter
    scatter = trade_scatter(trades_list)

    # Calibration
    cal = None
    if outcomes and trades_list:
        cal = calibration_plot(trades_list, outcomes, calibration_bins)

    return PlotBundle(
        equity=eq,
        underwater=uw,
        return_hist=ret_hist,
        pnl_hist=pnl_hist,
        monthly_heatmap=monthly,
        rolling=roll,
        trades=scatter,
        calibration=cal,
    )
