"""
Real labeled flood-event bookkeeping -- shared between
evaluate_flood_detection.py (which USES these events to flag real
false/true positives after training) and train_spatiotemporal_gnn.py
(which, as of the event-reweighted loss added here, USES the same real
events to upweight training examples near them). Moved out of
evaluate_flood_detection.py into its own module specifically so
train_spatiotemporal_gnn.py can import it too, without creating a
circular import -- evaluate_flood_detection.py already imports FROM
scripts.training.train_spatiotemporal_gnn (shared data-pipeline functions), so
the reverse import would have failed; this module has no dependency on
either script, so both can import it safely.

events_df is identify_flood_events.py's own output (a CSV with columns
station_code, start_date, end_date, threshold) -- both scripts treat
this the same way: real, labeled, human-verifiable flood windows, not
a heuristic re-derived at import time.
"""
from typing import Dict, List

import numpy as np
import pandas as pd


def build_event_lookup(events_df: pd.DataFrame) -> dict:
    """
    Real event date ranges, indexed by station, for fast membership
    checks -- {station_code: [(start_date, end_date, threshold), ...]}.
    A single station can have several real events across the record, so
    this stays a list per station rather than assuming at most one.
    """
    lookup = {}
    for _, row in events_df.iterrows():
        lookup.setdefault(row["station_code"], []).append(
            (pd.Timestamp(row["start_date"]), pd.Timestamp(row["end_date"]), row["threshold"])
        )
    return lookup


def is_real_event_day(station_code: str, date: pd.Timestamp, event_lookup: dict):
    """
    Returns the real threshold if `date` falls within a real event for
    this station, else None -- distinguishing "not an event" from
    "no threshold defined" matters for how the caller should treat it,
    so this returns the threshold itself rather than a bare boolean.
    """
    for start, end, threshold in event_lookup.get(station_code, []):
        if start <= date <= end:
            return threshold
    return None


def days_to_nearest_real_event(station_code: str, date: pd.Timestamp, event_lookup: dict) -> "float":
    """
    Real, signed distance in days from `date` to the nearest real event
    window for this station -- 0 if `date` falls inside one. Built to
    distinguish a real false positive that's a near-miss (a day or two
    from an actual real event) from one with no real event anywhere
    nearby, which are genuinely different situations worth telling
    apart when inspecting false positives directly.

    Returns float("inf") if this station has no real events at all.
    """
    best = float("inf")
    for start, end, _ in event_lookup.get(station_code, []):
        if start <= date <= end:
            return 0.0
        gap = min(abs((date - start).days), abs((date - end).days))
        best = min(best, gap)
    return best


def build_station_thresholds(events_df: pd.DataFrame) -> dict:
    """
    One real threshold per station -- identify_flood_events.py computes
    this once per station (that station's own historical percentile),
    so every real event row for a given station carries the same value;
    this collapses that back down to {station_code: threshold}.

    This threshold is what gets compared against EVERY day's
    prediction, event or not -- comparing only on real event days (and
    treating every other day as an automatic non-flag) would make false
    positives structurally impossible and precision meaningless by
    construction, not just optimistic.
    """
    return events_df.groupby("station_code")["threshold"].first().to_dict()


def build_event_weight_tensor(
    station_codes: List[str], anchor_dates: List[pd.Timestamp], horizons: List[int],
    event_lookup: dict, window_days: int = 5, boost: float = 4.0,
) -> np.ndarray:
    """
    Real, per-(example, node, horizon) sample weight for the SUPERVISED
    discharge quantile loss -- boost where the real TARGET date (anchor
    + horizon - 1 days, not the anchor date itself) falls within
    `window_days` of a real labeled flood event at that station, 1.0
    (neutral) everywhere else.

    Built as the direct, concrete response to the real, measured
    problem this project's own diagnostics confirmed: real flood
    events are rare (125 across the whole 2013-2025 record), so under
    plain (or log-space) pinball loss, the sheer VOLUME of quiet-day
    gradient updates dilutes the few real flood-day updates, even
    though pinball loss's own per-sample asymmetry already weights a
    flood-day underprediction more heavily than a quiet-day
    overprediction of the same size -- that per-sample weighting
    doesn't fix a per-EXAMPLE-COUNT imbalance. Direct, measured
    evidence this project already has: --check-quantile-calibration
    showed q0.95 coverage at 0.990 on low-flow days (should be ~0.95)
    and 0.698 on high-flow days (should also be ~0.95) on the SAME
    checkpoint -- opposite-direction miscalibration consistent with a
    model whose learned quantile function is dominated by the
    overwhelmingly more common quiet-day examples.

    Deliberately restricted to stations that appear in event_lookup at
    all (i.e. have at least one real labeled event) -- every other node
    (virtual/confluence nodes, or real gauges with zero labeled events)
    gets the neutral weight of 1.0 everywhere, without needing any real
    per-example work, since days_to_nearest_real_event would return
    float("inf") for them regardless. This keeps the real cost of
    building this tensor proportional to (stations with events) x
    (real examples) x (real horizons), not the full node count -- with
    ~4,500 nodes but only a handful of real gauges ever having labeled
    events, this is a real, large, deliberate savings, not a premature
    optimization.

    Args:
        station_codes: this run's real, ordered node->station-code list
            (combined["station_codes"]), same ordering Y's node axis
            uses.
        anchor_dates: this run's real anchor dates, same ordering X/Y's
            example axis uses (from prepare_graph_training_windows).
        horizons: this run's real HORIZONS list, same ordering Y's
            horizon axis uses.
        event_lookup: from build_event_lookup(events_df).
        window_days: how many real days from a labeled event's start/
            end the boost still applies -- 5 chosen as a real, moderate
            window (identify_flood_events.py's own docstring notes real
            event durations already run 3-10 days, so this catches the
            real rising/falling limb immediately around a labeled
            event, not just the event's own days).
        boost: the real weight multiplier applied to boosted examples
            (an untouched example keeps weight 1.0). 4.0 chosen as a
            real, moderate value in the 3-5x range this project's own
            discussion landed on -- large enough to meaningfully shift
            gradient share toward flood-adjacent examples, not so large
            it risks destabilizing training the way an aggressive
            reweighting scheme could. Treat as a real, tunable
            hyperparameter, not a fixed constant -- compare against
            boost=1.0 (a no-op, exactly reproducing the unweighted
            loss) as the baseline for any real run using this.
    """
    n_examples = len(anchor_dates)
    n_nodes = len(station_codes)
    n_horizons = len(horizons)
    weight = np.ones((n_examples, n_nodes, n_horizons), dtype=np.float32)

    stations_with_events = [code for code in station_codes if event_lookup.get(code)]
    if not stations_with_events:
        return weight

    for code in stations_with_events:
        node_idx = station_codes.index(code)
        for ex_idx, anchor in enumerate(anchor_dates):
            for h_idx, h in enumerate(horizons):
                target_date = anchor + pd.Timedelta(days=h - 1)
                if days_to_nearest_real_event(code, target_date, event_lookup) <= window_days:
                    weight[ex_idx, node_idx, h_idx] = boost
    return weight