"""Three measured days, a degree-five polynomial that interpolates them,
and what happens when each day is replaced by copies that lie within the
accuracy of the thermometer.

Purpose
-------
Numerical companion to the glossary entry 'dataaug'.  The GeoSphere
Austria weather station Krems (station id 3805) records the maximum air
temperature of each day.  This script downloads the 2024 records from
the GeoSphere data hub (dataset klima-v2-1d) and writes them to
dataaug_maxtemp.csv.

The learning task pairs consecutive days: the feature of a data point is
the maximum temperature measured today, its label is the maximum
temperature measured tomorrow.  Three days far apart in the year form
the training set, so the three features are far apart as well.

Three data points leave a polynomial of degree five underdetermined: it
has six coefficients, so infinitely many such polynomials pass through
all three and every one of them has zero training error.  ERM delivers any of them and the training error cannot choose.
The curve drawn is one of them, built by adding to the smoothest
interpolant a degree-five polynomial that vanishes at the three measured
features, scaled so that it reaches 15 degrees Celsius.  That curve
leaves the range the three labels occupy by far, which is what
overfitting looks like here.

A thermometer reports the air temperature only to within its accuracy,
assumed here to be 0.5 degrees Celsius.  Both the feature and the label
of a data point are readings of that instrument, so moving each of them
anywhere inside that tolerance gives a day that is just as consistent
with what was measured.  Replacing each of the three days by copies
drawn that way is the augmented training set.  The copies carry many
distinct feature values, so the feature matrix of the augmented set has
full column rank and ERM on it has a unique
solution.  A curve that is steep where the copies lie pays for it, since
a copy displaced in the feature direction then misses its label, so the
steep interpolants are the ones the augmented criterion rejects.

The demo checks the claims the entry makes.  (1) The three training days
lie in different seasons and their features span a wide range.  (2) Both
the smoothest interpolant and the drawn one have zero training error on
the three days, so the training error does not separate them.  (3) The
drawn one leaves the range of the three labels by more than ten degrees
between them.  (4) Every copy lies within the assumed accuracy of the
day it came from, and the augmented feature matrix has full column rank,
so the augmented problem has a unique solution.  (5) That solution stays
within one degree of the range of the three labels, an excursion smaller
by more than a factor of ten.  (6) On the day pairs of 2024 that were
held out and whose feature lies between the smallest and the largest
training feature, it has a far smaller average squared error than the
drawn interpolant, and it lands close to the straight line through the
three days.

Deterministic: the data are a fixed archive year and the copies are
drawn with a fixed seed.  Self-contained: numpy + matplotlib only
(stdlib urllib for the download).

Blocks
------
[B-data]      Download the maximum temperature of every day of 2024 and
              pair each day with the next one.
[B-three]     Pick the three training days, one per season, and report
              their features and labels.
[B-overfit]   Two of the infinitely many ERM solutions: both with zero
              training error, differing widely between the three days.
[B-augment]   Replace each day by copies inside the accuracy of the
              thermometer and fit the same degree-five model to them.
[B-compare]   The fits on the held-out day pairs whose feature lies
              between the smallest and the largest training feature,
              and the straight line through the three days.
[B-fig]       The preview figure: the picture on the left, the held-out
              errors on the right.

Outputs
-------
dataaug_maxtemp.csv : date and maximum temperature, 366 days
dataaug_points.csv  : the three training days, feature and label
dataaug_copies.csv  : the copies drawn inside the sensor accuracy
dataaug_curves.csv  : the two fitted curves on a grid of features
dataaug_zoom.csv    : copies of the middle day, relative to that day
dataaug_valerr.csv  : held-out error of the three fits
dataaug.png         : preview figure
"""
# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy==2.3.4", "matplotlib==3.10.7"]
# ///


import json
import urllib.request
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt                        # noqa: E402
import numpy as np                                     # noqa: E402

OUT_DIR = Path(__file__).parent
DEGREE = 5
TOL = 0.5                 # assumed accuracy of the thermometer, in Celsius
NR_COPIES = 200           # copies drawn per training day
SWING = 15.0              # amplitude, in Celsius, of the drawn ERM solution
RNG = np.random.default_rng(20240101)

FAILED = []


def check(name, ok):
    FAILED.append(name) if not ok else None
    print(f"  [{'ok' if ok else 'FAIL'}] {name}")


def design(x, lo, hi):
    """Powers of the feature, rescaled to [-1, 1] so the fit is stable."""
    u = 2.0 * (x - lo) / (hi - lo) - 1.0
    return np.stack([u ** k for k in range(DEGREE + 1)], axis=1)


# ---- [B-data] the maximum temperature of every day of 2024
URL = ("https://dataset.api.hub.geosphere.at/v1/station/historical/"
       "klima-v2-1d?parameters=tlmax&station_ids=3805"
       "&start=2024-01-01&end=2024-12-31")
with urllib.request.urlopen(URL, timeout=180) as resp:
    payload = json.load(resp)
tmax = np.array(
    payload["features"][0]["properties"]["parameters"]["tlmax"]["data"],
    dtype=float)
dates = [t[:10] for t in payload["timestamps"]]
with open(OUT_DIR / "dataaug_maxtemp.csv", "w") as f:
    f.write("date,tlmax\n")
    for day, v in zip(dates, tmax):
        f.write(f"{day},{v:g}\n")
# today's maximum predicts tomorrow's maximum
FEAT, LAB = tmax[:-1], tmax[1:]
print(f"[B-data] {len(FEAT)} day pairs of 2024 at Krems an der Donau")

# ---- [B-three] three days far apart in the year
IDX = [20, 105, 200]                      # late January, mid April, late July
xt, yt = FEAT[IDX], LAB[IDX]
for i, d in zip(IDX, [dates[i] for i in IDX]):
    print(f"[B-three] {d}: today {FEAT[i]:.1f} C, tomorrow {LAB[i]:.1f} C")
check("[B-three] the three training days lie in different seasons",
      len({dates[i][5:7] for i in IDX}) == 3)
check("[B-three] their features span more than twenty degrees",
      xt.max() - xt.min() > 20.0)

LO, HI = FEAT.min(), FEAT.max()
GRID = np.linspace(xt.min(), xt.max(), 400)
c_lin = np.polyfit(xt, yt, 1)
LINE = np.polyval(c_lin, GRID)

# ---- [B-overfit] one of the infinitely many ERM solutions
A = design(xt, LO, HI)
c_smooth = np.linalg.lstsq(A, yt, rcond=None)[0]   # smallest coefficients
# every polynomial that vanishes at the three features can be added to it
null = np.linalg.svd(A)[2][len(IDX):]              # directions spanning it
wiggle = design(GRID, LO, HI) @ null[0]
c_int = c_smooth + (SWING / np.max(np.abs(wiggle))) * null[0]
fit_smooth = design(GRID, LO, HI) @ c_smooth
fit_int = design(GRID, LO, HI) @ c_int
train_smooth = np.mean((A @ c_smooth - yt) ** 2)
train_int = np.mean((A @ c_int - yt) ** 2)
exc_int = max(fit_int.max() - yt.max(), yt.min() - fit_int.min())
print(f"[B-overfit] {null.shape[0]} independent directions leave the three "
      f"labels untouched, so ERM has infinitely many solutions")
print(f"[B-overfit] training error of the smoothest {train_smooth:.2e} and "
      f"of the drawn one {train_int:.2e}")
gap = np.max(np.abs(fit_int - fit_smooth))
MIDX = 12.0                       # a mild spring day, between the measured ones
mid = design(np.array([MIDX]), LO, HI)
span = np.max(np.abs(mid @ null.T))
print(f"[B-overfit] between the measured days the two solutions differ by "
      f"up to {gap:.1f} C")
print(f"[B-overfit] at a today-maximum of {MIDX:.0f} C the smoothest "
      f"predicts {float(mid @ c_smooth):.1f} C and the drawn one "
      f"{float(mid @ c_int):.1f} C")
check("[B-overfit] six coefficients against three data points, so ERM "
      "has infinitely many solutions", null.shape[0] == DEGREE + 1 - len(IDX))
check("[B-overfit] the training error does not separate them: both are "
      "zero", train_smooth < 1e-12 and train_int < 1e-12)
check("[B-overfit] yet they differ by more than ten degrees between the "
      "measured days", gap > 10.0)
check("[B-overfit] at a feature between the measured ones the solutions "
      "take every real value, since a direction that leaves the three "
      "labels untouched is nonzero there", span > 1e-6)

# ---- [B-augment] copies inside the accuracy of the thermometer
xa = np.repeat(xt, NR_COPIES) + RNG.uniform(-TOL, TOL, NR_COPIES * len(IDX))
ya = np.repeat(yt, NR_COPIES) + RNG.uniform(-TOL, TOL, NR_COPIES * len(IDX))
check("[B-augment] every copy lies within the assumed accuracy of the "
      "day it came from",
      np.all(np.abs(xa - np.repeat(xt, NR_COPIES)) <= TOL)
      and np.all(np.abs(ya - np.repeat(yt, NR_COPIES)) <= TOL))
Aa = design(xa, LO, HI)
c_aug = np.linalg.lstsq(Aa, ya, rcond=None)[0]
fit_aug = design(GRID, LO, HI) @ c_aug
exc_aug = max(fit_aug.max() - yt.max(), yt.min() - fit_aug.min())
print(f"[B-augment] the augmented feature matrix is {Aa.shape[0]} by "
      f"{Aa.shape[1]} of rank {np.linalg.matrix_rank(Aa)}")
dev_aug = np.max(np.abs(fit_aug - LINE))
dev_int = np.max(np.abs(fit_int - LINE))
print(f"[B-augment] the augmented fit leaves the range of the three "
      f"labels by {exc_aug:.2f} C")
print(f"[B-augment] it stays within {dev_aug:.1f} C of the straight line "
      f"through the three days, against {dev_int:.1f} C for the drawn "
      f"solution")
check("[B-augment] the augmented feature matrix has full column rank, so "
      "the augmented problem has a unique solution",
      np.linalg.matrix_rank(Aa) == DEGREE + 1)
check("[B-augment] its excursion beyond the range of the three labels "
      "stays below one degree", exc_aug < 1.0)
check("[B-augment] it departs from the straight line through the three "
      "days by less than the drawn solution does",
      dev_aug < dev_int)

# ---- [B-compare] the two fits on the days that were held out
# held out, and inside the feature range the three training days cover
inside = (FEAT >= xt.min()) & (FEAT <= xt.max())
inside[IDX] = False
rest = np.flatnonzero(inside)


def val(coef, poly1d=False):
    pred = (np.polyval(coef, FEAT[rest]) if poly1d
            else design(FEAT[rest], LO, HI) @ coef)
    return float(np.mean((pred - LAB[rest]) ** 2))


v_int, v_aug, v_lin = val(c_int), val(c_aug), val(c_lin, True)
v_smooth = val(c_smooth)
print(f"[B-compare] held-out error on {len(rest)} day pairs inside the "
      f"training range: drawn solution {v_int:.1f}, smoothest solution "
      f"{v_smooth:.1f}, augmented fit {v_aug:.1f}, straight line "
      f"{v_lin:.1f}")
check("[B-compare] the two ERM solutions differ widely on the held-out "
      "days although both have zero training error",
      abs(v_int - v_smooth) > 20.0)
check("[B-compare] the augmented fit beats the drawn solution",
      v_aug < v_int / 2.0)
check("[B-compare] and lands close to the straight line through the "
      "three days", abs(v_aug - v_lin) < 0.3 * v_lin)
with open(OUT_DIR / "dataaug_valerr.csv", "w") as f:
    f.write("fit,valerr\n")
    f.write(f"drawn,{v_int:.4f}\n")
    f.write(f"smoothest,{v_smooth:.4f}\n")
    f.write(f"augmented,{v_aug:.4f}\n")
    f.write(f"line,{v_lin:.4f}\n")

# ---- the figure data
with open(OUT_DIR / "dataaug_points.csv", "w") as f:
    f.write("x,y\n")
    for a, b in zip(xt, yt):
        f.write(f"{a:.4f},{b:.4f}\n")
with open(OUT_DIR / "dataaug_copies.csv", "w") as f:
    f.write("x,y\n")
    for a, b in zip(xa[::25], ya[::25]):       # a readable subset
        f.write(f"{a:.4f},{b:.4f}\n")
one = slice(NR_COPIES, 2 * NR_COPIES)          # the middle measured day
with open(OUT_DIR / "dataaug_zoom.csv", "w") as f:
    f.write("dx,dy\n")
    for a, b in zip(xa[one][:28], ya[one][:28]):
        f.write(f"{a - xt[1]:.4f},{b - yt[1]:.4f}\n")
with open(OUT_DIR / "dataaug_curves.csv", "w") as f:
    f.write("x,interpolant,augmented\n")
    for a, b, c in zip(GRID, fit_int, fit_aug):
        f.write(f"{a:.4f},{b:.4f},{c:.4f}\n")

# ---- [B-fig] the preview figure
fig, (axP, axV) = plt.subplots(1, 2, figsize=(9.6, 3.8))
axP.plot(GRID, fit_int, "k-", lw=1.6,
         label="degree-5 fit on the three days")
axP.plot(GRID, fit_aug, "k--", lw=1.6,
         label="degree-5 fit on the augmented set")
axP.plot(xa[::25], ya[::25], "s", mfc="none", mec="0.45", ms=4,
         label="copy inside the sensor accuracy")
axP.plot(xt, yt, "ko", ms=7, label="measured day")
axP.set_xlabel("maximum temperature today in $^\\circ$C")
axP.set_ylabel("maximum temperature tomorrow in $^\\circ$C")
axP.set_title("Three measured days and two fitted hypotheses",
              fontsize=10)
axP.set_ylim(yt.min() - 14, yt.max() + 14)
axP.legend(frameon=False, fontsize=7, loc="lower right")
names = ["drawn ERM\nsolution", "smoothest ERM\nsolution",
         "fit on the\naugmented set", "straight line\non three days"]
axV.bar(names, [v_int, v_smooth, v_aug, v_lin], color="0.55",
        edgecolor="black")
axV.set_xlabel("hypothesis compared")
axV.set_ylabel("average squared error in $(^\\circ$C$)^2$")
axV.set_title(f"Error on the {len(rest)} day pairs held out", fontsize=10)
axV.set_yscale("log")
axV.tick_params(axis="x", labelsize=7)
fig.tight_layout()
fig.savefig(OUT_DIR / "dataaug.png", dpi=110)
check("[B-fig] the preview figure was written",
      (OUT_DIR / "dataaug.png").exists())

print()
if FAILED:
    print(f"{len(FAILED)} check(s) FAILED: " + "; ".join(FAILED))
    raise SystemExit(1)
print("all checks passed")