# SPDX-License-Identifier: MIT
"""YS-EN-0006 · synthetic experiments (the 2023 report's setting, with baselines).

  python exp_synth.py main     # baselines vs ridge vs TCN (report head / direct head) vs GRU
  python exp_synth.py hist     # receptive field: history of 24 … 168 h, with and without calendar
  python exp_synth.py trend    # familiar vs unseen level, global vs per-window normalisation

Train 2021–2022, validate 2023-01..06, test 2023-07..12 (Ontario holidays, incl.
Thanksgiving 2023-10-09, the day the report plotted).  Results: results/synth_*.npz
"""
from __future__ import annotations

import datetime as dt
import sys
from pathlib import Path

import holidays
import numpy as np

import exp
import synth

HERE = Path(__file__).resolve().parent
OUT = HERE / "results"
START = dt.date(2021, 1, 1)
N_DAYS = (dt.date(2024, 1, 1) - START).days
HOL = set(holidays.Canada(subdiv="ON", years=[2021, 2022, 2023]))
SEEDS = (0, 1, 2)
N_SERIES = 2          # independent noise draws, pooled like several rooms


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


TR_HI, VA_HI = hour(dt.date(2023, 1, 1)), hour(dt.date(2023, 7, 1))


def spans(n):
    return ([(0, TR_HI)] * n, [(TR_HI, VA_HI)] * n, [(VA_HI, 24 * N_DAYS)] * n)


def series(level=None, seed0=0):
    out = []
    for k in range(N_SERIES):
        lv = level(k) if level else None
        y, wd, _ = synth.make(START, N_DAYS, HOL, level=lv, seed=seed0 + k)
        out.append({"y": y, "workday": wd, "extra": {}, "start": START, "holidays": HOL})
    return out


def main():
    S = series()
    configs = [
        {"name": "naive_week", "kind": "naive", "naive": "week"},
        {"name": "naive_day", "kind": "naive", "naive": "day"},
        {"name": "naive_sametype", "kind": "naive", "naive": "sametype"},
        {"name": "ridge", "kind": "ridge", "cal": False},
        {"name": "ridge_cal", "kind": "ridge", "cal": True},
        {"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": "gru_cal", "kind": "gru", "cal": True, "chans": ["y", "workday"]},
    ]
    run(S, configs, "synth_main")


def hist():
    S = series()
    configs = []
    for R in (24, 48, 72, 120, 168):
        configs.append({"name": f"tcn_h{R}", "kind": "tcn", "cal": False, "hist": R})
        configs.append({"name": f"tcn_cal_h{R}", "kind": "tcn", "cal": True, "hist": R,
                        "chans": ["y", "workday"]})
    run(S, configs, "synth_hist")


def trend():
    # Training/validation levels wander in [0, 10]; at test time the first series stays in
    # [0, 10] (familiar) and the second moves to [15, 25] (unseen).
    def level(k):
        n = 24 * N_DAYS
        lv = synth.random_level(n, 0, 10, seed=k)
        if k >= 1:
            lv[VA_HI:] = synth.random_level(n, 15, 25, seed=k)[VA_HI:]
        return lv

    S = series(level, seed0=100)
    configs = [
        {"name": "naive_sametype", "kind": "naive", "naive": "sametype"},
        {"name": "ridge_cal_global", "kind": "ridge", "cal": True, "norm": "global"},
        {"name": "ridge_cal_window", "kind": "ridge", "cal": True, "norm": "window"},
        {"name": "tcn_cal_global", "kind": "tcn", "cal": True, "norm": "global",
         "chans": ["y", "workday"]},
        {"name": "tcn_cal_window", "kind": "tcn", "cal": True, "norm": "window",
         "chans": ["y", "workday"]},
    ]
    run(S, configs, "synth_trend")


def run(S, configs, tag):
    OUT.mkdir(exist_ok=True)
    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:.3f}", flush=True)
        res.append(r)
    exp.save(res, OUT / f"{tag}.npz", S)


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