"""Tests for horizon.plotting - plot-ready data extraction."""

from __future__ import annotations

import math
from dataclasses import dataclass

import pytest

from horizon.plotting import (
    BandData,
    CalibrationPlotData,
    CurveData,
    HeatmapData,
    HistogramData,
    NormalizedEquityData,
    PlotBundle,
    RollingStatData,
    TradeScatterData,
    TradeScatterPoint,
    UnderwaterData,
    bollinger_bands,
    calibration_plot,
    correlation_heatmap,
    from_backtest,
    histogram,
    monthly_returns_heatmap,
    normalized_equity,
    pnl_distribution,
    return_distribution,
    rolling_stats,
    trade_scatter,
    underwater_curve,
)


# ---------------------------------------------------------------------------
# Mock Fill (same pattern as other test files)
# ---------------------------------------------------------------------------

@dataclass
class MockFill:
    market_id: str = "m1"
    order_side: str = "Buy"
    price: float = 0.5
    size: float = 10.0
    fee: float = 0.0
    timestamp: float = 1000.0


# ---------------------------------------------------------------------------
# Helper
# ---------------------------------------------------------------------------

def _all_finite(*tuples):
    """Assert every value in every tuple is finite."""
    for t in tuples:
        for v in t:
            if isinstance(v, (int, float)):
                assert math.isfinite(v), f"Non-finite value: {v}"


# ---------------------------------------------------------------------------
# BandData / bollinger_bands
# ---------------------------------------------------------------------------

class TestBollingerBands:
    def test_empty(self):
        b = bollinger_bands([])
        assert b == BandData((), (), (), (), ())

    def test_too_few_points(self):
        series = [(float(i), float(i)) for i in range(5)]
        b = bollinger_bands(series, window=20)
        assert len(b.timestamps) == 0

    def test_constant_series(self):
        series = [(float(i), 5.0) for i in range(30)]
        b = bollinger_bands(series, window=10, num_std=2.0)
        assert len(b.timestamps) == 21  # 30 - 10 + 1
        for m in b.middle:
            assert abs(m - 5.0) < 1e-10
        # std=0 → upper==lower==middle
        for u, l, m in zip(b.upper, b.lower, b.middle):
            assert abs(u - m) < 1e-10
            assert abs(l - m) < 1e-10
        for bw in b.bandwidth:
            assert abs(bw) < 1e-10

    def test_known_values(self):
        # 5-point window, simple data
        series = [(float(i), float(i + 1)) for i in range(5)]  # values 1..5
        b = bollinger_bands(series, window=5, num_std=1.0)
        assert len(b.timestamps) == 1
        mean = 3.0  # avg(1,2,3,4,5)
        var = sum((x - mean) ** 2 for x in range(1, 6)) / 5
        std = math.sqrt(var)
        assert abs(b.middle[0] - mean) < 1e-10
        assert abs(b.upper[0] - (mean + std)) < 1e-10
        assert abs(b.lower[0] - (mean - std)) < 1e-10
        _all_finite(b.timestamps, b.middle, b.upper, b.lower, b.bandwidth)

    def test_frozen(self):
        b = bollinger_bands([(0.0, 1.0), (1.0, 2.0)], window=2)
        with pytest.raises(AttributeError):
            b.timestamps = ()  # type: ignore


# ---------------------------------------------------------------------------
# UnderwaterData / underwater_curve
# ---------------------------------------------------------------------------

class TestUnderwaterCurve:
    def test_empty(self):
        u = underwater_curve([])
        assert u == UnderwaterData((), (), ())

    def test_no_drawdown(self):
        curve = [(float(i), 100.0 + i) for i in range(10)]
        u = underwater_curve(curve)
        assert all(d == 0.0 for d in u.drawdown_pct)

    def test_known_drawdown(self):
        curve = [(0.0, 100.0), (1.0, 90.0), (2.0, 80.0), (3.0, 100.0)]
        u = underwater_curve(curve)
        assert abs(u.drawdown_pct[0]) < 1e-10  # at peak
        assert abs(u.drawdown_pct[1] - 10.0) < 1e-10  # 10% down
        assert abs(u.drawdown_pct[2] - 20.0) < 1e-10  # 20% down
        assert abs(u.drawdown_pct[3]) < 1e-10  # recovered
        _all_finite(u.timestamps, u.drawdown_pct, u.peak)

    def test_peak_tracking(self):
        curve = [(0.0, 100.0), (1.0, 120.0), (2.0, 110.0)]
        u = underwater_curve(curve)
        assert u.peak[0] == 100.0
        assert u.peak[1] == 120.0
        assert u.peak[2] == 120.0  # still 120 during drawdown


# ---------------------------------------------------------------------------
# HistogramData / histogram
# ---------------------------------------------------------------------------

class TestHistogram:
    def test_empty(self):
        h = histogram([])
        assert h.counts == ()
        assert h.mean == 0.0

    def test_single_value(self):
        h = histogram([5.0])
        assert h.mean == 5.0
        assert h.median == 5.0
        assert h.std == 0.0
        assert h.counts == (1,)

    def test_known_distribution(self):
        vals = [1.0, 2.0, 3.0, 4.0, 5.0]
        h = histogram(vals, n_bins=5)
        assert sum(h.counts) == 5
        assert abs(h.mean - 3.0) < 1e-10
        assert abs(h.median - 3.0) < 1e-10
        _all_finite(h.bin_edges, h.bin_centers, h.frequencies)

    def test_nan_inf_filtered(self):
        vals = [1.0, float("nan"), 2.0, float("inf"), 3.0, float("-inf")]
        h = histogram(vals, n_bins=3)
        assert sum(h.counts) == 3
        assert abs(h.mean - 2.0) < 1e-10

    def test_frequencies_sum_to_one(self):
        vals = list(range(100))
        h = histogram([float(v) for v in vals], n_bins=10)
        assert abs(sum(h.frequencies) - 1.0) < 1e-10

    def test_skew_kurtosis(self):
        # Symmetric distribution → skew ≈ 0
        vals = [-2.0, -1.0, 0.0, 1.0, 2.0]
        h = histogram(vals, n_bins=5)
        assert abs(h.skew) < 1e-10


# ---------------------------------------------------------------------------
# return_distribution
# ---------------------------------------------------------------------------

class TestReturnDistribution:
    def test_empty(self):
        h = return_distribution([])
        assert h.counts == ()

    def test_too_few(self):
        h = return_distribution([(0.0, 100.0)])
        assert h.counts == ()

    def test_known_returns(self):
        curve = [(0.0, 100.0), (1.0, 110.0), (2.0, 99.0)]
        h = return_distribution(curve, n_bins=2)
        # Returns: 0.1, -0.1
        assert sum(h.counts) == 2
        assert abs(h.mean) < 1e-10  # symmetric


# ---------------------------------------------------------------------------
# pnl_distribution
# ---------------------------------------------------------------------------

class TestPnlDistribution:
    def test_empty(self):
        h = pnl_distribution([])
        assert h.counts == ()

    def test_with_trades(self):
        fills = [
            MockFill(order_side="Buy", price=0.4, size=10, timestamp=1.0),
            MockFill(order_side="Sell", price=0.6, size=10, timestamp=2.0),
        ]
        h = pnl_distribution(fills, n_bins=5)
        assert sum(h.counts) == 1  # one round-trip
        assert h.mean > 0  # profitable trade


# ---------------------------------------------------------------------------
# HeatmapData / monthly_returns_heatmap
# ---------------------------------------------------------------------------

class TestMonthlyReturnsHeatmap:
    def test_empty(self):
        hm = monthly_returns_heatmap([])
        assert hm.values == ()

    def test_single_month(self):
        # Two points in Jan 2025
        from datetime import datetime, timezone
        t1 = datetime(2025, 1, 1, tzinfo=timezone.utc).timestamp()
        t2 = datetime(2025, 1, 31, tzinfo=timezone.utc).timestamp()
        curve = [(t1, 100.0), (t2, 110.0)]
        hm = monthly_returns_heatmap(curve)
        assert len(hm.row_labels) == 1
        assert hm.row_labels[0] == "2025"
        assert len(hm.col_labels) == 12
        # Jan should have ~10% return
        jan_val = hm.values[0][0]
        assert abs(jan_val - 10.0) < 1e-10

    def test_multi_year(self):
        from datetime import datetime, timezone
        t1 = datetime(2024, 6, 1, tzinfo=timezone.utc).timestamp()
        t2 = datetime(2024, 6, 30, tzinfo=timezone.utc).timestamp()
        t3 = datetime(2025, 1, 1, tzinfo=timezone.utc).timestamp()
        t4 = datetime(2025, 1, 31, tzinfo=timezone.utc).timestamp()
        curve = [(t1, 100.0), (t2, 105.0), (t3, 105.0), (t4, 110.0)]
        hm = monthly_returns_heatmap(curve)
        assert len(hm.row_labels) == 2


# ---------------------------------------------------------------------------
# correlation_heatmap
# ---------------------------------------------------------------------------

class TestCorrelationHeatmap:
    def test_empty(self):
        hm = correlation_heatmap({})
        assert hm.values == ()

    def test_single_series(self):
        hm = correlation_heatmap({"a": [1.0, 2.0, 3.0]})
        assert hm.values == ()  # need >= 2

    def test_perfect_correlation(self):
        hm = correlation_heatmap({
            "a": [1.0, 2.0, 3.0, 4.0],
            "b": [2.0, 4.0, 6.0, 8.0],
        })
        assert len(hm.values) == 2
        # Diagonal should be 1.0
        assert abs(hm.values[0][0] - 1.0) < 1e-10
        assert abs(hm.values[1][1] - 1.0) < 1e-10
        # Off-diagonal should be 1.0 (perfect positive)
        assert abs(hm.values[0][1] - 1.0) < 1e-10
        assert abs(hm.values[1][0] - 1.0) < 1e-10

    def test_negative_correlation(self):
        hm = correlation_heatmap({
            "a": [1.0, 2.0, 3.0, 4.0],
            "b": [4.0, 3.0, 2.0, 1.0],
        })
        assert abs(hm.values[0][1] - (-1.0)) < 1e-10

    def test_mismatched_lengths(self):
        hm = correlation_heatmap({
            "a": [1.0, 2.0],
            "b": [1.0, 2.0, 3.0],
        })
        assert hm.values == ()


# ---------------------------------------------------------------------------
# RollingStatData / rolling_stats
# ---------------------------------------------------------------------------

class TestRollingStats:
    def test_empty(self):
        r = rolling_stats([])
        assert r.sharpe.timestamps == ()

    def test_too_few(self):
        curve = [(float(i), 100.0 + i) for i in range(5)]
        r = rolling_stats(curve, window=30)
        assert r.sharpe.timestamps == ()

    def test_enough_points(self):
        curve = [(float(i), 100.0 + i * 0.1) for i in range(50)]
        r = rolling_stats(curve, window=10)
        assert len(r.sharpe.timestamps) > 0
        assert len(r.sortino.timestamps) > 0
        assert r.sharpe.label == "Rolling Sharpe"
        assert r.sortino.label == "Rolling Sortino"
        _all_finite(r.sharpe.values, r.sortino.values)


# ---------------------------------------------------------------------------
# NormalizedEquityData / normalized_equity
# ---------------------------------------------------------------------------

class TestNormalizedEquity:
    def test_empty(self):
        ne = normalized_equity([])
        assert ne.timestamps == ()
        assert ne.bands is None

    def test_normalization(self):
        curve = [(0.0, 100.0), (1.0, 110.0), (2.0, 90.0)]
        ne = normalized_equity(curve, initial_capital=100.0)
        assert abs(ne.equity[0] - 1.0) < 1e-10
        assert abs(ne.equity[1] - 1.1) < 1e-10
        assert abs(ne.equity[2] - 0.9) < 1e-10

    def test_cumulative_pnl(self):
        curve = [(0.0, 100.0), (1.0, 120.0)]
        ne = normalized_equity(curve, initial_capital=100.0)
        assert abs(ne.cumulative_pnl[0]) < 1e-10
        assert abs(ne.cumulative_pnl[1] - 20.0) < 1e-10

    def test_with_bands(self):
        curve = [(float(i), 100.0 + i) for i in range(30)]
        ne = normalized_equity(curve, initial_capital=100.0, include_bands=True, band_window=10)
        assert ne.bands is not None
        assert len(ne.bands.timestamps) > 0

    def test_zero_capital_fallback(self):
        curve = [(0.0, 50.0), (1.0, 100.0)]
        ne = normalized_equity(curve, initial_capital=0.0)
        # Falls back to first equity point
        assert abs(ne.equity[0] - 1.0) < 1e-10


# ---------------------------------------------------------------------------
# TradeScatterData / trade_scatter
# ---------------------------------------------------------------------------

class TestTradeScatter:
    def test_empty(self):
        ts = trade_scatter([])
        assert ts.points == ()
        assert ts.buys == ()
        assert ts.sells == ()

    def test_long_round_trip(self):
        fills = [
            MockFill(order_side="Buy", price=0.4, size=10, timestamp=1.0),
            MockFill(order_side="Sell", price=0.6, size=10, timestamp=2.0),
        ]
        ts = trade_scatter(fills)
        assert len(ts.points) == 2
        assert len(ts.buys) == 1
        assert len(ts.sells) == 1
        # Buy entry
        assert ts.buys[0].side == "buy"
        assert ts.buys[0].price == 0.4
        # Sell exit
        assert ts.sells[0].side == "sell"
        assert ts.sells[0].price == 0.6
        assert ts.sells[0].pnl > 0

    def test_short_round_trip(self):
        fills = [
            MockFill(order_side="Sell", price=0.6, size=10, timestamp=1.0),
            MockFill(order_side="Buy", price=0.4, size=10, timestamp=2.0),
        ]
        ts = trade_scatter(fills)
        assert len(ts.points) == 2
        # Short entry is a sell, exit is a buy
        assert ts.sells[0].side == "sell"
        assert ts.buys[0].side == "buy"

    def test_frozen(self):
        ts = trade_scatter([])
        with pytest.raises(AttributeError):
            ts.points = ()  # type: ignore


# ---------------------------------------------------------------------------
# CalibrationPlotData / calibration_plot
# ---------------------------------------------------------------------------

class TestCalibrationPlot:
    def test_empty_no_outcomes(self):
        c = calibration_plot([], {}, n_bins=5)
        assert c.bin_centers == ()
        assert c.ece == 0.0

    def test_no_matching_trades(self):
        fills = [MockFill(market_id="m1", order_side="Buy", price=0.5)]
        c = calibration_plot(fills, {"m2": 1.0}, n_bins=5)
        assert c.bin_centers == ()

    def test_perfect_calibration(self):
        """Trades at 1.0 with outcome 1.0 → perfect."""
        fills = [MockFill(market_id="m1", order_side="Buy", price=1.0)]
        c = calibration_plot(fills, {"m1": 1.0}, n_bins=10)
        assert c.brier_score == 0.0

    def test_worst_calibration(self):
        """Trades at 1.0 with outcome 0.0 → worst."""
        fills = [MockFill(market_id="m1", order_side="Buy", price=1.0)]
        c = calibration_plot(fills, {"m1": 0.0}, n_bins=10)
        assert abs(c.brier_score - 1.0) < 1e-10

    def test_multiple_trades(self):
        fills = [
            MockFill(market_id="m1", order_side="Buy", price=0.3, timestamp=1.0),
            MockFill(market_id="m2", order_side="Buy", price=0.7, timestamp=2.0),
            MockFill(market_id="m3", order_side="Buy", price=0.9, timestamp=3.0),
        ]
        outcomes = {"m1": 0.0, "m2": 1.0, "m3": 1.0}
        c = calibration_plot(fills, outcomes, n_bins=5)
        assert len(c.bin_centers) == 5
        assert len(c.perfect_line) == 5
        # Perfect line should equal bin_centers
        assert c.perfect_line == c.bin_centers
        _all_finite(c.bin_centers, c.actual_freq, c.perfect_line)


# ---------------------------------------------------------------------------
# PlotBundle / from_backtest
# ---------------------------------------------------------------------------

class TestFromBacktest:
    def _make_result(self, n_points=50, with_trades=True, with_outcomes=False):
        """Create a mock BacktestResult-like object."""

        @dataclass
        class MockResult:
            equity_curve: list = None
            trades: list = None
            initial_capital: float = 1000.0
            outcomes: dict = None

        curve = [(float(i), 1000.0 + i * 10.0) for i in range(n_points)]
        trades = []
        outcomes = None

        if with_trades:
            for j in range(5):
                ts = float(j * 10)
                trades.append(MockFill(order_side="Buy", price=0.4, size=10, timestamp=ts))
                trades.append(MockFill(order_side="Sell", price=0.6, size=10, timestamp=ts + 5))

        if with_outcomes:
            outcomes = {"m1": 1.0}

        return MockResult(equity_curve=curve, trades=trades, outcomes=outcomes)

    def test_basic(self):
        r = self._make_result()
        b = from_backtest(r)
        assert isinstance(b, PlotBundle)
        assert len(b.equity.timestamps) == 50
        assert len(b.underwater.timestamps) == 50
        assert b.return_hist.counts != ()
        assert b.pnl_hist is not None
        assert b.trades.points != ()

    def test_empty_result(self):
        @dataclass
        class EmptyResult:
            equity_curve: list = None
            trades: list = None
            initial_capital: float = 0.0
            outcomes: dict = None

        r = EmptyResult(equity_curve=[], trades=[], outcomes=None)
        b = from_backtest(r)
        assert b.equity.timestamps == ()
        assert b.underwater.timestamps == ()
        assert b.pnl_hist is None
        assert b.monthly_heatmap is None
        assert b.rolling is None
        assert b.calibration is None

    def test_with_outcomes(self):
        r = self._make_result(with_outcomes=True)
        b = from_backtest(r)
        assert b.calibration is not None

    def test_no_rolling_for_short_curves(self):
        r = self._make_result(n_points=5)
        b = from_backtest(r, rolling_window=30)
        assert b.rolling is None

    def test_with_bands(self):
        r = self._make_result(n_points=50)
        b = from_backtest(r, include_bands=True, band_window=10)
        assert b.equity.bands is not None


# ---------------------------------------------------------------------------
# NaN/Inf safety
# ---------------------------------------------------------------------------

class TestNanInfSafety:
    def test_bollinger_nan(self):
        series = [(float(i), float("nan")) for i in range(30)]
        # NaN values will produce NaN in calculations, but _safe_float catches them
        b = bollinger_bands(series, window=5)
        _all_finite(b.middle, b.upper, b.lower, b.bandwidth)

    def test_underwater_zero_peak(self):
        curve = [(0.0, 0.0), (1.0, -10.0)]
        u = underwater_curve(curve)
        _all_finite(u.drawdown_pct, u.peak)

    def test_histogram_all_nan(self):
        h = histogram([float("nan"), float("nan")])
        assert h.counts == ()

    def test_correlation_zero_std(self):
        hm = correlation_heatmap({
            "a": [1.0, 1.0, 1.0],
            "b": [2.0, 4.0, 6.0],
        })
        # "a" has zero std → correlation with "b" should be 0
        assert abs(hm.values[0][1]) < 1e-10


# ---------------------------------------------------------------------------
# Frozen dataclass enforcement
# ---------------------------------------------------------------------------

class TestFrozen:
    def test_all_frozen(self):
        classes = [
            BandData, UnderwaterData, HistogramData, HeatmapData,
            CurveData, NormalizedEquityData, TradeScatterPoint,
            TradeScatterData, CalibrationPlotData, RollingStatData,
            PlotBundle,
        ]
        for cls in classes:
            assert cls.__dataclass_params__.frozen, f"{cls.__name__} is not frozen"
