# SPDX-License-Identifier: MIT
"""YS-EN-0006 · real office CO2: VTT SCOTT IAQ, 13 rooms, 2019 (hourly, from prep_vtt.py).

  python exp_vtt.py main    # baselines, ridge, the report's TCN, direct TCN variants, GRU
  python exp_vtt.py hist    # receptive field / history length 24 … 168 h
  python exp_vtt.py loss    # linear models trained with the TCN's L1 loss; TCN on squared error
  python exp_vtt.py check   # training-origin spacing: every hour vs every 3rd hour

Train 2019-02-11 → 08-31 (includes July holidays), validate September, test Oct–Dec
(Finnish public holidays from `holidays`: All Saints' Day, Independence Day, Christmas).
One global model across all 13 rooms.  Gaps of up to MAX_GAP hours in CO2 are linearly
interpolated; any window still containing a gap (history or target) is skipped.
"""
from __future__ import annotations

import datetime as dt
import sys
from pathlib import Path

import holidays
import numpy as np
import pandas as pd

import exp

HERE = Path(__file__).resolve().parent
OUT = HERE / "results"
START = dt.date(2019, 1, 1)
HOL = set(holidays.Finland(years=2019))
SEEDS = (0, 1, 2, 3, 4)
MAX_GAP = 3


def hour(d):
    return 24 * (d - START).days


TR = (hour(dt.date(2019, 2, 11)), hour(dt.date(2019, 9, 1)))
VA = (hour(dt.date(2019, 9, 1)), hour(dt.date(2019, 10, 1)))
TE = (hour(dt.date(2019, 10, 1)), hour(dt.date(2020, 1, 1)))


def fill_short(v, max_gap):
    s = pd.Series(v)
    filled = s.interpolate(limit=max_gap, limit_area="inside")
    # pandas fills the first max_gap hours of longer gaps too; undo that
    isna = s.isna()
    run = isna.groupby((~isna).cumsum()).transform("sum")
    return filled.where(~isna | (run <= max_gap)).to_numpy(np.float32)


def series():
    d = pd.read_csv(HERE / "results" / "vtt_hourly.csv.gz", parse_dates=["time"])
    days = [START + dt.timedelta(k) for k in range(365)]
    wd = np.repeat([x.weekday() < 5 and x not in HOL for x in days], 24).astype(np.float32)
    out = []
    for room, g in d.groupby("room"):
        g = g.sort_values("time")
        assert len(g) == 365 * 24
        out.append({"name": room, "y": fill_short(g["co2"].to_numpy(float), MAX_GAP),
                    "workday": wd, "extra": {"pir": fill_short(g["pir"].to_numpy(float), MAX_GAP)},
                    "start": START, "holidays": HOL})
    return out


def spans(n):
    return [TR] * n, [VA] * n, [TE] * n


def main():
    configs = [
        {"name": "naive_week", "kind": "naive", "naive": "week"},
        {"name": "naive_day", "kind": "naive", "naive": "day"},
        {"name": "naive_sametype", "kind": "naive", "naive": "sametype"},
        {"name": "sametype_mean", "kind": "naive", "naive": "sametype_mean"},
        {"name": "profile", "kind": "profile"},
        {"name": "ridge", "kind": "ridge", "cal": False},
        {"name": "ridge_cal", "kind": "ridge", "cal": True},
        {"name": "ridge_cal_window", "kind": "ridge", "cal": True, "norm": "window"},
        {"name": "tcn_report", "kind": "tcn", "head": "report", "cal": False},
        {"name": "tcn_report_cal", "kind": "tcn", "head": "report", "cal": True,
         "chans": ["y", "workday"]},
        {"name": "tcn", "kind": "tcn", "cal": False},
        {"name": "tcn_cal", "kind": "tcn", "cal": True, "chans": ["y", "workday"]},
        {"name": "tcn_cal_window", "kind": "tcn", "cal": True, "norm": "window",
         "chans": ["y", "workday"]},
        {"name": "tcn_cal_pir", "kind": "tcn", "cal": True, "chans": ["y", "workday", "pir"]},
        {"name": "gru_cal", "kind": "gru", "cal": True, "chans": ["y", "workday"]},
    ]
    run(configs, "vtt_main")


def hist():
    configs = [{"name": f"tcn_cal_h{R}", "kind": "tcn", "cal": True, "hist": R,
                "chans": ["y", "workday"]} for R in (24, 48, 72, 120, 168)]
    configs += [{"name": f"ridge_cal_h{R}", "kind": "ridge", "cal": True, "hist": R}
                for R in (24, 48, 72, 120, 168)]
    run(configs, "vtt_hist")


def loss():
    """Is it the convolutions or the loss?  Linear models trained like the TCN (L1, Adam,
    early stopping), the TCN trained on squared error, and a median-of-days baseline."""
    configs = [
        {"name": "sametype_median", "kind": "naive", "naive": "sametype_median"},
        {"name": "linear_l1", "kind": "linear", "cal": False},
        {"name": "linear_l1_cal", "kind": "linear", "cal": True, "chans": ["y", "workday"]},
        {"name": "linear_l2_cal", "kind": "linear", "cal": True, "chans": ["y", "workday"],
         "loss": "l2"},
        {"name": "tcn_cal_l2", "kind": "tcn", "cal": True, "chans": ["y", "workday"], "loss": "l2"},
    ]
    run(configs, "vtt_loss")


def check():
    """Training on every 3rd hourly origin (the default) vs every origin."""
    configs = [{"name": f"tcn_cal_step{k}", "kind": "tcn", "cal": True, "train_step": k,
                "chans": ["y", "workday"]} for k in (1, 3)]
    run(configs, "vtt_check", seeds=(0, 1, 2))


def run(configs, tag, seeds=SEEDS):
    OUT.mkdir(exist_ok=True)
    S = series()
    tr, va, te = spans(len(S))
    res = []
    for c in configs:
        r = exp.run_config(c, S, tr, va, te, seeds=seeds)
        mae = np.abs(r["pred"] - r["truth"][None]).mean()
        print(f"{tag} {c['name']:18s} MAE {mae:.2f} ppm", flush=True)
        res.append(r)
    exp.save(res, OUT / f"{tag}.npz", S)


if __name__ == "__main__":
    {"main": main, "hist": hist, "loss": loss, "check": check}[sys.argv[1]]()