"""WeatherFeatures — scikit-learn Transformer for temperature prediction features.

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

Target variable (y) is the user's responsibility — this transformer only produces X.
"""

from __future__ import annotations

import warnings
from collections.abc import Sequence
from datetime import UTC, 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]"
    )

from therminal.client import TherminalClient
from therminal.ml.cache import get_cached, set_cached

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

_DEFAULT_OMO_AGGS = ("min", "max", "mean", "std")
_DEFAULT_METAR_FIELDS = (
    "temp_f", "dewpoint_f", "relative_humidity",
    "wind_speed_kt", "wind_dir_degrees", "wind_gust_kt",
    "altimeter_inhg", "visibility_miles",
)
_CALENDAR_FEATURES = ("day_of_year", "month", "weekday", "hour")


# ── Helper: parse dates ──────────────────────────────────────────────────────

def _to_dates(X: Any) -> list[datetime]:
    """Coerce input to a list of timezone-aware UTC datetimes.

    Accepts: list[str], list[datetime], numpy array of strings/datetime64,
    pandas DatetimeIndex, or anything iterable of date-like objects.
    """
    dates: list[datetime] = []
    items = np.asarray(X).ravel()

    for item in items:
        if isinstance(item, datetime):
            dt = item if item.tzinfo else item.replace(tzinfo=UTC)
        elif isinstance(item, str):
            raw = item.strip()
            if len(raw) == 10:  # YYYY-MM-DD
                raw += "T12:00:00Z"  # noon UTC default for daily targets
            dt = datetime.fromisoformat(raw.replace("Z", "+00:00"))
        elif hasattr(item, "isoformat"):  # numpy datetime64 or pandas Timestamp
            s = str(item)
            try:
                dt = datetime.fromisoformat(s)
                if dt.tzinfo is None:
                    dt = dt.replace(tzinfo=UTC)
            except ValueError:
                dt = datetime.fromisoformat(s[:19]).replace(tzinfo=UTC)
        else:
            raise ValueError(f"Cannot parse date: {item!r}")
        dates.append(dt)

    return dates


# ── Helper: fetch with cache ────────────────────────────────────────────────

def _fetch_omo(
    client: TherminalClient,
    station: str,
    from_dt: str,
    to_dt: str,
    use_cache: bool,
) -> list[dict[str, Any]]:
    """Fetch 1-minute OMO data, with optional caching."""
    params = {"from": from_dt, "to": to_dt, "resolution": "1min"}
    if use_cache:
        cached = get_cached(station, "omo", params)
        if cached is not None:
            return cached

    # Use max API limit to avoid silent truncation (1-min data = 1440 rows/day)
    data = client.observations(
        station=station, from_date=from_dt, to_date=to_dt,
        resolution="1min", limit=50000,
    )
    if use_cache:
        set_cached(station, "omo", params, data)
    return data  # type: ignore[return-value]


def _fetch_metar(
    client: TherminalClient,
    station: str,
    from_dt: str,
    to_dt: str,
    use_cache: bool,
) -> list[dict[str, Any]]:
    """Fetch hourly METAR data, with optional caching."""
    params = {"from": from_dt, "to": to_dt}
    if use_cache:
        cached = get_cached(station, "metar", params)
        if cached is not None:
            return cached

    data = client.observations(station=station, from_date=from_dt, to_date=to_dt, limit=50000)
    if use_cache:
        set_cached(station, "metar", params, data)
    return data  # type: ignore[return-value]


# ── Feature engineering ──────────────────────────────────────────────────────

def _omo_features_for_date(
    records: list[dict[str, Any]],
    target_dt: datetime,
    lookback_hours: int,
    aggs: tuple[str, ...],
) -> dict[str, float]:
    """Compute rolling aggregations on OMO temp_c for a single target date."""
    cutoff = target_dt - timedelta(hours=lookback_hours)

    temps: list[float] = []
    dewpoints: list[float] = []
    pressures: list[float] = []

    for r in records:
        obs_str = r.get("observed_at", "")
        if not obs_str:
            continue
        obs_dt = datetime.fromisoformat(obs_str.replace("Z", "+00:00"))
        if cutoff <= obs_dt < target_dt:
            if r.get("temp_c") is not None:
                temps.append(float(r["temp_c"]))
            if r.get("dewpoint_c") is not None:
                dewpoints.append(float(r["dewpoint_c"]))
            p1 = r.get("pressure1_inhg")
            if p1 is not None:
                pressures.append(float(p1))

    features: dict[str, float] = {}
    window = f"{lookback_hours}h"

    with warnings.catch_warnings():
        warnings.simplefilter("ignore", RuntimeWarning)  # empty slice expected for no-data windows
        for name, values in [("temp_c", temps), ("dewpoint_c", dewpoints), ("pres_inhg", pressures)]:
            arr = np.array(values) if values else np.array([np.nan])
            for agg in aggs:
                key = f"omo_{name}_{agg}_{window}"
                if agg == "min":
                    features[key] = float(np.nanmin(arr))
                elif agg == "max":
                    features[key] = float(np.nanmax(arr))
                elif agg == "mean":
                    features[key] = float(np.nanmean(arr))
                elif agg == "std":
                    features[key] = float(np.nanstd(arr))
                elif agg == "range":
                    features[key] = float(np.nanmax(arr) - np.nanmin(arr))
                elif agg == "last":
                    features[key] = float(arr[-1]) if len(values) > 0 else np.nan

    features[f"omo_count_{window}"] = float(len(temps))
    return features


def _metar_features_for_date(
    records: list[dict[str, Any]],
    target_dt: datetime,
    lookback_hours: int,
    fields: tuple[str, ...],
) -> dict[str, float]:
    """Extract last-known METAR values before target date + basic stats."""
    cutoff = target_dt - timedelta(hours=lookback_hours)
    window = f"{lookback_hours}h"

    # Collect all METAR obs in the lookback window
    window_obs: list[dict[str, Any]] = []
    for r in records:
        obs_str = r.get("observed_at", "")
        if not obs_str:
            continue
        obs_dt = datetime.fromisoformat(obs_str.replace("Z", "+00:00"))
        if cutoff <= obs_dt < target_dt:
            window_obs.append(r)

    features: dict[str, float] = {}

    with warnings.catch_warnings():
        warnings.simplefilter("ignore", RuntimeWarning)
        for field in fields:
            values = [float(r[field]) for r in window_obs if r.get(field) is not None]
            arr = np.array(values) if values else np.array([np.nan])

            features[f"metar_{field}_last_{window}"] = float(arr[-1]) if values else np.nan
            features[f"metar_{field}_mean_{window}"] = float(np.nanmean(arr))
            features[f"metar_{field}_min_{window}"] = float(np.nanmin(arr))
            features[f"metar_{field}_max_{window}"] = float(np.nanmax(arr))

    features[f"metar_count_{window}"] = float(len(window_obs))
    return features


def _calendar_features(dt: datetime) -> dict[str, float]:
    """Time-based features that don't require API calls."""
    return {
        "day_of_year": float(dt.timetuple().tm_yday),
        "month": float(dt.month),
        "weekday": float(dt.weekday()),
        "hour": float(dt.hour),
        # Cyclical encoding for seasonal patterns
        "day_of_year_sin": float(np.sin(2 * np.pi * dt.timetuple().tm_yday / 365.25)),
        "day_of_year_cos": float(np.cos(2 * np.pi * dt.timetuple().tm_yday / 365.25)),
        "month_sin": float(np.sin(2 * np.pi * dt.month / 12)),
        "month_cos": float(np.cos(2 * np.pi * dt.month / 12)),
    }


# ── The Transformer ─────────────────────────────────────────────────────────

class WeatherFeatures(BaseEstimator, TransformerMixin):
    """scikit-learn Transformer that fetches weather data and engineers features.

    Conforms to the estimator API: ``fit(X) -> self``, ``transform(X) -> np.ndarray``,
    ``get_feature_names_out() -> np.ndarray``. Compatible with ``Pipeline``,
    ``GridSearchCV``, ``cross_val_score``, etc.

    Parameters
    ----------
    station : str
        Therminal station code (e.g., "ATL", "MDW").
    sources : list of str, default ["omo", "metar"]
        Data sources to include. Options: "omo" (1-min ASOS), "metar" (hourly).
    lookback_hours : int, default 24
        Hours of data before each target date to compute features from.
    omo_aggs : tuple of str, default ("min", "max", "mean", "std")
        Aggregations to compute on OMO numeric fields.
        Options: "min", "max", "mean", "std", "range", "last".
    metar_fields : tuple of str or None
        METAR fields to include. None uses defaults (temp, dewpoint, humidity,
        wind, pressure, visibility).
    include_calendar : bool, default True
        Include temporal features (day_of_year, month, cyclical encodings).
    cache : bool, default True
        Cache API responses locally to avoid re-fetching.
    client : TherminalClient or None
        Custom client instance. None creates a default client.
    base_url : str or None
        API base URL. Ignored if client is provided.

    Attributes
    ----------
    feature_names_out_ : np.ndarray
        Feature names after fitting.
    n_features_ : int
        Number of features.

    Examples
    --------
    >>> from therminal.ml import WeatherFeatures
    >>> from sklearn.pipeline import make_pipeline
    >>> from sklearn.preprocessing import StandardScaler
    >>> from sklearn.ensemble import GradientBoostingRegressor
    >>>
    >>> pipe = make_pipeline(
    ...     WeatherFeatures(station="ATL"),
    ...     StandardScaler(),
    ...     GradientBoostingRegressor(),
    ... )
    >>> pipe.fit(dates_train, y_train)
    >>> pipe.predict(dates_test)
    """

    def __init__(
        self,
        station: str = "ATL",
        sources: Sequence[str] = ("omo", "metar"),
        lookback_hours: int = 24,
        omo_aggs: tuple[str, ...] = _DEFAULT_OMO_AGGS,
        metar_fields: tuple[str, ...] | None = None,
        include_calendar: bool = True,
        cache: bool = True,
        client: TherminalClient | None = None,
        base_url: str | None = None,
    ):
        self.station = station
        self.sources = sources
        self.lookback_hours = lookback_hours
        self.omo_aggs = omo_aggs
        self.metar_fields = metar_fields  # None → resolved at use time (sklearn clone compat)
        self.include_calendar = include_calendar
        self.cache = cache
        self.client = client
        self.base_url = base_url

    def _get_client(self) -> TherminalClient:
        if self.client is not None:
            return self.client
        kwargs: dict[str, Any] = {}
        if self.base_url is not None:
            kwargs["base_url"] = self.base_url
        return TherminalClient(**kwargs)

    def fit(self, X: Any, y: Any = None) -> WeatherFeatures:
        """Fit the transformer — determine feature names and prefetch data.

        Parameters
        ----------
        X : array-like of dates
            Target dates (str "YYYY-MM-DD", datetime, or numpy datetime64).
        y : ignored
            Present for sklearn API compatibility.
        """
        dates = _to_dates(X)
        if not dates:
            raise ValueError("X must contain at least one date")

        # Build a single feature row to determine column names
        sample_features = self._build_features_for_date(dates[0], prefetch=False)
        self.feature_names_out_ = np.array(sorted(sample_features.keys()))
        self.n_features_ = len(self.feature_names_out_)

        # Prefetch all data for the full date range (single API call)
        self._prefetch(dates)

        return self

    def transform(self, X: Any) -> np.ndarray:
        """Transform dates into a feature matrix.

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

        Returns
        -------
        np.ndarray of shape (n_samples, n_features)
        """
        check_is_fitted(self, "feature_names_out_")
        dates = _to_dates(X)

        rows: list[list[float]] = []
        for dt in dates:
            feat = self._build_features_for_date(dt, prefetch=True)
            row = [feat.get(name, np.nan) for name in self.feature_names_out_]
            rows.append(row)

        return np.array(rows, dtype=np.float64)

    def get_feature_names_out(self, input_features: Any = None) -> np.ndarray:
        """Return feature names for the output of transform."""
        check_is_fitted(self, "feature_names_out_")
        return self.feature_names_out_

    # ── Internal ──────────────────────────────────────────────────────────

    def _prefetch(self, dates: list[datetime]) -> None:
        """Fetch all data for the date range in bulk (fewer API calls)."""
        if not dates:
            return

        client = self._get_client()
        earliest = min(dates) - timedelta(hours=self.lookback_hours)
        latest = max(dates)
        from_str = earliest.strftime("%Y-%m-%dT%H:%M:%SZ")
        to_str = latest.strftime("%Y-%m-%dT%H:%M:%SZ")

        # Track prefetch range so transform can detect out-of-range dates
        self._prefetch_from = earliest
        self._prefetch_to = latest

        if "omo" in self.sources:
            self._omo_cache = _fetch_omo(
                client, self.station, from_str, to_str, self.cache
            )
        if "metar" in self.sources:
            self._metar_cache = _fetch_metar(
                client, self.station, from_str, to_str, self.cache
            )

    def _build_features_for_date(
        self, dt: datetime, prefetch: bool
    ) -> dict[str, float]:
        """Build the full feature dict for a single target date."""
        client = self._get_client()
        features: dict[str, float] = {}

        from_str = (dt - timedelta(hours=self.lookback_hours)).strftime(
            "%Y-%m-%dT%H:%M:%SZ"
        )
        to_str = dt.strftime("%Y-%m-%dT%H:%M:%SZ")

        # Use prefetched data if available and date is within the prefetch range
        in_range = (
            prefetch
            and hasattr(self, "_prefetch_from")
            and self._prefetch_from <= dt - timedelta(hours=self.lookback_hours)
            and dt <= self._prefetch_to
        )

        if "omo" in self.sources:
            if in_range and hasattr(self, "_omo_cache"):
                records = self._omo_cache
            else:
                records = _fetch_omo(
                    client, self.station, from_str, to_str, self.cache
                )
            features.update(
                _omo_features_for_date(records, dt, self.lookback_hours, self.omo_aggs)
            )

        if "metar" in self.sources:
            if in_range and hasattr(self, "_metar_cache"):
                records = self._metar_cache
            else:
                records = _fetch_metar(
                    client, self.station, from_str, to_str, self.cache
                )
            effective_metar_fields = self.metar_fields if self.metar_fields is not None else _DEFAULT_METAR_FIELDS
            features.update(
                _metar_features_for_date(
                    records, dt, self.lookback_hours, effective_metar_fields
                )
            )

        if self.include_calendar:
            features.update(_calendar_features(dt))

        return features
