"""Tests for LOBFeatures transformer."""

from __future__ import annotations

from datetime import UTC
from unittest.mock import MagicMock

import numpy as np
import pytest

from therminal.ml.features import _to_dates
from therminal.ml.lob import LOBFeatures

# ── Helpers ──────────────────────────────────────────────────────────────────

def _mock_client(snapshots: list[dict] | None = None) -> MagicMock:
    """Create a mock TherminalClient returning LOB snapshots."""
    client = MagicMock()
    default_snaps = [
        {
            "ts": 1704110340_000_000_000,  # 2024-01-01 11:59:00 UTC (1 min before target)
            "market_id": "KXTEST-B45",
            "best_bid": 45, "best_ask": 55, "spread": 10, "mid": 50.0,
            "bids": [{"price": 45, "qty": 100}, {"price": 44, "qty": 200}],
            "asks": [{"price": 55, "qty": 150}, {"price": 56, "qty": 300}],
        },
        {
            "ts": 1704110400_000_000_000,  # 2024-01-01 12:00:00 UTC (at target)
            "market_id": "KXTEST-B45",
            "best_bid": 46, "best_ask": 54, "spread": 8, "mid": 50.0,
            "bids": [{"price": 46, "qty": 120}, {"price": 45, "qty": 180}],
            "asks": [{"price": 54, "qty": 160}, {"price": 55, "qty": 280}],
        },
    ]
    client.lob.return_value = snapshots if snapshots is not None else default_snaps
    return client


# ── fit / transform ──────────────────────────────────────────────────────────

class TestLOBFeatures:
    def test_fit_returns_self(self):
        client = _mock_client()
        lf = LOBFeatures(market="X", client=client, cache=False)
        result = lf.fit(["2024-01-01T12:00:00Z"])
        assert result is lf

    def test_feature_names(self):
        client = _mock_client()
        lf = LOBFeatures(market="X", client=client, cache=False)
        lf.fit(["2024-01-01T12:00:00Z"])
        names = lf.get_feature_names_out()
        assert "lob_spread_mean" in names
        assert "lob_mid_last" in names
        assert "lob_imbalance_std" in names
        assert "lob_snapshot_count" in names
        assert "day_of_year" in names
        assert "month_sin" in names

    def test_n_features(self):
        client = _mock_client()
        lf = LOBFeatures(market="X", client=client, cache=False)
        lf.fit(["2024-01-01T12:00:00Z"])
        # 7 features × 5 aggs + 1 count + 3 raw cal + 4 cyclical = 43
        assert lf.n_features_ == 43

    def test_n_features_no_calendar(self):
        client = _mock_client()
        lf = LOBFeatures(market="X", client=client, cache=False, include_calendar=False)
        lf.fit(["2024-01-01T12:00:00Z"])
        # 7 features × 5 aggs + 1 count = 36
        assert lf.n_features_ == 36

    def test_transform_shape(self):
        client = _mock_client()
        lf = LOBFeatures(market="X", client=client, cache=False)
        dates = ["2024-01-01T12:00:00Z", "2024-01-02T12:00:00Z"]
        X = lf.fit_transform(dates)
        assert X.shape == (2, lf.n_features_)

    def test_transform_dtype(self):
        client = _mock_client()
        lf = LOBFeatures(market="X", client=client, cache=False)
        X = lf.fit_transform(["2024-01-01T12:00:00Z"])
        assert X.dtype == np.float64

    def test_spread_features(self):
        client = _mock_client()
        lf = LOBFeatures(
            market="X", client=client, cache=False,
            features=("spread",), include_calendar=False,
        )
        X = lf.fit_transform(["2024-01-01T12:00:00Z"])
        names = list(lf.get_feature_names_out())

        # With 2 snapshots: spread = [10, 8]
        min_idx = names.index("lob_spread_min")
        max_idx = names.index("lob_spread_max")
        mean_idx = names.index("lob_spread_mean")
        last_idx = names.index("lob_spread_last")

        assert X[0, min_idx] == 8.0
        assert X[0, max_idx] == 10.0
        assert X[0, mean_idx] == 9.0
        assert X[0, last_idx] == 8.0

    def test_imbalance_features(self):
        client = _mock_client()
        lf = LOBFeatures(
            market="X", client=client, cache=False,
            features=("imbalance",), include_calendar=False,
        )
        X = lf.fit_transform(["2024-01-01T12:00:00Z"])
        names = list(lf.get_feature_names_out())

        last_idx = names.index("lob_imbalance_last")
        # Last snap (at target): bid_vol=120+180=300, ask_vol=160+280=440
        # imbalance = (300-440)/(300+440) = -140/740
        expected = (300 - 440) / (300 + 440)
        np.testing.assert_almost_equal(X[0, last_idx], expected, decimal=4)

    def test_empty_snapshots(self):
        client = _mock_client(snapshots=[])
        lf = LOBFeatures(market="X", client=client, cache=False)
        X = lf.fit_transform(["2024-01-01T12:00:00Z"])
        assert X.shape == (1, lf.n_features_)
        # Should not raise, values should be 0 or NaN

    def test_custom_features(self):
        client = _mock_client()
        lf = LOBFeatures(
            market="X", client=client, cache=False,
            features=("spread", "mid"), include_calendar=False,
        )
        lf.fit(["2024-01-01T12:00:00Z"])
        # 2 features × 5 aggs + 1 count = 11
        assert lf.n_features_ == 11

    def test_source_param(self):
        client = _mock_client()
        lf = LOBFeatures(market="0xabc", source="polymarket", client=client, cache=False)
        lf.fit_transform(["2024-01-01T12:00:00Z"])
        call_kwargs = client.lob.call_args[1]
        assert call_kwargs["source"] == "polymarket"


# ── sklearn integration ──────────────────────────────────────────────────────

class TestSKLearnIntegration:
    def test_get_params(self):
        lf = LOBFeatures(market="X", interval=300)
        params = lf.get_params()
        assert params["market"] == "X"
        assert params["interval"] == 300

    def test_set_params(self):
        lf = LOBFeatures(market="X")
        lf.set_params(market="Y")
        assert lf.market == "Y"

    def test_clone(self):
        from sklearn.base import clone
        lf = LOBFeatures(market="X", interval=120)
        lf2 = clone(lf)
        assert lf2.market == "X"
        assert lf2.interval == 120

    def test_not_fitted_raises(self):
        from sklearn.exceptions import NotFittedError
        lf = LOBFeatures(market="X")
        with pytest.raises(NotFittedError):
            lf.transform(["2024-01-01"])


# ── Date parsing ─────────────────────────────────────────────────────────────

class TestDateParsing:
    def test_string_dates(self):
        dates = _to_dates(["2024-01-01", "2024-06-15T12:00:00Z"])
        assert len(dates) == 2

    def test_datetime_dates(self):
        from datetime import datetime
        dates = _to_dates([datetime(2024, 1, 1, tzinfo=UTC)])
        assert len(dates) == 1

    def test_numpy_array(self):
        dates = _to_dates(np.array(["2024-01-01", "2024-06-15"]))
        assert len(dates) == 2
