"""Three tiny models and the metrics the game uses. Everything here is deterministic
(no random numbers, fixed iteration counts) so that web/models.js can reproduce the
numbers exactly. Keep the two files in step: tests/parity.py checks them.

    logreg_fit / logreg_predict   full-batch gradient descent on standardized features
    knn1                          1-nearest neighbour on standardized features
    stump_fit / stump_predict     the best single threshold on a single feature
    auc, accuracy, balanced_accuracy, confusion
"""
from __future__ import annotations

import numpy as np

STD_EPS = 1e-12


def standardize_fit(X: np.ndarray):
    """Mean and 1/std per column; constant columns get 1/std = 0 so they contribute nothing."""
    X = np.asarray(X, dtype=float)
    mean = X.mean(axis=0)
    std = np.sqrt(((X - mean) ** 2).mean(axis=0))
    inv = np.where(std > STD_EPS, 1.0 / np.where(std > STD_EPS, std, 1.0), 0.0)
    return mean, inv


def standardize(X: np.ndarray, mean, inv) -> np.ndarray:
    return (np.asarray(X, dtype=float) - mean) * inv


def lipschitz(Z: np.ndarray, l2: float, iters: int = 60) -> float:
    """Largest eigenvalue of Z'Z/n by power iteration, then L = lambda/4 + l2
    (the Lipschitz constant of the logistic-loss gradient). Fixed start, fixed iterations."""
    n, d = Z.shape
    v = np.ones(d) / np.sqrt(d)
    lam = 0.0
    for _ in range(iters):
        u = Z @ v
        w = Z.T @ u
        norm = float(np.sqrt(np.dot(w, w)))
        if norm < 1e-300:
            return l2
        v = w / norm
        lam = norm / n
    return lam / 4.0 + l2


def sigmoid(s: np.ndarray) -> np.ndarray:
    with np.errstate(over="ignore"):
        return 1.0 / (1.0 + np.exp(-s))


def logreg_fit(X, y, iters: int = 500, l2: float = 1e-3) -> dict:
    """Logistic regression by full-batch gradient descent with step 1/L, zero start.
    Returns standardized-space weights (what the charts show) plus the mean/scale to apply."""
    X = np.asarray(X, dtype=float)
    y = np.asarray(y, dtype=float)
    mean, inv = standardize_fit(X)
    Z = standardize(X, mean, inv)
    n, d = Z.shape
    L = lipschitz(Z, l2)
    lr = min(10.0, 1.0 / L)
    w = np.zeros(d)
    b = 0.0
    for _ in range(iters):
        p = sigmoid(Z @ w + b)
        g = p - y
        w = w - lr * (Z.T @ g / n + l2 * w)
        b = b - lr * float(g.mean())
    return {"w": w, "b": b, "mean": mean, "inv": inv, "lr": lr, "iters": iters, "l2": l2}


def logreg_scores(model: dict, X) -> np.ndarray:
    Z = standardize(X, model["mean"], model["inv"])
    return sigmoid(Z @ model["w"] + model["b"])


def logreg_predict(model: dict, X) -> np.ndarray:
    return (logreg_scores(model, X) >= 0.5).astype(int)


def knn1(Xtr, ytr, Xte):
    """1-nearest neighbour in standardized space. Ties go to the lowest train index.
    Returns (predicted labels, distance to the nearest training row)."""
    Xtr = np.asarray(Xtr, dtype=float)
    Xte = np.asarray(Xte, dtype=float)
    ytr = np.asarray(ytr)
    mean, inv = standardize_fit(Xtr)
    A = standardize(Xtr, mean, inv)
    B = standardize(Xte, mean, inv)
    pred = np.empty(len(B), dtype=int)
    dist = np.empty(len(B))
    for i in range(len(B)):
        d2 = ((A - B[i]) ** 2).sum(axis=1)
        j = int(np.argmin(d2))          # first minimum = lowest index
        pred[i] = ytr[j]
        dist[i] = float(np.sqrt(d2[j]))
    return pred, dist


def stump_fit(X, y) -> dict:
    """Best single-feature threshold rule: predict `above` when x_j > t, else 1 - above.
    Ties keep the earlier feature, then the lower threshold, then above = 1."""
    X = np.asarray(X, dtype=float)
    y = np.asarray(y, dtype=int)
    n, d = X.shape
    best = {"feature": 0, "threshold": 0.0, "above": 1, "acc": -1.0}
    per_feature = []
    for j in range(d):
        order = np.argsort(X[:, j], kind="stable")
        xs, ys = X[order, j], y[order]
        pos_total = int(ys.sum())
        pos_left = 0
        fbest = -1.0
        for k in range(n - 1):          # split between position k and k+1
            pos_left += int(ys[k])
            if xs[k + 1] == xs[k]:
                continue
            left = k + 1
            # above=1: right side predicted 1, left predicted 0
            correct1 = (left - pos_left) + (pos_total - pos_left)
            correct0 = pos_left + ((n - left) - (pos_total - pos_left))
            t = 0.5 * (xs[k] + xs[k + 1])
            for above, c in ((1, correct1), (0, correct0)):
                acc = c / n
                if acc > fbest:
                    fbest = acc
                if acc > best["acc"]:
                    best = {"feature": j, "threshold": float(t), "above": above, "acc": acc}
        if fbest < 0:                   # constant feature: predict the majority
            fbest = max(pos_total, n - pos_total) / n
        per_feature.append(fbest)
    best["per_feature"] = per_feature
    return best


def stump_predict(model: dict, X) -> np.ndarray:
    X = np.asarray(X, dtype=float)
    above = X[:, model["feature"]] > model["threshold"]
    return np.where(above, model["above"], 1 - model["above"]).astype(int)


def auc(scores, labels) -> float:
    """Area under the ROC curve by the rank statistic, with average ranks for ties."""
    s = np.asarray(scores, dtype=float)
    yb = np.asarray(labels, dtype=int)
    n1 = int(yb.sum())
    n0 = len(yb) - n1
    if n1 == 0 or n0 == 0:
        return 0.5
    order = np.argsort(s, kind="stable")
    ranks = np.empty(len(s))
    i = 0
    while i < len(s):
        j = i
        while j + 1 < len(s) and s[order[j + 1]] == s[order[i]]:
            j += 1
        ranks[order[i:j + 1]] = 0.5 * (i + j) + 1.0   # average 1-based rank
        i = j + 1
    return float((ranks[yb == 1].sum() - n1 * (n1 + 1) / 2.0) / (n1 * n0))


def accuracy(pred, y) -> float:
    pred = np.asarray(pred)
    y = np.asarray(y)
    return float((pred == y).mean())


def confusion(pred, y) -> dict:
    pred = np.asarray(pred).astype(int)
    y = np.asarray(y).astype(int)
    return {"tp": int(((pred == 1) & (y == 1)).sum()), "fp": int(((pred == 1) & (y == 0)).sum()),
            "fn": int(((pred == 0) & (y == 1)).sum()), "tn": int(((pred == 0) & (y == 0)).sum())}


def balanced_accuracy(pred, y) -> float:
    c = confusion(pred, y)
    tpr = c["tp"] / (c["tp"] + c["fn"]) if c["tp"] + c["fn"] else 0.0
    tnr = c["tn"] / (c["tn"] + c["fp"]) if c["tn"] + c["fp"] else 0.0
    return 0.5 * (tpr + tnr)


def recall(pred, y) -> float:
    c = confusion(pred, y)
    return c["tp"] / (c["tp"] + c["fn"]) if c["tp"] + c["fn"] else 0.0