# SPDX-License-Identifier: MIT
"""YS-EN-0006 · day-ahead forecasting toolkit: windows, baselines, ridge, TCN, GRU.

A series is a dict with hourly arrays of equal length, starting at local midnight:
  y        target (e.g. CO2 ppm), NaN where missing
  workday  1 on working days, 0 on weekends and public holidays
  extra    optional dict of further past-only channels (e.g. PIR activity)
Every forecast uses the 168 hours before the origin and predicts the next 24.
"""
from __future__ import annotations

import copy
import math
import os

import numpy as np

L, H = 168, 24


# ---------------------------------------------------------------- windows
def origins(s, lo, hi, step=1):
    """Origin indices t (forecast starts at hour t) with full, finite history and target,
    lo <= t and t + H <= hi. step=24 gives midnight origins only."""
    y = s["y"]
    ok = np.isfinite(y).astype(np.int32)
    c = np.concatenate([[0], np.cumsum(ok)])
    t = np.arange(max(lo, L), hi - H + 1)
    t = t[t % step == 0]
    full = (c[t + H] - c[t - L]) == L + H
    return t[full]


def calendar(s, t):
    """Future-known features for origin t: workday flag and hour-of-day sin/cos, 24 h each."""
    hours = np.arange(t, t + H)
    hod = 2 * np.pi * (hours % 24) / 24
    return np.concatenate([s["workday"][hours], np.sin(hod), np.cos(hod)]).astype(np.float32)


def past_channels(s, t, chans):
    sl = slice(t - L, t)
    out = []
    for c in chans:
        if c == "y":
            out.append(s["y"][sl])
        elif c == "workday":
            out.append(s["workday"][sl])
        elif c == "hod":
            hod = 2 * np.pi * (np.arange(t - L, t) % 24) / 24
            out += [np.sin(hod), np.cos(hod)]
        else:
            out.append(s["extra"][c][sl])
    return np.stack(out).astype(np.float32)


class Scaler:
    """'global': (y - mu) / sd with train-set mu, sd.  'window': subtract each window's own
    history mean first (a RevIN-style per-window normalisation), divide by the same sd."""

    def __init__(self, mode, mu, sd, extra_stats):
        self.mode, self.mu, self.sd, self.extra = mode, mu, sd, extra_stats

    def offset(self, hist):
        return hist.mean() if self.mode == "window" else self.mu


def fit_scaler(series, idx, mode, chans):
    ys = np.concatenate([s["y"][:hi] for s, (lo, hi) in zip(series, idx)])
    mu, sd = float(np.nanmean(ys)), float(np.nanstd(ys))
    extra = {}
    for c in chans:
        if c not in ("y", "workday", "hod"):
            v = np.concatenate([s["extra"][c][:hi] for s, (lo, hi) in zip(series, idx)])
            extra[c] = (float(np.nanmean(v)), float(np.nanstd(v)) or 1.0)
    return Scaler(mode, mu, sd, extra)


def build(series, spans, chans, scaler, step=1):
    """Stack windows. spans[i] = (lo, hi) hour range for series i.
    Returns X [N, C, L], F [N, 72], Y [N, 24] (scaled), off [N], keys [(series, t)]."""
    X, F, Y, off, keys = [], [], [], [], []
    for i, (s, (lo, hi)) in enumerate(zip(series, spans)):
        for t in origins(s, lo, hi, step):
            p = past_channels(s, t, chans)
            o = scaler.offset(s["y"][t - L:t])
            row = 0
            for c in chans:
                if c == "y":
                    p[row] = (p[row] - o) / scaler.sd
                elif c in scaler.extra:
                    m, d = scaler.extra[c]
                    p[row] = np.nan_to_num((p[row] - m) / d)
                row += 2 if c == "hod" else 1
            X.append(p)
            F.append(calendar(s, t))
            Y.append((s["y"][t:t + H] - o) / scaler.sd)
            off.append(o)
            keys.append((i, int(t)))
    return (np.asarray(X, np.float32), np.asarray(F, np.float32),
            np.asarray(Y, np.float32), np.asarray(off, np.float32), keys)


# ---------------------------------------------------------------- baselines (raw units)
def naive(s, t, kind):
    y, wd = s["y"], s["workday"]
    if kind == "week":                       # same hour, one week earlier
        return y[t - 168:t - 144]
    if kind == "day":                        # same hour yesterday
        return y[t - 24:t]
    if kind in ("sametype_mean", "sametype_median"):   # every same-type day in the history
        want, days = wd[t], []
        for back in range(1, 8):
            a = t - 24 * back
            if wd[a] == want:
                days.append(y[a:a + 24])
        if not days:
            return y[t - 168:t - 144]
        return np.mean(days, 0) if kind == "sametype_mean" else np.median(days, 0)
    # 'sametype': for each future day, the most recent day in the history of the same
    # type (workday / non-workday), using the known future calendar.  Falls back to week.
    out = y[t - 168:t - 144].copy()
    want = wd[t]
    for back in range(1, 8):
        a = t - 24 * back
        if wd[a] == want:
            out = y[a:a + 24]
            break
    return out


def profile_fit(s, lo, hi):
    """Median by (workday, hour of day) over the series' own training span."""
    y, wd = s["y"][lo:hi], s["workday"][lo:hi]
    hod = np.arange(lo, hi) % 24
    prof = np.full((2, 24), np.nanmedian(y))
    for w in (0, 1):
        for h in range(24):
            v = y[(wd == w) & (hod == h)]
            if np.isfinite(v).any():
                prof[w, h] = np.nanmedian(v)
    return prof


def profile_predict(prof, s, t):
    hours = np.arange(t, t + H)
    return prof[s["workday"][hours].astype(int), hours % 24]


# ---------------------------------------------------------------- ridge
def ridge_fit(X, F, Y, alpha, use_cal=True):
    Z = X.reshape(len(X), -1)
    if use_cal:
        Z = np.hstack([Z, F])
    Z = np.hstack([Z, np.ones((len(Z), 1), np.float32)])
    A = Z.T @ Z + alpha * np.eye(Z.shape[1], dtype=np.float32)
    A[-1, -1] -= alpha
    W = np.linalg.solve(A, Z.T @ Y)
    return W, use_cal


def ridge_predict(model, X, F):
    W, use_cal = model
    Z = X.reshape(len(X), -1)
    if use_cal:
        Z = np.hstack([Z, F])
    Z = np.hstack([Z, np.ones((len(Z), 1), np.float32)])
    return Z @ W


# ---------------------------------------------------------------- neural models
def _torch():
    import torch
    import torch.nn as nn
    return torch, nn


def make_model(kind, n_ch, use_cal=True, channels=32, n_layers=8, kernel=2, dropout=0.1,
               head="direct"):
    torch, nn = _torch()

    class CausalBlock(nn.Module):
        def __init__(self, cin, cout, d):
            super().__init__()
            self.pad = (kernel - 1) * d
            self.conv = nn.Conv1d(cin, cout, kernel, dilation=d)
            self.drop = nn.Dropout(dropout)
            self.skip = nn.Conv1d(cin, cout, 1) if cin != cout else nn.Identity()

        def forward(self, x):
            h = self.conv(nn.functional.pad(x, (self.pad, 0)))     # left pad only: causal
            return torch.relu(self.drop(torch.relu(h)) + self.skip(x))

    class TCN(nn.Module):
        def __init__(self):
            super().__init__()
            blocks, c = [], n_ch
            for i in range(n_layers):
                blocks.append(CausalBlock(c, channels, 2 ** i))
                c = channels
            self.net = nn.Sequential(*blocks)
            nf = 72 if use_cal else 0
            if head == "report":
                # The 2023 report's head (its Keras diagrams): a per-timestep dense stack
                # (128 units; 8 in the exogenous model), keep the last 24 positions as the 24
                # forecasts.  Exogenous model: pair each forecast hour with that hour's own
                # future workday flag and apply a per-hour Dense(8) -> Dense(1).  Output k only
                # sees inputs up to position L-24+k, i.e. up to the same hour the day before.
                hid = 8 if use_cal else 128
                self.pt = nn.Sequential(nn.Linear(channels, hid), nn.ReLU(), nn.Linear(hid, 1))
                self.mix = (nn.Sequential(nn.Linear(2, 8), nn.ReLU(), nn.Linear(8, 1))
                            if use_cal else None)
            elif head == "last24":
                # Variant used in an earlier draft: report head, then one MLP over all 24
                # hours plus the 72 calendar features (every hour sees the whole history).
                self.pt = nn.Sequential(nn.Linear(channels, 128), nn.ReLU(), nn.Linear(128, 1))
                self.mix = (nn.Sequential(nn.Linear(H + nf, 64), nn.ReLU(), nn.Linear(64, H))
                            if use_cal else None)
            else:
                self.out = nn.Sequential(nn.Linear(channels * H + nf, 128), nn.ReLU(),
                                         nn.Dropout(dropout), nn.Linear(128, H))

        def forward(self, x, f):
            z = self.net(x)                                        # [N, C, L]
            if head == "report":
                o = self.pt(z[:, :, -H:].transpose(1, 2))                 # [N, 24, 1]
                if self.mix is None:
                    return o.squeeze(-1)
                return self.mix(torch.cat([o, f[:, :H, None]], 2)).squeeze(-1)
            if head == "last24":
                o = self.pt(z[:, :, -H:].transpose(1, 2)).squeeze(-1)
                return self.mix(torch.cat([o, f], 1)) if self.mix is not None else o
            z = z[:, :, -H:].flatten(1)
            return self.out(torch.cat([z, f], 1) if use_cal else z)

    class GRU(nn.Module):
        def __init__(self):
            super().__init__()
            self.rnn = nn.GRU(n_ch, 64, batch_first=True)
            nf = 72 if use_cal else 0
            self.out = nn.Sequential(nn.Linear(64 + nf, 128), nn.ReLU(), nn.Dropout(dropout),
                                     nn.Linear(128, H))

        def forward(self, x, f):
            _, h = self.rnn(x.transpose(1, 2))
            h = h[-1]
            return self.out(torch.cat([h, f], 1) if use_cal else h)

    class Linear(nn.Module):
        """Direct linear map from the flattened window (+ calendar) to 24 outputs."""

        def __init__(self):
            super().__init__()
            self.lin = nn.Linear(n_ch * L + (72 if use_cal else 0), H)

        def forward(self, x, f):
            z = x.flatten(1)
            return self.lin(torch.cat([z, f], 1) if use_cal else z)

    return {"tcn": TCN, "gru": GRU, "linear": Linear}[kind]()


def n_params(model):
    return sum(p.numel() for p in model.parameters())


def train(model, tr, va, seed=0, lr=1e-3, batch=256, max_epochs=40, patience=6, threads=None,
          loss="l1"):
    """tr, va = (X, F, Y). L1 loss (or "l2"); early stopping on validation MAE (scaled units)."""
    torch, nn = _torch()
    torch.manual_seed(seed)
    torch.set_num_threads(threads or int(os.environ.get("YS_THREADS", "4")))
    rng = np.random.default_rng(seed)
    dev = "cuda" if torch.cuda.is_available() else "cpu"
    model.to(dev)
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    Xt, Ft, Yt = (torch.from_numpy(a).to(dev) for a in tr)
    Xv, Fv, Yv = (torch.from_numpy(a).to(dev) for a in va)
    best, best_state, bad, hist = math.inf, None, 0, []
    for ep in range(max_epochs):
        model.train()
        perm = rng.permutation(len(Xt))
        for b in range(0, len(perm), batch):
            j = torch.from_numpy(perm[b:b + batch]).to(dev)
            opt.zero_grad()
            out = model(Xt[j], Ft[j])
            err = (nn.functional.l1_loss(out, Yt[j]) if loss == "l1"
                   else nn.functional.mse_loss(out, Yt[j]))
            err.backward()
            opt.step()
        v = float(nn.functional.l1_loss(predict(model, Xv, Fv, raw=True), Yv))
        hist.append(v)
        if v < best - 1e-5:
            best, best_state, bad = v, copy.deepcopy(model.state_dict()), 0
        else:
            bad += 1
            if bad >= patience:
                break
    model.load_state_dict(best_state)
    model.cpu()
    return model, {"epochs": ep + 1, "best_val": best, "val_curve": hist}


def predict(model, X, F, raw=False):
    torch, _ = _torch()
    model.eval()
    if not raw:
        X, F = torch.from_numpy(X), torch.from_numpy(F)
    with torch.no_grad():
        out = torch.cat([model(X[b:b + 2048], F[b:b + 2048]) for b in range(0, len(X), 2048)])
    return out if raw else out.numpy()


def unscale(pred, off, scaler):
    return pred * scaler.sd + off[:, None]