"""
valset.py -- numerical companion to the entry 'validation set'.

The weather narrative of the 'training set' entry, continued: the same
three training days and the same two hypotheses (straight line,
degree-two polynomial), plus three validation days held back from model
training.  The validation error ranks the two hypotheses opposite to
the training error, the hypothesis with the smaller validation error is
selected (model selection), and three further test days assess the
selected hypothesis.  The demo also evaluates the Hoeffding lower bound
on the validation-set size behind the entry's third figure.
Self-contained (numpy/matplotlib only), deterministic.

Blocks
------
[B-data]     Recreate the three training days of the 'training set'
             entry and learn the two hypotheses from them; check that
             the fitted coefficients match those of trainset.py, so the
             figures of the two entries show the same curves.
[B-valset]   Three validation days held back from model training. The
             validation error of the polynomial far exceeds that of the
             line, reversing the training-error ranking.
[B-modelsel] Select the hypothesis with the smaller validation error
             (the line); evaluate the selected hypothesis on three test
             days that entered neither training nor selection.
[B-bound]    The Hoeffding lower bound ln(2/delta) / (2 Delta^2) on the
             validation-set size: check the entry's example (185 data
             points for Delta = 0.1, delta = 0.05) and print the bound
             for the three curves of the entry's figure.
[B-kfold]    3-fold cross-validation on the six pooled days: each fold
             yields a noisy validation error and their average is the CV
             estimate. Over 300 replicate datasets, the CV estimate
             varies far less than a single split with folds of the same
             size, backing the entry's remedy for scarce data points.

Outputs
-------
pythondemos/valset.png          : preview figure (checking only).
pythondemos/valset_train.csv    : the three training days (x = morning
                                  minimum temperature, y = maximum
                                  daytime temperature)
pythondemos/valset_val.csv      : the three validation days
pythondemos/valset_test.csv     : the three test days
pythondemos/valset_line.csv     : the fitted line on a two-point grid
pythondemos/valset_poly.csv     : the fitted polynomial on a dense grid
"""

import numpy as np
import matplotlib

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

from pathlib import Path

OUT_DIR = Path(__file__).parent

report = []


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


def days(n, gen, lo=-15.0, hi=5.0, w1=0.8, w0=4.0, noise=2.0):
    """n days: morning minimum x and maximum daytime temperature y."""
    x = np.sort(gen.uniform(lo, hi, n))
    y = w0 + w1 * x + gen.normal(0.0, noise, n)
    return x, y


def avg_sqerr(x, y, coeffs):
    return float(np.mean((y - np.polyval(coeffs, x)) ** 2))


# ------------------------------------------------------------------ [B-data]
print("[B-data] the three training days and the two hypotheses of the "
      "'training set' entry")

x, y = days(3, np.random.default_rng(100))
c_line = np.polyfit(x, y, 1)
c_poly = np.polyfit(x, y, 2)
trainerr_line = avg_sqerr(x, y, c_line)
trainerr_poly = avg_sqerr(x, y, c_poly)
print(f"    fitted line: max temperature = {c_line[1]:.2f} + {c_line[0]:.2f}"
      f" * min temperature; training error {trainerr_line:.3f}")
check("the fitted line matches trainset.py (same coefficients)",
      abs(c_line[0] - 0.8688) < 5e-3 and abs(c_line[1] - 4.6785) < 5e-3)
check("the polynomial's training error is zero (up to round-off)",
      trainerr_poly < 1e-10)

# ---------------------------------------------------------------- [B-valset]
print("\n[B-valset] three validation days held back from model training")

xv, yv = days(3, np.random.default_rng(1100))
for xr, yr in zip(xv, yv):
    print(f"    morning minimum {xr:6.1f} C, maximum daytime {yr:6.1f} C")
valerr_line = avg_sqerr(xv, yv, c_line)
valerr_poly = avg_sqerr(xv, yv, c_poly)
print(f"    validation error: line {valerr_line:.2f}, "
      f"polynomial {valerr_poly:.2f}")
check("the validation days are disjoint from the training days",
      len(set(np.round(xv, 6)) & set(np.round(x, 6))) == 0)
check("the validation error reverses the training-error ranking",
      trainerr_poly < trainerr_line and valerr_poly > valerr_line)
check("the polynomial's validation error far exceeds the line's",
      valerr_poly > 3 * valerr_line)

# -------------------------------------------------------------- [B-modelsel]
print("\n[B-modelsel] the validation error selects the line; test days "
      "assess it")

selected = "line" if valerr_line < valerr_poly else "polynomial"
xt, yt = days(3, np.random.default_rng(3100))
testerr_line = avg_sqerr(xt, yt, c_line)
print(f"    selected hypothesis: {selected}; "
      f"test error of the line {testerr_line:.2f}")
check("the smaller validation error selects the line", selected == "line")
check("the test days entered neither training nor selection",
      len(set(np.round(xt, 6)) & set(np.round(np.r_[x, xv], 6))) == 0)

# ----------------------------------------------------------------- [B-bound]
print("\n[B-bound] the Hoeffding lower bound on the validation-set size")


def bound(delta, band):
    return np.log(2.0 / delta) / (2.0 * band ** 2)


print(f"    Delta = 0.1, delta = 0.05: {bound(0.05, 0.1):.1f} "
      f"-> 185 data points")
for delta in (0.01, 0.1, 0.5):
    print(f"    delta = {delta}: bound at Delta = 0.1 is "
          f"{bound(delta, 0.1):.0f}")
check("the entry's example holds: 185 data points suffice",
      int(np.ceil(bound(0.05, 0.1))) == 185)
check("the bound grows as the band narrows (0.05 needs 4x more than 0.1)",
      abs(bound(0.05, 0.05) / bound(0.05, 0.1) - 4.0) < 1e-12)

# ----------------------------------------------------------------- [B-kfold]
print("\n[B-kfold] averaging noisy per-fold validation errors is more "
      "reliable")

X_all, Y_all = np.r_[x, xv], np.r_[y, yv]


def cv3(xs, ys, gen):
    """Per-fold validation errors of 3-fold CV of the linear model."""
    folds = gen.permutation(len(xs)).reshape(3, 2)
    errs = []
    for hold in folds:
        keep = np.ones(len(xs), bool)
        keep[hold] = False
        c = np.polyfit(xs[keep], ys[keep], 1)
        errs.append(avg_sqerr(xs[hold], ys[hold], c))
    return errs


fold_errs = cv3(X_all, Y_all, np.random.default_rng(5))
print(f"    per-fold validation errors: "
      f"{', '.join(f'{v:.2f}' for v in fold_errs)}; "
      f"average {np.mean(fold_errs):.2f}")
check("the per-fold validation errors are noisy (spread over a 3x range)",
      max(fold_errs) > 3 * min(fold_errs))

gen = np.random.default_rng(9)
singles, cv_avgs = [], []
for _ in range(300):
    xr, yr = days(6, gen)
    c = np.polyfit(xr[:4], yr[:4], 1)
    singles.append(avg_sqerr(xr[4:], yr[4:], c))
    cv_avgs.append(float(np.mean(cv3(xr, yr, gen))))
std_single, std_cv = float(np.std(singles)), float(np.std(cv_avgs))
print(f"    over 300 replicate datasets: std of a single split "
      f"{std_single:.1f}, std of the 3-fold CV average {std_cv:.1f}")
check("the CV average varies far less than a single split",
      std_single > 2 * std_cv)

# --------------------------------------------------- outputs: CSVs, preview
xs_line = np.array([-14.0, 5.0])
xs_poly = np.linspace(-14.0, 5.0, 200)


def write_csv(name, cols, header):
    np.savetxt(OUT_DIR / name, np.column_stack(cols), delimiter=",",
               header=header, comments="", fmt="%.4f")


write_csv("valset_train.csv", (x, y), "x,y")
write_csv("valset_val.csv", (xv, yv), "x,y")
write_csv("valset_test.csv", (xt, yt), "x,y")
write_csv("valset_line.csv", (xs_line, np.polyval(c_line, xs_line)),
          "x,yhat")
write_csv("valset_poly.csv", (xs_poly, np.polyval(c_poly, xs_poly)),
          "x,yhat")

fig, axes = plt.subplots(1, 2, figsize=(9.6, 3.6))
ax = axes[0]
ax.plot(xs_poly, np.polyval(c_poly, xs_poly), "--", color="black", lw=1.2,
        label="polynomial")
ax.plot(xs_line, np.polyval(c_line, xs_line), "-", color="black", lw=1.2,
        label="line (selected)")
ax.plot(x, y, "o", ms=6, color="black", label="training set")
ax.plot(xv, yv, "^", ms=7, markerfacecolor="none", markeredgecolor="black",
        label="validation set")
ax.plot(xt, yt, "s", ms=6, markerfacecolor="none", markeredgecolor="black",
        label="test set")
ax.set_xlabel("morning minimum temperature")
ax.set_ylabel("maximum daytime temperature")
ax.set_title("validation days rank the two hypotheses", fontsize=9)
ax.legend(frameon=False, fontsize=7)

ax = axes[1]
band = np.linspace(0.02, 0.5, 200)
for delta, style in ((0.01, "-"), (0.1, "--"), (0.5, ":")):
    ax.semilogy(band, bound(delta, band), style, color="black", lw=1.2,
                label=f"confidence level {1 - delta:g}")
ax.set_xlabel("uncertainty band half-width")
ax.set_ylabel("required validation-set size")
ax.set_title("Hoeffding lower bound on the validation-set size", fontsize=9)
ax.legend(frameon=False, fontsize=8)

fig.tight_layout()
fig.savefig(OUT_DIR / "valset.png", dpi=110)

n_ok = sum(ok for _, ok in report)
print(f"\n{n_ok}/{len(report)} checks pass")
print("wrote valset.png")
if n_ok != len(report):
    raise SystemExit(1)