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

Purpose
-------
Verifies the entry's two definitions and its ML claim on concrete objects:
segment membership for a convex set, the chord inequality and epigraph
convexity for the average squared error loss of linear regression, and the
local-equals-global-minimum property via gradient descent from many random
initializations.  Self-contained (numpy only), fixed seed.

Blocks
------
[B-set]    The ellipse set C = {w : w^T diag(1, 4) w <= 4} contains
           beta w + (1-beta) w' for 500 random pairs w, w' in C and
           21 values beta in [0, 1].
[B-fn]     The linear-regression objective f(w) = (1/m)||y - X w||^2
           (m = 10, d = 2, seed 0) satisfies the chord inequality
           f(beta w + (1-beta) w') <= beta f(w) + (1-beta) f(w')
           for 500 random pairs and 21 values of beta.
[B-epi]    Convex combinations of 500 random point pairs of the epigraph
           {(w, t) : t >= f(w)} stay in the epigraph.
[B-global] Gradient descent on f from 20 random initializations always
           reaches the same minimizer (pairwise distance < 1e-6) with the
           least-squares optimum as its objective value: every local
           minimum is global.
[B-half]   Halfspace representation: C is contained in the intersection
           P_k of k supporting halfspaces, and the area of P_k minus C
           shrinks as k grows (k = 8, 32, 128): the intersection of all
           supporting halfspaces recovers C.

Outputs
-------
convex.png : matplotlib preview — 1-D slice of f with a chord above the
             graph (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}")


rng = np.random.default_rng(0)
betas = np.linspace(0.0, 1.0, 21)

# ---------------------------------------------------------------- [B-set]
D = np.diag([1.0, 4.0])


def in_set(w):
    return np.einsum("...i,ij,...j->...", w, D, w) <= 4.0 + 1e-12


pts = rng.uniform(-2.0, 2.0, (5000, 2))
pts = pts[in_set(pts)][:1000]
pairs = pts.reshape(-1, 2, 2)[:500]
ok_set = all(in_set(b * p[0] + (1 - b) * p[1])
             for p in pairs for b in betas)
check("[B-set]    segments between set points stay in the set", ok_set)

# ----------------------------------------------------------------- [B-fn]
m, d = 10, 2
X = rng.standard_normal((m, d))
y = X @ np.array([1.0, -2.0]) + 0.1 * rng.standard_normal(m)


def f(w):
    r = y - X @ w
    return float(r @ r) / m


W = rng.standard_normal((500, 2, 2)) * 3.0
ok_fn = all(f(b * w1 + (1 - b) * w2) <= b * f(w1) + (1 - b) * f(w2) + 1e-10
            for w1, w2 in W for b in betas)
check("[B-fn]     chord inequality for the linreg objective", ok_fn)

# ---------------------------------------------------------------- [B-epi]
ok_epi = True
for w1, w2 in W:
    t1 = f(w1) + abs(rng.standard_normal())      # points above the graph
    t2 = f(w2) + abs(rng.standard_normal())
    for b in betas:
        wb, tb = b * w1 + (1 - b) * w2, b * t1 + (1 - b) * t2
        if tb < f(wb) - 1e-10:
            ok_epi = False
check("[B-epi]    epigraph is a convex set", ok_epi)

# ------------------------------------------------------------- [B-global]
L = 2.0 * np.linalg.eigvalsh(X.T @ X / m).max()
eta = 1.0 / L
minimizers = []
for _ in range(20):
    w = rng.standard_normal(2) * 5.0
    for _ in range(2000):
        w = w - eta * (2.0 / m) * X.T @ (X @ w - y)
    minimizers.append(w)
minimizers = np.array(minimizers)
spread = np.max(np.linalg.norm(minimizers - minimizers[0], axis=1))
w_star, *_ = np.linalg.lstsq(X, y, rcond=None)
ok_glob = spread < 1e-6 and abs(f(minimizers[0]) - f(w_star)) < 1e-10
check("[B-global] GD from 20 random inits reaches the global minimum",
      ok_glob)

# ---------------------------------------------------------------- [B-half]
# Supporting halfspace of C = {w : w^T D w <= 4} at boundary point p:
# normal n = D p, halfspace {w : n^T w <= n^T p}.
def polygon_gap_area(k, n_mc=200000):
    """Monte-Carlo area of P_k \\ C for the intersection P_k of k
    supporting halfspaces at equally spaced boundary points."""
    th = np.linspace(0.0, 2.0 * np.pi, k, endpoint=False)
    bnd = np.c_[2.0 * np.cos(th), np.sin(th)]          # boundary of C
    N = bnd @ D                                        # outward normals
    c = np.einsum("ki,ki->k", N, bnd)
    S = rng_mc.uniform(-2.4, 2.4, (n_mc, 2))
    in_pk = np.all(S @ N.T <= c + 1e-12, axis=1)
    in_c = in_set(S)
    box = 4.8 * 4.8
    contain_ok = not np.any(in_c & ~in_pk)             # C subset of P_k
    return box * np.mean(in_pk & ~in_c), contain_ok


rng_mc = np.random.default_rng(1)
gaps, contains = zip(*(polygon_gap_area(k) for k in (8, 32, 128)))
check("[B-half]   C inside every halfspace intersection; gap area shrinks",
      all(contains) and gaps[0] > gaps[1] > gaps[2] and gaps[2] < 0.02)

# --------------------------------------------------------------- [B-sep]
# Linear separability of a binary trainset holds exactly when the convex
# hulls of the two classes do not intersect. Both directions are checked:
# hulls disjoint -> a separating (w, b) exists; hulls overlapping -> none
# does. Hull intersection is decided by the LP feasibility of
#   sum_i a_i x_i^+ = sum_j c_j x_j^-,  a, c >= 0,  sum a = sum c = 1,
# solved here by projected gradient on the squared distance between the
# two hulls (zero distance <=> the hulls meet).
def hull_distance(P, N, iters=5000):
    """Distance between conv(P) and conv(N) by Frank-Wolfe with exact line
    search on z = p - n, which stays in the (convex) Minkowski difference:
    each step moves z toward the vertex p_i - n_j that the current z sees as
    steepest, so ||z|| decreases monotonically to the true distance."""
    z = P[0] - N[0]
    for _ in range(iters):
        # vertex of conv(P) - conv(N) minimizing the linear approximation
        s_vert = P[np.argmin(P @ z)] - N[np.argmax(N @ z)]
        diff = z - s_vert
        denom = diff @ diff
        if denom < 1e-18:
            break
        gam = min(1.0, max(0.0, (z @ diff) / denom))   # exact line search
        z = z - gam * diff
    return float(np.linalg.norm(z))


def separable(P, N, iters=20000, lr=0.1):
    """Perceptron-style search for (w, b) with y (w^T x + b) > 0."""
    X = np.vstack([P, N])
    y = np.r_[np.ones(len(P)), -np.ones(len(N))]
    w, b = np.zeros(X.shape[1]), 0.0
    for _ in range(iters):
        m = y * (X @ w + b) <= 0
        if not m.any():
            return True
        w += lr * (y[m] @ X[m])
        b += lr * y[m].sum()
    return bool(np.all(y * (X @ w + b) > 0))


rng_sep = np.random.default_rng(3)
P_far = rng_sep.normal(size=(12, 2)) * 0.3 + np.array([2.5, 0.0])
N_far = rng_sep.normal(size=(12, 2)) * 0.3 + np.array([-2.5, 0.0])
P_mix = rng_sep.normal(size=(12, 2)) * 0.9
N_mix = rng_sep.normal(size=(12, 2)) * 0.9
check(f"[B-sep]    disjoint hulls (distance {hull_distance(P_far, N_far):.3f}) "
      f"-> the trainset is linearly separable",
      hull_distance(P_far, N_far) > 1e-3 and separable(P_far, N_far))
check(f"[B-sep]    overlapping hulls (distance {hull_distance(P_mix, N_mix):.3f}) "
      f"-> no linear classifier separates the trainset",
      hull_distance(P_mix, N_mix) < 1e-3 and not separable(P_mix, N_mix))

# ------------------------------------------------------------- [B-expfam]
# Exponential family p(x; w) ~ h(x) exp(w^T T(x) - A(w)) on x in {0,1,2}
# with T(x) = (x, x^2), h = 1. Two claims of the entry: the log-partition
# function A is convex in w, and its gradient is the expectation of the
# sufficient statistics.
Tstat = np.c_[np.arange(3), np.arange(3) ** 2].astype(float)


def logpart(w):
    return float(np.log(np.exp(Tstat @ w).sum()))


rng_ef = np.random.default_rng(5)
viol = 0
for _ in range(20000):                       # chord test for convexity of A
    u, v = rng_ef.normal(size=2) * 1.5, rng_ef.normal(size=2) * 1.5
    t = rng_ef.uniform()
    if logpart(t * u + (1 - t) * v) > t * logpart(u) + (1 - t) * logpart(v) + 1e-12:
        viol += 1
check(f"[B-expfam] the log-partition function A is convex "
      f"({viol} chord violations in 20000 random pairs)", viol == 0)

w_ef = np.array([0.3, -0.2])
h_ef = 1e-6
gradA = np.array([(logpart(w_ef + h_ef * e) - logpart(w_ef - h_ef * e)) / (2 * h_ef)
                  for e in np.eye(2)])
pw = np.exp(Tstat @ w_ef)
pw /= pw.sum()
check(f"[B-expfam] grad A equals the expected sufficient statistics "
      f"({gradA.round(4)} vs {(pw @ Tstat).round(4)})",
      np.allclose(gradA, pw @ Tstat, atol=1e-5))

# -------------------------------------------------------------- preview
w0, w1v = np.array([-2.0, -3.0]), np.array([3.0, 1.0])
ts = np.linspace(0.0, 1.0, 100)
fig, ax = plt.subplots(figsize=(4.6, 3.0))
ax.plot(ts, [f(t * w1v + (1 - t) * w0) for t in ts], "k-",
        label="$f$ along the segment")
ax.plot([0, 1], [f(w0), f(w1v)], "k--", label="chord")
ax.set_xlabel(r"$\beta$"), ax.set_ylabel("objective value")
ax.set_title("[B-fn] the chord lies above the function")
ax.legend(frameon=False, fontsize=8)
fig.tight_layout()
fig.savefig("pythondemos/convex.png", dpi=110)

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