"""Flood forecast inference.

Integration contract (see get_forecasts / get_latest_forecasts): callers pass
raw station readings, get back one forecast dict per station x horizon. A
station with a stale, missing, or version-mismatched model transparently
falls back to a simple persistence heuristic instead of raising -- this
module must never crash the caller (e.g. the web API).
"""

import datetime
import logging
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Union

import joblib
import numpy as np
import pandas as pd

from . import features

logger = logging.getLogger(__name__)

# Anchored to the repo root so the API finds trained bundles regardless of CWD.
DEFAULT_MODELS_DIR = Path(__file__).resolve().parents[2] / "models"

DEFAULT_HORIZONS: Tuple[int, ...] = (6, 12, 24)
STALE_AFTER_H = 6.0
HEURISTIC_SIGMA = 0.3
HEURISTIC_VERSION = "heuristic-v1"

# Keyed by (path, mtime) so a retrained model (new mtime) invalidates the old entry.
_MODEL_CACHE: Dict[Tuple[str, float], dict] = {}


def _load_bundle(path: Path) -> dict:
    # joblib.load runs arbitrary pickle code; safe here because `path` is always
    # models/flood_{station}.joblib, an artifact this pipeline's own train.py wrote --
    # never a user- or network-supplied file.
    key = (str(path), path.stat().st_mtime)
    cached = _MODEL_CACHE.get(key)
    if cached is not None:
        return cached
    bundle = joblib.load(path)
    for stale_key in [k for k in _MODEL_CACHE if k[0] == str(path)]:
        del _MODEL_CACHE[stale_key]
    _MODEL_CACHE[key] = bundle
    return bundle


def _readings_to_long_df(readings_by_station: Dict[str, List[dict]]) -> pd.DataFrame:
    rows = []
    for station_code, readings in readings_by_station.items():
        for reading in readings:
            timestamp = reading.get("timestamp")
            if isinstance(timestamp, str):
                timestamp = pd.to_datetime(timestamp)
            rows.append(
                {
                    "timestamp": timestamp,
                    "station_code": station_code,
                    "water_level": reading.get("water_level"),
                    "discharge": reading.get("discharge"),
                }
            )
    if not rows:
        return pd.DataFrame(
            columns=["timestamp", "station_code", "water_level", "discharge"]
        )
    df = pd.DataFrame(rows)
    return df.dropna(subset=["timestamp"])


def _clip_probability(value: float) -> float:
    return float(min(max(value, 0.0), 1.0))


def _sigmoid_probability(predicted_max: float, threshold: float, sigma: float) -> float:
    return 1.0 / (1.0 + np.exp(-(predicted_max - threshold) / sigma))


def _heuristic_forecast(
    station_code: str,
    as_of: pd.Timestamp,
    current_level: float,
    level_t_minus_3: Optional[float],
    warn_thr: float,
    danger_thr: float,
    horizons: Tuple[int, ...],
) -> List[dict]:
    if level_t_minus_3 is None:
        rate = 0.0
    else:
        rate = max(0.0, (current_level - level_t_minus_3) / 3.0)

    results = []
    for horizon_h in horizons:
        predicted_max = max(current_level + rate * horizon_h * 0.7, current_level)
        p_warning = _clip_probability(
            _sigmoid_probability(predicted_max, warn_thr, HEURISTIC_SIGMA)
        )
        p_danger = _clip_probability(
            _sigmoid_probability(predicted_max, danger_thr, HEURISTIC_SIGMA)
        )
        p_danger = min(p_danger, p_warning)
        results.append(
            {
                "station_code": station_code,
                "horizon_hours": horizon_h,
                "p_warning": p_warning,
                "p_danger": p_danger,
                "predicted_max_level": predicted_max,
                "current_level": current_level,
                "as_of": as_of.isoformat(),
                "model_version": HEURISTIC_VERSION,
                "trained_at": None,
                "source": "heuristic",
                "threshold_warning": warn_thr,
                "threshold_danger": danger_thr,
            }
        )
    return results


def _model_forecast(
    station_code: str,
    grid: features.HourlyGrid,
    bundle: dict,
    as_of: pd.Timestamp,
    current_level: float,
    rain: Optional[pd.Series] = None,
    dam: Optional[pd.DataFrame] = None,
) -> List[dict]:
    warn_thr = bundle["thresholds"]["warning"]
    danger_thr = bundle["thresholds"]["danger"]
    # If thresholds changed since this bundle was trained, its classifier heads
    # answer the OLD question (labels for the old levels) while stages/config use
    # the new ones — a silent contradiction on the dashboard. Until a retrain,
    # answer the current question consistently: use the regression + sigma against
    # the configured thresholds and skip the stale heads.
    cfg_warn, cfg_danger = features.get_thresholds(station_code)
    thresholds_stale = (cfg_warn, cfg_danger) != (warn_thr, danger_thr)
    if thresholds_stale:
        logger.warning(
            f"{station_code}: bundle thresholds ({warn_thr}, {danger_thr}) differ from "
            f"configured ({cfg_warn}, {cfg_danger}); using regression-derived probabilities "
            "until the model is retrained"
        )
        warn_thr, danger_thr = cfg_warn, cfg_danger

    feature_row = features.build_features(grid, station_code, rain=rain, dam=dam).loc[
        [as_of]
    ]
    expected_columns = bundle["feature_names"]
    missing = [c for c in expected_columns if c not in feature_row.columns]
    if missing:
        logger.error(
            f"Feature mismatch for {station_code} (missing {missing}); falling back to heuristic"
        )
        return None
    feature_row = feature_row[expected_columns]

    results = []
    for horizon_h in bundle["horizons"]:
        reg = bundle["heads"].get(f"max_{horizon_h}")
        if reg is None:
            results.append(None)
            continue
        raw_prediction = float(reg.predict(feature_row)[0])
        if bundle.get("regression_target") == "rise":
            # v2 bundles predict the rise over the current level
            raw_prediction += current_level
        predicted_max = max(raw_prediction, current_level)
        sigma_h = bundle["sigma"].get(horizon_h, HEURISTIC_SIGMA)

        # Belt-and-braces: the classifier head OR the regression-sigmoid path,
        # whichever is more alarmed. The 2026-08-11 backtest showed a trained
        # classifier staying silent through the 2024 record flood while the
        # regression head tracked it — alerting must never be worse than the
        # regression fallback.
        warn_head = (
            None if thresholds_stale else bundle["heads"].get(f"warn_{horizon_h}")
        )
        p_warning = _sigmoid_probability(predicted_max, warn_thr, sigma_h)
        if warn_head is not None:
            p_warning = max(
                p_warning, float(warn_head.predict_proba(feature_row)[0][1])
            )

        danger_head = (
            None if thresholds_stale else bundle["heads"].get(f"danger_{horizon_h}")
        )
        p_danger = _sigmoid_probability(predicted_max, danger_thr, sigma_h)
        if danger_head is not None:
            p_danger = max(
                p_danger, float(danger_head.predict_proba(feature_row)[0][1])
            )

        p_warning = _clip_probability(p_warning)
        p_danger = min(_clip_probability(p_danger), p_warning)

        row = {
            "station_code": station_code,
            "horizon_hours": horizon_h,
            "p_warning": p_warning,
            "p_danger": p_danger,
            "predicted_max_level": predicted_max,
            "current_level": current_level,
            "as_of": as_of.isoformat(),
            "model_version": bundle["model_version"],
            "trained_at": bundle["trained_at"],
            "source": "model",
            "threshold_warning": warn_thr,
            "threshold_danger": danger_thr,
        }
        stages = features.FLOOD_STAGES.get(station_code)
        if stages:
            # Exceedance probability per official inundation stage, from the
            # regression head and its validation-residual sigma. These are
            # threshold-agnostic, so no retraining is needed to serve them.
            row["stages"] = [
                {
                    "stage": s["stage"],
                    "level": s["level"],
                    "p_exceed": _clip_probability(
                        _sigmoid_probability(predicted_max, s["level"], sigma_h)
                    ),
                }
                for s in stages
            ]
        results.append(row)
    return results


def _forecast_station(
    station_code: str,
    grid: features.HourlyGrid,
    models_dir: Path,
    now: pd.Timestamp,
    horizons: Tuple[int, ...],
    rain: Optional[pd.Series] = None,
    dam: Optional[pd.DataFrame] = None,
) -> List[dict]:
    level_col = (station_code, "water_level")
    if level_col not in grid.observed.columns:
        logger.warning(f"No data for station {station_code}; omitting")
        return []
    observed_level = grid.observed[level_col].dropna()
    if observed_level.empty:
        logger.warning(f"No observed readings for station {station_code}; omitting")
        return []

    as_of = observed_level.index.max()
    current_level = float(observed_level.loc[as_of])
    staleness_h = (pd.Timestamp(now) - as_of).total_seconds() / 3600.0

    warn_thr, danger_thr = features.get_thresholds(station_code)
    t_minus_3 = as_of - pd.Timedelta(hours=3)
    level_t_minus_3 = (
        float(observed_level.loc[t_minus_3])
        if t_minus_3 in observed_level.index
        else None
    )

    bundle_path = models_dir / f"flood_{station_code}.joblib"
    if not bundle_path.exists() or staleness_h > STALE_AFTER_H:
        return _heuristic_forecast(
            station_code,
            as_of,
            current_level,
            level_t_minus_3,
            warn_thr,
            danger_thr,
            horizons,
        )

    bundle = _load_bundle(bundle_path)
    model_results = _model_forecast(
        station_code, grid, bundle, as_of, current_level, rain=rain, dam=dam
    )
    if model_results is None:
        return _heuristic_forecast(
            station_code,
            as_of,
            current_level,
            level_t_minus_3,
            warn_thr,
            danger_thr,
            horizons,
        )

    # Per-horizon heads that were skipped at train time (e.g. too few positives) still
    # need a forecast row -- fall back to the single-horizon heuristic for just that row.
    filled = []
    for horizon_h, row in zip(bundle["horizons"], model_results):
        if row is not None:
            filled.append(row)
        else:
            filled.extend(
                _heuristic_forecast(
                    station_code,
                    as_of,
                    current_level,
                    level_t_minus_3,
                    warn_thr,
                    danger_thr,
                    (horizon_h,),
                )
            )
    return filled


def get_forecasts(
    readings_by_station: Dict[str, List[dict]],
    models_dir: Union[str, Path] = DEFAULT_MODELS_DIR,
    now: Optional[Union[datetime.datetime, str]] = None,
    rain: Optional[pd.Series] = None,
    dam: Optional[pd.DataFrame] = None,
) -> List[dict]:
    """Produce flood forecasts for every station present in `readings_by_station`.

    Each reading dict needs at least {timestamp, water_level, discharge}; extra
    keys are ignored so raw API/DB rows can be passed straight through. At
    least 96 hours of span is required to populate every feature; 336 hours
    (14 days) is recommended.
    """
    models_dir = Path(models_dir)
    if now is None:
        now = datetime.datetime.now()
    now = pd.Timestamp(now)

    df_long = _readings_to_long_df(readings_by_station)
    if df_long.empty:
        return []

    grid = features.make_hourly_grid(df_long)
    results: List[dict] = []
    for station_code in readings_by_station.keys():
        try:
            results.extend(
                _forecast_station(
                    station_code,
                    grid,
                    models_dir,
                    now,
                    DEFAULT_HORIZONS,
                    rain=rain,
                    dam=dam,
                )
            )
        except Exception as error:
            logger.error(f"Forecast failed for station {station_code}: {error}")
    return results


def get_latest_forecasts(
    db_url: Optional[str] = None,
    models_dir: Union[str, Path] = DEFAULT_MODELS_DIR,
    hours: int = 336,
) -> List[dict]:
    """Convenience wrapper for web_api: load the latest window from the DB/API and forecast.

    Raises FileNotFoundError when no trained model bundle exists at all, so the
    API can 503 instead of serving purely heuristic output as if it were a forecast.
    """
    from .data import load_latest

    if not sorted(Path(models_dir).glob("flood_*.joblib")):
        raise FileNotFoundError(f"no trained model bundles in {models_dir}")

    df_long = load_latest(db_url=db_url, hours=hours)
    readings_by_station: Dict[str, List[dict]] = {}
    if not df_long.empty:
        for station_code, group in df_long.groupby("station_code"):
            readings_by_station[station_code] = group[
                ["timestamp", "water_level", "discharge"]
            ].to_dict("records")

    expected_stations = set(features.UPSTREAM_LEADS.keys())
    for missing_station in expected_stations - set(readings_by_station.keys()):
        logger.warning(
            f"No recent data for station {missing_station}; omitting from forecasts"
        )

    # Live rain: trailing days + next-48h forecast. On fetch failure pass an
    # EMPTY series (not None) so rain-trained bundles still find their columns
    # (as NaN) and serve model output instead of tripping the feature guard.
    from .rain import serving_series

    rain = serving_series()
    if rain is None:
        logger.warning("live rain unavailable; rain features will be NaN")
        rain = pd.Series(dtype=float)

    # Recent Mae Ngat reservoir state; same empty-not-None contract so
    # dam-trained bundles keep their columns (NaN) when the DB read fails.
    from . import dam as dam_mod

    try:
        dam = dam_mod.serving_frame(db_url=db_url)
    except Exception as error:
        logger.warning(f"dam serving frame failed: {error}")
        dam = None
    if dam is None:
        logger.warning("dam state unavailable; dam features will be NaN")
        dam = pd.DataFrame()

    return get_forecasts(readings_by_station, models_dir=models_dir, rain=rain, dam=dam)