"""
logreg.py — numerical companion to the glossary entry 'logistic regression'.

Purpose
-------
A one-feature binary classification trainset on which logistic regression
is fit by a hand-written gradient descent loop: the average logistic loss
decreases monotonically, the gradient vanishes at the learned parameters,
the GD update leaves them (approximately) unchanged — its fixed point —
and thresholding the learned hypothesis at zero classifies the trainset
better than always answering with the majority label.  Self-contained
(numpy/matplotlib only), fixed seed.

Blocks
------
[B-data]    m = 40 data points with a single feature x drawn uniformly
            from [-3, 3]; the binary label y in {-1, +1} is +1 with
            probability sigmoid(2 x - 1), so the labels are noisy around
            the point where 2 x - 1 = 0.
[B-gd]      Gradient descent on the average logistic loss
            f(w) = (1/m) sum_r log(1 + exp(-y^(r) w^T x^(r))) with the
            constant feature 1 absorbing the intercept: the loss never
            increases along the run, the gradient norm at the learned
            parameters is small, and one further GD update moves them by
            (approximately) nothing — the learned parameters are a fixed
            point of the update.
[B-clf]     The classifier sign(w^T x) obtained by thresholding the
            learned hypothesis classifies a larger fraction of the
            trainset correctly than the constant rule that always
            answers with the majority label.
[B-csv]     The two CSVs the entry's pgfplots figure reads.
[B-preview] The matplotlib preview of the figure.

Outputs
-------
logreg_points.csv : the trainset, columns x,y01,cls (y01 = (y+1)/2 for
                    drawing labels at heights 0 and 1; cls in {pos,neg}).
logreg_curve.csv  : the fitted probability curve sigmoid(w^T x) on a
                    grid, columns x,p.
logreg.png        : matplotlib preview of the figure (checking only).
"""

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 sigmoid(z):
    return 1.0 / (1.0 + np.exp(-z))


# ---- [B-data] one feature, binary labels from a noisy sigmoid rule --------
rng = np.random.default_rng(3)
m = 40
x = rng.uniform(-3.0, 3.0, m)
p_true = sigmoid(2.0 * x - 1.0)
y = np.where(rng.uniform(size=m) < p_true, 1.0, -1.0)
X = np.c_[x, np.ones(m)]                    # constant feature absorbs the intercept
check("[B-data]    both labels occur", (y > 0).any() and (y < 0).any())

# ---- [B-gd] gradient descent on the average logistic loss -----------------
def avg_logloss(w):
    return float(np.mean(np.log1p(np.exp(-y * (X @ w)))))


def grad(w):
    return -(X * (y * sigmoid(-y * (X @ w)))[:, None]).mean(axis=0)


eta = 0.5
w = np.zeros(2)
losses = [avg_logloss(w)]
for _ in range(2000):
    w = w - eta * grad(w)                   # the GD update w <- T(w)
    losses.append(avg_logloss(w))
check("[B-gd]      the average logistic loss never increases and ends "
      "below its start",
      all(b <= a + 1e-12 for a, b in zip(losses, losses[1:]))
      and losses[-1] < losses[0])
check("[B-gd]      the gradient vanishes at the learned parameters",
      float(np.linalg.norm(grad(w))) < 1e-4)
w_next = w - eta * grad(w)                  # one further update
check("[B-gd]      the learned parameters are a fixed point of the update",
      float(np.linalg.norm(w_next - w)) < 1e-4)

# ---- [B-clf] thresholding beats the constant majority rule ----------------
frac_lr = float(np.mean(np.sign(X @ w) == y))
majority = 1.0 if (y > 0).sum() >= (y < 0).sum() else -1.0
frac_const = float(np.mean(y == majority))
print(f"    correctly classified: logistic regression {frac_lr:.2f}, "
      f"majority label {frac_const:.2f}")
check("[B-clf]     sign(w^T x) beats the constant majority rule",
      frac_lr > frac_const)

# ---- [B-csv] the CSVs the entry's pgfplots figure reads -------------------
with open(OUT_DIR / "logreg_points.csv", "w") as f:
    f.write("x,y01,cls\n")
    for xi, yi in zip(x, y):
        f.write(f"{xi:.4f},{(yi + 1) / 2:.0f},{'pos' if yi > 0 else 'neg'}\n")

gx = np.linspace(-3.2, 3.2, 161)
gp = sigmoid(w[0] * gx + w[1])
with open(OUT_DIR / "logreg_curve.csv", "w") as f:
    f.write("x,p\n")
    for xi, pi in zip(gx, gp):
        f.write(f"{xi:.4f},{pi:.4f}\n")

# ---- [B-preview] matplotlib preview of the figure -------------------------
fig, ax = plt.subplots(figsize=(5.2, 3.4))
ax.plot(gx, gp, "k-", lw=1.6, label="sigmoid($\\hat{w}^{\\top} x$)")
ax.plot(x[y > 0], np.ones((y > 0).sum()), "ko", ms=5, label="label $+1$")
ax.plot(x[y < 0], np.zeros((y < 0).sum()), "ks", mfc="none", ms=5,
        label="label $-1$")
ax.set_xlabel("feature $x$")
ax.set_ylabel("label / probability of label $+1$")
ax.set_title("logistic regression: fitted probability curve")
ax.legend(frameon=False, loc="center right", fontsize=8)
fig.tight_layout()
fig.savefig(OUT_DIR / "logreg.png", dpi=110)

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