"""Build the supervised table: weather up to day t -> log(turbidity FNU) at day t+h.

Features only use information up to day t (no future weather), so inference needs no
weather forecast. Caveat: training weather is ERA5 / wave reanalysis, while live use
feeds the real-time Open-Meteo analysis, a small distribution shift.
"""
import numpy as np, pandas as pd
from pathlib import Path
from spots import SPOTS

HORIZONS = (1, 2, 3)  # days ahead (24/48/72 h)


def daily_weather(name, data=Path("data")):
    return daily_weather_from_hourly(pd.read_parquet(data / f"wx_{name}.parquet"))


def daily_weather_from_hourly(w):
    w = w.set_index("time").sort_index()
    w["wave_energy"] = w["wave_height"] ** 2 * w["wave_period"]
    d = pd.DataFrame({
        "hs_max": w["wave_height"].resample("D").max(),
        "hs_mean": w["wave_height"].resample("D").mean(),
        "tp_mean": w["wave_period"].resample("D").mean(),
        "wave_e": w["wave_energy"].resample("D").mean(),
        "wind": w["wind_speed_10m"].resample("D").mean(),
        "wind_max": w["wind_speed_10m"].resample("D").max(),
        "rain": w["precipitation"].resample("D").sum(),
    })
    for c, wins in (("hs_max", (2, 3, 7)), ("wave_e", (3, 7)), ("rain", (3, 7)), ("wind", (3,))):
        for k in wins:
            d[f"{c}_r{k}"] = d[c].rolling(k, min_periods=1).mean()
    return d


def build(h, data=Path("data")):
    rows = []
    for i, (name, (lat, lon, reg)) in enumerate(SPOTS.items()):
        p = data / f"hr_{name}.parquet"
        if not p.exists() or not (data / f"wx_{name}.parquet").exists():
            continue
        hr = pd.read_parquet(p).set_index("date")
        hr = hr[hr["valid_frac"] > 0.2]  # require enough valid pixels in the box
        kd = hr["tur"].groupby(level=0).mean().clip(lower=0.01)
        wx = daily_weather(name, data)
        idx = pd.date_range("2020-01-02", "2026-08-25")
        kdf = np.log(kd.reindex(idx))
        df = wx.reindex(idx).copy()
        df["y"] = kdf.shift(-h)  # target at t+h
        # last known (cloud-free) observation up to t, and its age
        last = kdf.ffill()
        age = kdf.notna().astype(int).groupby(kdf.notna().cumsum()).cumcount()
        df["last_kd"] = last
        df["last_age"] = age
        df["doy_sin"] = np.sin(2 * np.pi * idx.dayofyear / 365.25)
        df["doy_cos"] = np.cos(2 * np.pi * idx.dayofyear / 365.25)
        df["lat"], df["lon"], df["spot"] = lat, lon, i
        df["date"] = idx
        rows.append(df)
    out = pd.concat(rows).dropna(subset=["y", "last_kd"])
    return out.reset_index(drop=True)


def split(df):
    t = df["date"]
    return (df[t < "2025-01-01"], df[(t >= "2025-01-01") & (t < "2026-01-01")], df[t >= "2026-01-01"])