"""LOBFeatures — scikit-learn Transformer for limit order book microstructure features.

Fetches LOB snapshot data from the Therminal API and engineers features suitable
for market prediction pipelines. Conforms to the scikit-learn estimator API
(BaseEstimator + TransformerMixin) for use in Pipeline, GridSearchCV, etc.

Features include spread, mid-price, book imbalance, depth, and their
aggregations over a lookback window.
"""

from __future__ import annotations

from datetime import datetime, timedelta
from typing import Any

import numpy as np

try:
    from sklearn.base import BaseEstimator, TransformerMixin
    from sklearn.utils.validation import check_is_fitted
except ImportError:
    raise ImportError(
        "scikit-learn is required for therminal.ml. "
        "Install with: pip install therminal-py[ml]"
    )

import logging

from therminal.client import TherminalClient
from therminal.exceptions import NotFoundError
from therminal.ml.cache import get_cached, set_cached
from therminal.ml.features import _to_dates

logger = logging.getLogger(__name__)

# ── Constants ────────────────────────────────────────────────────────────────

_DEFAULT_FEATURES = (
    "spread", "mid", "imbalance", "depth_bid", "depth_ask",
    "weighted_mid", "pressure",
)

_AGGS = ("min", "max", "mean", "std", "last")


# ── LOBFeatures ──────────────────────────────────────────────────────────────

class LOBFeatures(BaseEstimator, TransformerMixin):
    """Extract order book microstructure features from LOB snapshots.

    For each input timestamp, fetches LOB data over a lookback window and
    computes aggregated features (spread dynamics, book imbalance, depth, etc.).

    Parameters
    ----------
    market : str
        Market ticker (e.g., 'KXHIGHNY-26MAR04-B51').
    source : str, default='kalshi'
        Data source — 'kalshi' or 'polymarket'.
    lookback_minutes : int, default=60
        Minutes of LOB history to aggregate for each sample.
    interval : int, default=60
        Snapshot interval in seconds (passed to API).
    levels : int, default=10
        Book depth per side (passed to API).
    features : tuple[str, ...], default=('spread', 'mid', 'imbalance', ...)
        Which base features to extract.
    include_calendar : bool, default=True
        Include calendar features (day_of_year, month, hour — raw + cyclical).
    cache : bool, default=True
        Cache API responses locally.
    client : TherminalClient or None
        Pre-configured client. Created automatically if None.
    base_url : str or None
        API base URL (only used if client is None).
    """

    def __init__(
        self,
        market: str = "",
        source: str = "kalshi",
        lookback_minutes: int = 60,
        interval: int = 60,
        levels: int = 10,
        features: tuple[str, ...] = _DEFAULT_FEATURES,
        include_calendar: bool = True,
        cache: bool = True,
        client: TherminalClient | None = None,
        base_url: str | None = None,
    ):
        self.market = market
        self.source = source
        self.lookback_minutes = lookback_minutes
        self.interval = interval
        self.levels = levels
        self.features = features
        self.include_calendar = include_calendar
        self.cache = cache
        self.client = client
        self.base_url = base_url

    def fit(self, X: Any, y: Any = None) -> LOBFeatures:
        """Determine feature names and prefetch data for the date range.

        Parameters
        ----------
        X : array-like of dates (str, datetime)
            Target timestamps to build features for.
        y : ignored
        """
        dates = _to_dates(X)
        if not dates:
            raise ValueError("X must contain at least one date")

        # Build feature name list
        names: list[str] = []
        for feat in self.features:
            for agg in _AGGS:
                names.append(f"lob_{feat}_{agg}")
        names.append("lob_snapshot_count")

        if self.include_calendar:
            for feat in ("day_of_year", "month", "hour"):
                names.append(feat)
            for feat in ("day_of_year", "month"):
                names.append(f"{feat}_sin")
                names.append(f"{feat}_cos")

        self.feature_names_out_: list[str] = names
        self.n_features_: int = len(names)

        # Store prefetch range
        earliest = min(dates) - timedelta(minutes=self.lookback_minutes)
        self._prefetch_from = earliest.strftime("%Y-%m-%d")
        self._prefetch_to = max(dates).strftime("%Y-%m-%d")

        return self

    def transform(self, X: Any) -> np.ndarray:
        """Generate feature matrix from LOB data.

        Parameters
        ----------
        X : array-like of dates
            Target timestamps.

        Returns
        -------
        np.ndarray of shape (n_samples, n_features)
        """
        check_is_fitted(self, "feature_names_out_")
        dates = _to_dates(X)
        n = len(dates)
        result = np.full((n, self.n_features_), np.nan, dtype=np.float64)

        client = self.client or TherminalClient(
            base_url=self.base_url or "https://api.mostlyright.xyz"
        )

        for i, dt in enumerate(dates):
            row = self._features_for_date(client, dt)
            result[i, :len(row)] = row

        return result

    def get_feature_names_out(self, input_features: Any = None) -> np.ndarray:
        """Return feature names."""
        check_is_fitted(self, "feature_names_out_")
        return np.array(self.feature_names_out_)

    def _features_for_date(self, client: TherminalClient, dt: datetime) -> list[float]:
        """Compute features for a single timestamp."""
        date_str = dt.strftime("%Y-%m-%d")

        cache_params = {
            "market": self.market,
            "source": self.source,
            "date": date_str,
            "interval": self.interval,
            "levels": self.levels,
        }
        snapshots = None
        if self.cache:
            snapshots = get_cached(self.market, "lob", cache_params)

        if snapshots is None:
            try:
                snapshots = client.lob(
                    market=self.market,
                    source=self.source,
                    date=date_str,
                    interval=self.interval,
                    levels=self.levels,
                )
            except NotFoundError:
                snapshots = []  # legitimate: no data for this date
            except Exception as exc:
                logger.warning("LOB fetch failed for %s: %s", date_str, exc)
                snapshots = []

            if self.cache and snapshots:
                set_cached(self.market, "lob", cache_params, snapshots)

        # Filter snapshots to lookback window ending at dt
        target_ns = int(dt.timestamp()) * 1_000_000_000
        lookback_ns = self.lookback_minutes * 60 * 1_000_000_000
        window_start = target_ns - lookback_ns

        window = [s for s in snapshots if window_start <= s.get("ts", 0) <= target_ns]

        vals: list[float] = []

        # Extract base feature time series from window
        feature_series: dict[str, list[float]] = {f: [] for f in self.features}
        for snap in window:
            bids = snap.get("bids", [])
            asks = snap.get("asks", [])
            bid_vol = sum(b.get("qty", 0) for b in bids)
            ask_vol = sum(a.get("qty", 0) for a in asks)
            total_vol = bid_vol + ask_vol

            for feat in self.features:
                if feat == "spread":
                    feature_series[feat].append(float(snap.get("spread", 0)))
                elif feat == "mid":
                    feature_series[feat].append(float(snap.get("mid", 0)))
                elif feat == "imbalance":
                    imb = (bid_vol - ask_vol) / total_vol if total_vol > 0 else 0.0
                    feature_series[feat].append(imb)
                elif feat == "depth_bid":
                    feature_series[feat].append(float(bid_vol))
                elif feat == "depth_ask":
                    feature_series[feat].append(float(ask_vol))
                elif feat == "weighted_mid":
                    if bids and asks:
                        bb = bids[0].get("price", 0)
                        ba = asks[0].get("price", 0)
                        bq = bids[0].get("qty", 1)
                        aq = asks[0].get("qty", 1)
                        wmid = (bb * aq + ba * bq) / (bq + aq) if (bq + aq) > 0 else 0.0
                        feature_series[feat].append(wmid)
                    else:
                        feature_series[feat].append(0.0)
                elif feat == "pressure":
                    top3_bid = sum(b.get("qty", 0) for b in bids[:3])
                    top3_ask = sum(a.get("qty", 0) for a in asks[:3])
                    ratio = top3_bid / top3_ask if top3_ask > 0 else 0.0
                    feature_series[feat].append(ratio)

        # Compute aggregations
        for feat in self.features:
            series = feature_series[feat]
            arr = np.array(series) if series else np.array([0.0])
            for agg in _AGGS:
                if agg == "min":
                    vals.append(float(np.min(arr)))
                elif agg == "max":
                    vals.append(float(np.max(arr)))
                elif agg == "mean":
                    vals.append(float(np.mean(arr)))
                elif agg == "std":
                    vals.append(float(np.std(arr)))
                elif agg == "last":
                    vals.append(float(arr[-1]))

        vals.append(float(len(window)))  # snapshot_count

        # Calendar features
        if self.include_calendar:
            vals.append(float(dt.timetuple().tm_yday))  # day_of_year
            vals.append(float(dt.month))
            vals.append(float(dt.hour))
            # Cyclical encoding
            doy = dt.timetuple().tm_yday
            vals.append(np.sin(2 * np.pi * doy / 365.25))
            vals.append(np.cos(2 * np.pi * doy / 365.25))
            vals.append(np.sin(2 * np.pi * dt.month / 12))
            vals.append(np.cos(2 * np.pi * dt.month / 12))

        return vals
