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

Purpose
-------
Quantifies interpretability as simulatability, following the entry (and
its figure): a user who is familiar with linear maps anticipates the
predictions of two learned hypotheses on a test set located OUTSIDE the
region of the labeled examples shown to the user.  The linear hypothesis
h is anticipated almost exactly; the kinked hypothesis h' (identical to h
on the example region, but bending afterwards) defeats the user's
anticipation.  Self-contained (numpy/matplotlib only), deterministic.

Setup
-----
h(x)  = 0.4 x + 2                       (linear everywhere)
h'(x) = h(x)              for x <= 4    (identical on the example region)
      = h(x) - 0.5 (x-4)  for x  > 4    (kink at x = 4)
The user sees labeled examples of each hypothesis at x in [1, 3] (the
region left of the kink), fits a linear map mentally, and anticipates
predictions on the test set x in [5.5, 6.5].

Blocks
------
[B-lin]  Anticipation of h on the test set is exact (max deviation
         < 1e-10): the ML method delivering h is interpretable for the
         user.
[B-kink] Anticipation of h' deviates by at least 0.5 at every test
         point: the method delivering h' is not interpretable, although
         h' agrees with h on all examples the user has seen.
[B-same] On the example region, h and h' are indistinguishable
         (max |h - h'| = 0): interpretability concerns behavior on
         ARBITRARY test sets, beyond the seen examples.

Outputs
-------
interpretability.png : matplotlib preview mirroring the entry's figure
                       (checking only; the entry's figure is schematic
                       TikZ).
"""

import numpy as np
import matplotlib

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

report = []


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


def h(x):
    return 0.4 * x + 2.0


def h_prime(x):
    x = np.asarray(x, dtype=float)
    return h(x) - 0.5 * np.maximum(x - 4.0, 0.0)


x_ex = np.linspace(1.0, 3.0, 6)           # examples shown to the user
x_te = np.linspace(5.5, 6.5, 5)           # test set to anticipate


def user_anticipation(hyp):
    """The user fits a linear map to the labeled examples and
    extrapolates it to the test points."""
    A = np.c_[x_ex, np.ones_like(x_ex)]
    w, *_ = np.linalg.lstsq(A, hyp(x_ex), rcond=None)
    return np.c_[x_te, np.ones_like(x_te)] @ w


# ---------------------------------------------------------------- [B-lin]
dev_lin = float(np.max(np.abs(user_anticipation(h) - h(x_te))))
check(f"[B-lin]  linear hypothesis anticipated exactly (max dev "
      f"{dev_lin:.1e})", dev_lin < 1e-10)

# --------------------------------------------------------------- [B-kink]
dev_kink = float(np.min(np.abs(user_anticipation(h_prime)
                               - h_prime(x_te))))
check(f"[B-kink] kinked hypothesis defeats anticipation (min dev "
      f"{dev_kink:.2f})", dev_kink >= 0.5)

# --------------------------------------------------------------- [B-same]
gap = float(np.max(np.abs(h(x_ex) - h_prime(x_ex))))
check("[B-same] both hypotheses agree on all seen examples", gap == 0.0)

# -------------------------------------------------------------- preview
xs = np.linspace(0.0, 7.0, 200)
fig, ax = plt.subplots(figsize=(4.8, 3.2))
ax.plot(xs, h(xs), "k-", lw=1.4, label="$\\hat{h}$ (linear)")
ax.plot(xs, h_prime(xs), "k--", lw=1.4, label="$\\hat{h}'$ (kinked)")
ax.plot(x_ex, h(x_ex), "ko", ms=5, mfc="none", label="examples seen")
ax.plot(x_te, user_anticipation(h_prime), "kx", ms=6,
        label="user anticipation")
ax.plot(x_te, h_prime(x_te), "k^", ms=5, mfc="none",
        label="$\\hat{h}'$ on test set")
ax.set_xlabel("$x$"), ax.set_ylabel("$y$")
ax.set_title("anticipating predictions: a linear hypothesis against a kinked one",
             fontsize=9)
ax.legend(frameon=False, fontsize=7, loc="upper left")
fig.tight_layout()
fig.savefig("pythondemos/interpretability.png", dpi=110)

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