"""Distribution shift between two FMI weather stations: a linear
hypothesis learned from Helsinki summer days fails on winter days at
Utsjoki Nuorgam.

Purpose
-------
Numerical companion to the glossary entry 'distshift' (distribution
shift).  The Finnish Meteorological Institute (FMI) station Helsinki
Kaisaniemi (fmisid 100971, 60.18 N, 24.94 E) and the northernmost FMI
station, Utsjoki Nuorgam (fmisid 102036, 70.08 N, 27.90 E), record the
minimum and maximum air temperature of each day.  This script downloads
the daily records for June-August 2026 (Kaisaniemi) and December
2025 - February 2026 (Nuorgam) from the FMI open data service and
writes them to CSV files, so the exact numbers behind the figure stay
on record.  Checks pin the downloaded values to the archive, so a
change on the server side is caught rather than silently absorbed.

Each data point is one day: its feature is the minimum temperature of
that day, its label the maximum temperature of the same day.  A linear
hypothesis is learned from Helsinki summer days by minimizing the average
squared error on the training set; its average squared error on a
held-out test set from the same station and season is close to the
training error.  Applied to the winter days at Nuorgam, the same
hypothesis errs dozens of times worse: every winter feature lies far
below every summer feature (the distribution of the feature values
has shifted), and the relation between the two temperatures differs as
well (the conditional distribution has shifted too), so the summer line
extrapolates to above-zero maxima for Arctic winter mornings.

Deterministic: the only randomness is the train/test split of the
Helsinki days, drawn with a fixed seed.  Self-contained: numpy +
matplotlib only (stdlib urllib for the download).

Blocks
------
[B-fetch] Download the daily minimum and maximum temperature for
          Helsinki Kaisaniemi (June-August 2026) and Utsjoki Nuorgam
          (December 2025 - February 2026) from the FMI open data
          service; write them to distshift_helsinki.csv and
          distshift_nuorgam.csv and check them against the archive.
[B-split] Split the Helsinki days into a training set and a test set
          with a fixed seed; check the sizes.
[B-learn] Learn the linear hypothesis by minimizing the average squared
          error on the training set; check that the fitted coefficients
          match the archive.
[B-shift] Average squared error on the training set, the test set, and
          the Nuorgam winter days; check that the test error stays
          close to the training error while the Nuorgam error is
          several times larger, that the two feature ranges do not
          overlap, that the hypothesis overestimates every winter
          maximum, and that a line fitted to the winter days has a
          different slope (the conditional distribution shifted too).

Outputs
-------
distshift_helsinki.csv   : date, tmin, tmax of the Helsinki summer days
distshift_nuorgam.csv    : date, tmin, tmax of the Nuorgam winter days
distshift_train.csv      : tmin, tmax of the Helsinki training set
distshift_test.csv       : tmin, tmax of the Helsinki test set
distshift_deploy.csv     : tmin, tmax of the Nuorgam winter days
distshift_hypothesis.csv : tmin, tmax along the learned hypothesis
distshift_errors.csv     : average squared error on the three sets
distshift.png            : preview (checking only) -- the data points,
                           the learned hypothesis, and the three errors
"""

import re
import urllib.request
from pathlib import Path

import numpy as np
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt

OUT_DIR = Path(__file__).parent

report = []                         # collects (check name, pass/fail) pairs


def check(name, ok):                # records and prints one verification
    report.append((name, bool(ok)))
    print(f"  [{'ok' if ok else 'FAIL'}] {name}")


# ---- [B-fetch] download the daily temperatures from the FMI open data
def fetch_days(fmisid, start, end):
    """Daily (date, tmin, tmax) records of one station, NaN days dropped."""
    url = ("https://opendata.fmi.fi/wfs?service=WFS&version=2.0.0"
           "&request=getFeature"
           "&storedquery_id=fmi::observations::weather::daily::simple"
           f"&fmisid={fmisid}&starttime={start}T00:00:00Z"
           f"&endtime={end}T00:00:00Z&parameters=tmin,tmax")
    with urllib.request.urlopen(url, timeout=120) as resp:
        text = resp.read().decode()
    values = {}                     # date -> {"tmin": v, "tmax": v}
    member = (r"<BsWfs:Time>(\d{4}-\d{2}-\d{2})T[^<]*</BsWfs:Time>\s*"
              r"<BsWfs:ParameterName>(tmin|tmax)</BsWfs:ParameterName>\s*"
              r"<BsWfs:ParameterValue>([^<]+)</BsWfs:ParameterValue>")
    for day, name, value in re.findall(member, text):
        values.setdefault(day, {})[name] = float(value)
    return [(day, rec["tmin"], rec["tmax"])
            for day, rec in sorted(values.items())
            if len(rec) == 2 and np.isfinite([rec["tmin"], rec["tmax"]]).all()]


helsinki = fetch_days(100971, "2026-06-01", "2026-08-31")
nuorgam = fetch_days(102036, "2025-12-01", "2026-02-28")
for fname, records in (("distshift_helsinki.csv", helsinki),
                       ("distshift_nuorgam.csv", nuorgam)):
    with open(OUT_DIR / fname, "w") as f:
        f.write("date,tmin,tmax\n")
        for day, lo, hi in records:
            f.write(f"{day},{lo},{hi}\n")
check("[B-fetch] 92 Helsinki summer days downloaded", len(helsinki) == 92)
check("[B-fetch] 90 Nuorgam winter days downloaded", len(nuorgam) == 90)
check("[B-fetch] Helsinki record matches the archive (Jun 1: 8.1 to 16.3)",
      helsinki[0] == ("2026-06-01", 8.1, 16.3))
check("[B-fetch] Nuorgam record matches the archive (Dec 1: -15.3 to -0.5)",
      nuorgam[0] == ("2025-12-01", -15.3, -0.5))

# ---- [B-split] fixed-seed split of the Helsinki days into train and test
temps_hel = np.array([[lo, hi] for _, lo, hi in helsinki])
temps_nuo = np.array([[lo, hi] for _, lo, hi in nuorgam])
rng = np.random.default_rng(42)
order = rng.permutation(len(temps_hel))
train, test = temps_hel[order[:62]], temps_hel[order[62:]]
check("[B-split] 62 training and 30 test days", len(train) == 62
      and len(test) == 30)

# ---- [B-learn] linear hypothesis minimizing the avg squared training error
A = np.stack([train[:, 0], np.ones(len(train))], 1)
w1, w0 = np.linalg.lstsq(A, train[:, 1], rcond=None)[0]
print(f"  learned hypothesis: tmax = {w1:.3f} * tmin + {w0:.3f}")
check("[B-learn] fitted coefficients match the archive",
      abs(w1 - 0.4806) < 5e-4 and abs(w0 - 14.5005) < 5e-4)


def avg_sq_err(points):
    """Average squared error of the learned hypothesis on (tmin, tmax)."""
    return float(np.mean((w1 * points[:, 0] + w0 - points[:, 1]) ** 2))


# ---- [B-shift] the test error transfers, the Nuorgam error does not
err_train, err_test = avg_sq_err(train), avg_sq_err(test)
err_deploy = avg_sq_err(temps_nuo)
print(f"  avg squared error: train {err_train:.2f}, test {err_test:.2f},"
      f" Nuorgam winter {err_deploy:.2f}")
check("[B-shift] the test error is within a factor 2 of the training error",
      err_test < 2 * err_train)
check("[B-shift] the Nuorgam error exceeds ten times the test error",
      err_deploy > 10 * err_test)
check("[B-shift] the two feature ranges do not overlap",
      temps_nuo[:, 0].max() < train[:, 0].min())
check("[B-shift] every winter maximum is overestimated",
      bool((w1 * temps_nuo[:, 0] + w0 > temps_nuo[:, 1]).all()))
check("[B-shift] the mean winter day reads minimum -19.8, maximum -9.6",
      round(float(temps_nuo[:, 0].mean()), 1) == -19.8
      and round(float(temps_nuo[:, 1].mean()), 1) == -9.6)
w1_win, w0_win = np.linalg.lstsq(
    np.stack([temps_nuo[:, 0], np.ones(len(temps_nuo))], 1),
    temps_nuo[:, 1], rcond=None)[0]
print(f"  line fitted to the winter days: tmax = {w1_win:.3f} * tmin"
      f" + {w0_win:.3f}")
check("[B-shift] the winter days follow a different line (slope 0.766)",
      abs(w1_win - 0.766) < 2e-3)

# ---- outputs: CSVs for the entry's pgfplots figure, preview PNG
header = "tmin,tmax"
np.savetxt(OUT_DIR / "distshift_train.csv", train, delimiter=",",
           header=header, comments="", fmt="%.1f")
np.savetxt(OUT_DIR / "distshift_test.csv", test, delimiter=",",
           header=header, comments="", fmt="%.1f")
np.savetxt(OUT_DIR / "distshift_deploy.csv", temps_nuo, delimiter=",",
           header=header, comments="", fmt="%.1f")
t_line = np.array([temps_nuo[:, 0].min() - 1.0, temps_hel[:, 0].max() + 1.0])
np.savetxt(OUT_DIR / "distshift_hypothesis.csv",
           np.stack([t_line, w1 * t_line + w0], 1), delimiter=",",
           header=header, comments="", fmt="%.3f")
with open(OUT_DIR / "distshift_errors.csv", "w") as f:
    f.write("set,avg_sq_err\n")
    f.write(f"train,{err_train:.2f}\ntest,{err_test:.2f}\n")
    f.write(f"deploy,{err_deploy:.2f}\n")

fig, ax = plt.subplots(figsize=(7, 4.5))
ax.scatter(train[:, 0], train[:, 1], s=22, c="black", marker="o",
           label="training set (Helsinki, summer)")
ax.scatter(test[:, 0], test[:, 1], s=26, facecolors="none",
           edgecolors="black", marker="s",
           label="test set (Helsinki, summer)")
ax.scatter(temps_nuo[:, 0], temps_nuo[:, 1], s=26, facecolors="none",
           edgecolors="black", marker="D",
           label="after deployment (Nuorgam, winter)")
ax.plot(t_line, w1 * t_line + w0, "k-", lw=1.5,
        label="learned hypothesis")
ax.set_xlabel("daily minimum temperature (deg C)")
ax.set_ylabel("daily maximum temperature (deg C)")
ax.set_title("Learned from Helsinki summer days, deployed on Nuorgam winter days\n"
             f"avg squared error: train {err_train:.1f}, test {err_test:.1f},"
             f" deployment {err_deploy:.1f}")
ax.legend(frameon=False, fontsize=9)
fig.tight_layout()
fig.savefig(OUT_DIR / "distshift.png", dpi=120)

failed = [name for name, ok in report if not ok]
print(f"{len(report) - len(failed)}/{len(report)} checks passed"
      + (f"; FAILED: {failed}" if failed else ""))
raise SystemExit(1 if failed else 0)