"""
classification.py — numerical companion to the glossary entry
'classification'.

Purpose
-------
The entry's first figure is produced here from data rather than drawn by
hand: a training set of data points with known labels and two features,
the linear classifier learned from it, and the two decision regions its decision
boundary cuts the feature space into.  The classifier is built the way the
entry describes: a real-valued hypothesis h(x) = w^T x is trained by
minimizing the average logistic loss over the training set, and the
prediction is obtained by comparing h(x) against the threshold 0.

The two label classes overlap on purpose: data points with nearly the same
feature vector can carry different labels, so no classifier reaches zero
error and the accuracy on a test set stays below one — the entry's point
that some misclassifications are unavoidable.

Self-contained (numpy + matplotlib only), fixed seed.

Blocks
------
[B-data]      A training set and a test set of data points, two features
              each and a label from {-1, 1}.  The two classes overlap.
[B-train]     The real-valued hypothesis h(x) = w^T x, trained by
              minimizing the average logistic loss over the training set.
[B-threshold] The two-step construction: the prediction is 1 where
              h(x) >= 0 and -1 otherwise, so the two decision regions are
              the half-spaces on either side of the hyperplane w^T x = 0,
              and the average 0/1 loss equals one minus the accuracy.
[B-losses]    The 0/1 loss and its two surrogates as functions of the
              margin y*h(x): the logistic loss is convex and
              differentiable, the hinge loss is convex with a kink at
              margin 1, and the hinge loss never falls below the 0/1 loss.
[B-fig]       The figure data: training set, decision boundary and the
              shaded decision region of the label 1, written to
              classification_points.csv, classification_boundary.csv,
              classification_region.csv and classification.png.
[B-test]      The learned classifier judged by its accuracy on the test
              set: clearly above chance, below one because of the overlap.

Outputs
-------
classification_points.csv   : x1, x2, grp for the training set
classification_boundary.csv : the decision boundary of the learned classifier
classification_region.csv   : the decision region of label 1, as a polygon
classification.png          : preview (checking only)
"""

from pathlib import Path

import numpy as np
import matplotlib

OUT_DIR = Path(__file__).parent

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

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}")


rng = np.random.default_rng(20260911)
M_TRAIN = 15                        # training data points per label value
M_TEST = 300                        # test data points in total
MU = np.array([1.1, 1.4])           # center of the label-1 class; the
                                    # label -1 class sits at -MU
SPREAD = 1.1
BOX = 4.5                           # the plotted part of the feature space

# ---------------------------------------------------------------- [B-data]
# A training set and a test set.  Each data point has two features and a
# label from {-1, 1}; the two classes are overlapping clouds centered at
# MU and -MU, so nearby feature vectors can carry different labels.
def draw(m_per_class):
    xp = MU + SPREAD * rng.standard_normal((m_per_class, 2))
    xn = -MU + SPREAD * rng.standard_normal((m_per_class, 2))
    X = np.vstack([xp, xn])
    y = np.hstack([np.ones(m_per_class), -np.ones(m_per_class)])
    return X, y


X_tr, y_tr = draw(M_TRAIN)
X_te, y_te = draw(M_TEST // 2)

check("[B-data] training set holds 15 data points per label value",
      (y_tr == 1).sum() == M_TRAIN and (y_tr == -1).sum() == M_TRAIN)
check("[B-data] training set fits the plotted box",
      np.abs(X_tr).max() < BOX)
# Overlap: even the boundary that separates the two class centers best
# misclassifies part of the test set.
best_err = np.mean(np.sign(X_te @ MU) != y_te)
check("[B-data] the two classes overlap (no error-free boundary)",
      best_err > 0)

# --------------------------------------------------------------- [B-train]
# The real-valued hypothesis h(x) = w^T x, trained by minimizing the
# average logistic loss over the training set (a fixed number of update
# steps).
def avg_logistic_loss(w, X, y):
    return np.mean(np.log1p(np.exp(-y * (X @ w))))


w = np.zeros(2)
loss_path = [avg_logistic_loss(w, X_tr, y_tr)]
for _ in range(3000):
    p = 1.0 / (1.0 + np.exp(-(X_tr @ w)))          # one update step on the
    w -= 0.3 * X_tr.T @ (p - (y_tr + 1) / 2) / len(y_tr)   # average loss
    loss_path.append(avg_logistic_loss(w, X_tr, y_tr))

check("[B-train] the average logistic loss decreased during training",
      loss_path[-1] < loss_path[0])
check("[B-train] the learned weights are finite", np.all(np.isfinite(w)))

# ----------------------------------------------------------- [B-threshold]
# Two-step construction: h(x) quantifies the confidence in the label 1,
# and the prediction compares it against the threshold 0.
h_tr = X_tr @ w
yhat_tr = np.where(h_tr >= 0, 1, -1)

check("[B-threshold] every prediction lies in the label space {-1, 1}",
      set(np.unique(yhat_tr)) <= {-1, 1})
check("[B-threshold] the prediction depends only on the side of the "
      "hyperplane w^T x = 0",
      np.all((yhat_tr == 1) == (h_tr >= 0)))
acc_tr = np.mean(yhat_tr == y_tr)
zeroone_tr = np.mean(yhat_tr != y_tr)
check("[B-threshold] average 0/1 loss = 1 - accuracy on the training set",
      np.isclose(zeroone_tr, 1.0 - acc_tr))

# -------------------------------------------------------------- [B-losses]
# The 0/1 loss and its two surrogates as functions of the margin y*h(x).
t = np.linspace(-2.0, 3.0, 1001)
zeroone = np.where(t <= 0, 1.0, 0.0)
logistic = np.log1p(np.exp(-t))
hinge = np.maximum(0.0, 1.0 - t)

check("[B-losses] the hinge loss never falls below the 0/1 loss",
      np.all(hinge >= zeroone - 1e-12))
check("[B-losses] the logistic loss is convex and strictly decreasing",
      np.all(np.diff(logistic, 2) > -1e-12) and np.all(np.diff(logistic) < 0))
check("[B-losses] the hinge loss is convex",
      np.all(np.diff(hinge, 2) > -1e-12))
left = (hinge[t < 1][-1] - hinge[t < 0.9][-1]) / (t[t < 1][-1] - t[t < 0.9][-1])
right = (hinge[t < 2][-1] - hinge[t < 1.1][-1]) / (t[t < 2][-1] - t[t < 1.1][-1])
check("[B-losses] the hinge loss has a kink at margin 1 "
      "(slope -1 on the left, 0 on the right)",
      np.isclose(left, -1.0, atol=1e-6) and np.isclose(right, 0.0, atol=1e-6))

# ----------------------------------------------------------------- [B-fig]
# Figure data.  The decision boundary w^T x = 0 is a line through the
# origin; the decision region of the label 1 is the half-space above it,
# clipped to the plotted box.
slope = -w[0] / w[1]
check("[B-fig] the boundary leaves the box through its left/right edges",
      w[1] > 0 and abs(slope) * BOX < BOX)

bnd = np.array([[-BOX, slope * -BOX], [BOX, slope * BOX]])
region = np.array([[-BOX, slope * -BOX], [BOX, slope * BOX],
                   [BOX, BOX], [-BOX, BOX]])
check("[B-fig] the interior corners of the region carry the prediction 1",
      np.all(region[2:] @ w > 0))

with open(OUT_DIR / "classification_points.csv", "w") as fh:
    fh.write("x1,x2,grp\n")
    for (a, b), lab in zip(X_tr, y_tr):
        fh.write(f"{a:.4f},{b:.4f},k{1 if lab > 0 else 0}\n")
with open(OUT_DIR / "classification_boundary.csv", "w") as fh:
    fh.write("x1,x2\n")
    for a, b in bnd:
        fh.write(f"{a:.4f},{b:.4f}\n")
with open(OUT_DIR / "classification_region.csv", "w") as fh:
    fh.write("x1,x2\n")
    for a, b in region:
        fh.write(f"{a:.4f},{b:.4f}\n")
print(f"  wrote {OUT_DIR / 'classification_points.csv'}")
print(f"  wrote {OUT_DIR / 'classification_boundary.csv'}")
print(f"  wrote {OUT_DIR / 'classification_region.csv'}")

fig, axes = plt.subplots(1, 2, figsize=(9.2, 4.0))
ax = axes[0]
ax.fill(region[:, 0], region[:, 1], color="0.92",
        label="decision region of label 1")
ax.plot(bnd[:, 0], bnd[:, 1], "k-", lw=1.6, label="decision boundary")
pos, neg = y_tr > 0, y_tr < 0
ax.plot(X_tr[neg, 0], X_tr[neg, 1], "ko", ms=5, label="label -1")
ax.plot(X_tr[pos, 0], X_tr[pos, 1], "ks", mfc="none", ms=6, label="label 1")
ax.set_xlim(-BOX, BOX)
ax.set_ylim(-BOX, BOX)
ax.set_xlabel("feature x1")
ax.set_ylabel("feature x2")
ax.set_title("training set and learned decision regions")
ax.legend(frameon=False, fontsize=8, loc="lower left")

ax = axes[1]
ax.plot(t[t <= 0], zeroone[t <= 0], "k-", lw=1.6, label="0/1 loss")
ax.plot(t[t > 0], zeroone[t > 0], "k-", lw=1.6)
ax.plot(t, logistic, "k--", lw=1.4, label="logistic loss")
ax.plot(t, hinge, "k:", lw=1.8, label="hinge loss")
ax.set_xlabel("margin y h(x)")
ax.set_ylabel("loss")
ax.set_title("0/1 loss and two surrogates")
ax.legend(frameon=False, fontsize=8)

fig.tight_layout()
fig.savefig(OUT_DIR / "classification.png", dpi=110)
print(f"  wrote {OUT_DIR / 'classification.png'}")

# ---------------------------------------------------------------- [B-test]
# The learned classifier judged on the test set.
yhat_te = np.where(X_te @ w >= 0, 1, -1)
acc_te = np.mean(yhat_te == y_te)
print(f"  accuracy on the test set: {acc_te:.3f}")
check("[B-test] test accuracy is clearly above chance", acc_te > 0.8)
check("[B-test] test accuracy stays below one (the classes overlap)",
      acc_te < 1.0)

n_fail = sum(not ok for _, ok in report)
print(f"\n{len(report) - n_fail}/{len(report)} checks pass.")
raise SystemExit(1 if n_fail else 0)