import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import make_moons
from sklearn.metrics import confusion_matrix, precision_recall_fscore_support, roc_curve, roc_auc_score
from sklearn.model_selection import train_test_split

plt.style.use("seaborn-v0_8-whitegrid")
rng = np.random.default_rng(42)

n_per_class = 100
x0 = rng.multivariate_normal([-1.2, -0.7], [[0.9, 0.25], [0.25, 0.7]], n_per_class)
x1 = rng.multivariate_normal([1.0, 0.9], [[0.8, -0.2], [-0.2, 0.8]], n_per_class)
X = np.vstack([x0, x1])
t = np.r_[np.zeros(n_per_class), np.ones(n_per_class)]

X_train_raw, X_test_raw, t_train, t_test = train_test_split(
    X, t, test_size=0.3, random_state=42, stratify=t
)
mean = X_train_raw.mean(axis=0)
std = X_train_raw.std(axis=0)
X_train = (X_train_raw - mean) / std
X_test = (X_test_raw - mean) / std
Phi_train = np.c_[np.ones(len(X_train)), X_train]
Phi_test = np.c_[np.ones(len(X_test)), X_test]

fig, ax = plt.subplots(figsize=(7, 5))
for label, color in [(0, "#2878B5"), (1, "#D95319")]:
    points = X_train[t_train == label]
    ax.scatter(points[:, 0], points[:, 1], s=35, alpha=0.8, label=f"class {label}", color=color)
ax.set(xlabel="$x_1$ (standardized)", ylabel="$x_2$ (standardized)", title="Binary classification data")
ax.legend()
plt.show()

print("train/test:", len(X_train), len(X_test), "design matrix:", Phi_train.shape)

def plot_binary_boundary(ax, X_data, labels, weights, title, probability=False):
    pad = 0.7
    x_min, x_max = X_data[:, 0].min() - pad, X_data[:, 0].max() + pad
    y_min, y_max = X_data[:, 1].min() - pad, X_data[:, 1].max() + pad
    xx, yy = np.meshgrid(np.linspace(x_min, x_max, 250), np.linspace(y_min, y_max, 250))
    grid_phi = np.c_[np.ones(xx.size), xx.ravel(), yy.ravel()]
    scores = (grid_phi @ weights).reshape(xx.shape)
    if probability:
        values = stable_sigmoid(scores)
        ax.contourf(xx, yy, values, levels=np.linspace(0, 1, 11), cmap="RdBu_r", alpha=0.28)
    else:
        ax.contourf(xx, yy, scores >= 0, levels=[-0.5, 0.5, 1.5], cmap="RdBu_r", alpha=0.18)
    ax.contour(xx, yy, scores, levels=[0], colors="black", linewidths=2)
    ax.scatter(X_data[:, 0], X_data[:, 1], c=labels, cmap="RdBu_r", edgecolor="white", s=38)
    ax.set(title=title, xlabel="$x_1$", ylabel="$x_2$")

w_demo = np.array([-0.1, 1.2, 0.8])
fig, ax = plt.subplots(figsize=(7, 5))
plot_binary_boundary(ax, X_train, t_train, w_demo, "A linear discriminant")
plt.show()

signed_distance = Phi_train @ w_demo / np.linalg.norm(w_demo[1:])
print("first five signed distances:", np.round(signed_distance[:5], 3))

w_ls = np.linalg.pinv(Phi_train) @ t_train
ls_scores = Phi_test @ w_ls
ls_accuracy = np.mean((ls_scores >= 0.5) == t_test)

# High-leverage, contradictory samples: class 1 points placed deep on the class 0 side.
outliers = np.array([[-4.2, -0.15], [-4.0, 0.0], [-3.8, 0.15]])
X_outlier = np.vstack([X_train, outliers])
t_outlier = np.r_[t_train, np.ones(len(outliers))]
Phi_outlier = np.c_[np.ones(len(X_outlier)), X_outlier]
w_ls_outlier = np.linalg.pinv(Phi_outlier) @ t_outlier

# plot_binary_boundary uses a zero threshold, so shift the intercept for the 0.5 LS threshold.
w_ls_boundary = w_ls.copy()
w_ls_boundary[0] -= 0.5
w_ls_outlier_boundary = w_ls_outlier.copy()
w_ls_outlier_boundary[0] -= 0.5

fig, axes = plt.subplots(1, 2, figsize=(12, 4.5), sharex=True, sharey=True)
plot_binary_boundary(axes[0], X_train, t_train, w_ls_boundary, "Clean data")
plot_binary_boundary(axes[1], X_outlier, t_outlier, w_ls_outlier_boundary, "With contradictory outliers")

# Overlay the clean-data boundary on the second panel for a direct visual comparison.
x_limits = (-4.6, 3.2)
y_limits = (-3.0, 3.2)
x_line = np.linspace(*x_limits, 300)
y_line = -(w_ls_boundary[0] + w_ls_boundary[1] * x_line) / w_ls_boundary[2]
axes[1].plot(x_line, y_line, "k--", linewidth=2, label="clean boundary")
axes[1].scatter(outliers[:, 0], outliers[:, 1], marker="X", s=130, color="#D95319", edgecolor="black", label="outliers")
for ax in axes:
    ax.set_xlim(x_limits)
    ax.set_ylim(y_limits)
axes[1].legend(loc="upper left")
plt.tight_layout()
plt.show()

clean_normal = w_ls_boundary[1:]
outlier_normal = w_ls_outlier_boundary[1:]
cosine = np.dot(clean_normal, outlier_normal) / (np.linalg.norm(clean_normal) * np.linalg.norm(outlier_normal))
boundary_rotation = np.degrees(np.arccos(np.clip(cosine, -1.0, 1.0)))
xx, yy = np.meshgrid(np.linspace(*x_limits, 250), np.linspace(*y_limits, 250))
comparison_grid = np.c_[np.ones(xx.size), xx.ravel(), yy.ravel()]
changed_fraction = np.mean(
    (comparison_grid @ w_ls_boundary >= 0) != (comparison_grid @ w_ls_outlier_boundary >= 0)
)

print(f"test accuracy without outliers: {ls_accuracy:.3f}")
print(f"decision-boundary rotation: {boundary_rotation:.1f} degrees")
print(f"changed predictions in the shown region: {changed_fraction:.1%}")
print(f"raw score range: [{ls_scores.min():.3f}, {ls_scores.max():.3f}] (not constrained to [0, 1])")

def fit_fisher(X_data, labels):
    class0, class1 = X_data[labels == 0], X_data[labels == 1]
    m0, m1 = class0.mean(axis=0), class1.mean(axis=0)
    S0 = (class0 - m0).T @ (class0 - m0)
    S1 = (class1 - m1).T @ (class1 - m1)
    direction = np.linalg.solve(S0 + S1 + 1e-8 * np.eye(X_data.shape[1]), m1 - m0)
    threshold = 0.5 * (direction @ m0 + direction @ m1)
    return direction, threshold

w_fisher, fisher_threshold = fit_fisher(X_train, t_train)
fisher_projection = X_train @ w_fisher
fisher_test_pred = (X_test @ w_fisher >= fisher_threshold).astype(int)
w_fisher_boundary = np.r_[-fisher_threshold, w_fisher]

fig, axes = plt.subplots(1, 2, figsize=(12, 4.5))
plot_binary_boundary(axes[0], X_train, t_train, w_fisher_boundary, "Fisher boundary")
axes[1].hist(fisher_projection[t_train == 0], bins=18, alpha=0.65, label="class 0")
axes[1].hist(fisher_projection[t_train == 1], bins=18, alpha=0.65, label="class 1")
axes[1].axvline(fisher_threshold, color="black", linestyle="--", label="threshold")
axes[1].set(title="Projection onto Fisher direction", xlabel="$w^T x$")
axes[1].legend()
plt.tight_layout()
plt.show()
print(f"Fisher test accuracy: {np.mean(fisher_test_pred == t_test):.3f}")

def fit_perceptron(Phi, labels, learning_rate=0.05, epochs=80):
    signed_labels = 2 * labels.astype(int) - 1
    weights = np.zeros(Phi.shape[1])
    mistakes_history, loss_history = [], []
    local_rng = np.random.default_rng(7)
    for _ in range(epochs):
        mistakes = 0
        for index in local_rng.permutation(len(Phi)):
            margin = signed_labels[index] * (Phi[index] @ weights)
            if margin <= 0:
                weights += learning_rate * signed_labels[index] * Phi[index]
                mistakes += 1
        margins = signed_labels * (Phi @ weights)
        mistakes_history.append(mistakes)
        loss_history.append(np.maximum(0, -margins).sum())
    return weights, np.array(mistakes_history), np.array(loss_history)

w_perceptron, perceptron_mistakes, perceptron_loss = fit_perceptron(Phi_train, t_train)
perceptron_pred = (Phi_test @ w_perceptron >= 0).astype(int)

fig, axes = plt.subplots(1, 2, figsize=(12, 4.5))
plot_binary_boundary(axes[0], X_train, t_train, w_perceptron, "Perceptron boundary")
axes[1].plot(perceptron_mistakes, label="mistakes per epoch")
axes[1].plot(perceptron_loss, label="perceptron criterion")
axes[1].set(title="Perceptron training", xlabel="epoch")
axes[1].legend()
plt.tight_layout()
plt.show()
print(f"Perceptron test accuracy: {np.mean(perceptron_pred == t_test):.3f}")

def fit_gaussian_lda(X_data, labels):
    classes = np.array([0, 1])
    means = np.vstack([X_data[labels == k].mean(axis=0) for k in classes])
    priors = np.array([np.mean(labels == k) for k in classes])
    covariance = sum(
        (X_data[labels == k] - means[k]).T @ (X_data[labels == k] - means[k])
        for k in classes
    ) / len(X_data)
    covariance += 1e-8 * np.eye(X_data.shape[1])
    precision_mean = np.linalg.solve(covariance, (means[1] - means[0]))
    solved_means = np.linalg.solve(covariance, means.T).T
    intercept = (
        -0.5 * means[1] @ solved_means[1]
        + 0.5 * means[0] @ solved_means[0]
        + np.log(priors[1] / priors[0])
    )
    return np.r_[intercept, precision_mean], means, covariance, priors

w_lda, lda_means, lda_covariance, lda_priors = fit_gaussian_lda(X_train, t_train)
lda_pred = (Phi_test @ w_lda >= 0).astype(int)

fig, ax = plt.subplots(figsize=(7, 5))
plot_binary_boundary(ax, X_train, t_train, w_lda, "Gaussian generative classifier (LDA)")
ax.scatter(lda_means[:, 0], lda_means[:, 1], marker="X", s=180, c=[0, 1], cmap="RdBu_r", edgecolor="black")
plt.show()
print("estimated priors:", lda_priors)
print(f"LDA test accuracy: {np.mean(lda_pred == t_test):.3f}")

def stable_sigmoid(values):
    values = np.asarray(values, dtype=float)
    result = np.empty_like(values)
    positive = values >= 0
    result[positive] = 1.0 / (1.0 + np.exp(-values[positive]))
    exp_values = np.exp(values[~positive])
    result[~positive] = exp_values / (1.0 + exp_values)
    return result


def binary_cross_entropy_from_logits(logits, labels):
    # log(1 + exp(a)) - t*a is BCE without taking log(sigmoid(a)).
    return np.mean(np.logaddexp(0.0, logits) - labels * logits)


logit_grid = np.linspace(-8, 8, 400)
fig, ax = plt.subplots(figsize=(7, 4))
ax.plot(logit_grid, stable_sigmoid(logit_grid), linewidth=2)
ax.axhline(0.5, color="black", linestyle="--")
ax.axvline(0, color="black", linestyle="--")
ax.set(xlabel="logit a", ylabel="$\\sigma(a)$", title="Sigmoid maps scores to probabilities")
plt.show()

def fit_logistic_gd(Phi, labels, learning_rate=0.2, iterations=6000, l2=0.0, tolerance=1e-9):
    weights = np.zeros(Phi.shape[1])
    losses = []
    regularizer = np.r_[0.0, np.ones(Phi.shape[1] - 1)]  # Do not penalize the intercept.
    for _ in range(iterations):
        logits = Phi @ weights
        probabilities = stable_sigmoid(logits)
        gradient = Phi.T @ (probabilities - labels) / len(Phi)
        gradient += l2 * regularizer * weights
        weights -= learning_rate * gradient
        losses.append(
            binary_cross_entropy_from_logits(Phi @ weights, labels)
            + 0.5 * l2 * np.sum(weights[1:] ** 2)
        )
        if np.linalg.norm(gradient) < tolerance:
            break
    return weights, np.array(losses)

w_logistic, logistic_losses = fit_logistic_gd(Phi_train, t_train)
logistic_test_probability = stable_sigmoid(Phi_test @ w_logistic)
logistic_pred = (logistic_test_probability >= 0.5).astype(int)

fig, axes = plt.subplots(1, 2, figsize=(12, 4.5))
plot_binary_boundary(axes[0], X_train, t_train, w_logistic, "Logistic regression probability", probability=True)
axes[1].plot(logistic_losses)
axes[1].set(title="Binary cross-entropy", xlabel="iteration", ylabel="loss")
plt.tight_layout()
plt.show()
print(f"Logistic regression test accuracy: {np.mean(logistic_pred == t_test):.3f}")
print(f"Final BCE: {logistic_losses[-1]:.4f}")

def fit_logistic_irls(Phi, labels, max_iterations=30, damping=1e-6, tolerance=1e-9):
    weights = np.zeros(Phi.shape[1])
    losses = []
    for _ in range(max_iterations):
        probabilities = stable_sigmoid(Phi @ weights)
        gradient = Phi.T @ (probabilities - labels)
        r = np.clip(probabilities * (1 - probabilities), 1e-9, None)
        hessian = Phi.T @ (r[:, None] * Phi) + damping * np.eye(Phi.shape[1])
        step = np.linalg.solve(hessian, gradient)
        weights -= step
        losses.append(binary_cross_entropy_from_logits(Phi @ weights, labels))
        if np.linalg.norm(step) < tolerance:
            break
    return weights, np.array(losses)

w_irls, irls_losses = fit_logistic_irls(Phi_train, t_train)
print(f"Gradient descent: {len(logistic_losses)} steps, final BCE={logistic_losses[-1]:.6f}")
print(f"IRLS: {len(irls_losses)} steps, final BCE={irls_losses[-1]:.6f}")

fig, ax = plt.subplots(figsize=(7, 4))
ax.plot(np.arange(1, len(irls_losses) + 1), irls_losses, marker="o")
ax.set(xlabel="IRLS iteration", ylabel="BCE", title="Newton / IRLS convergence")
plt.show()

def stable_softmax(logits):
    shifted = logits - logits.max(axis=1, keepdims=True)
    exponentials = np.exp(shifted)
    return exponentials / exponentials.sum(axis=1, keepdims=True)


def fit_softmax(Phi, labels, class_count, learning_rate=0.2, iterations=1500, l2=1e-4):
    targets = np.eye(class_count)[labels]
    weights = np.zeros((Phi.shape[1], class_count))
    losses = []
    regularizer = np.r_[np.zeros((1, class_count)), np.ones((Phi.shape[1] - 1, class_count))]
    for _ in range(iterations):
        probabilities = stable_softmax(Phi @ weights)
        gradient = Phi.T @ (probabilities - targets) / len(Phi) + l2 * regularizer * weights
        weights -= learning_rate * gradient
        weights -= weights.mean(axis=1, keepdims=True)  # Choose the zero-sum representative.
        updated_probabilities = stable_softmax(Phi @ weights)
        cross_entropy = -np.mean(
            np.sum(targets * np.log(np.clip(updated_probabilities, 1e-12, 1.0)), axis=1)
        )
        losses.append(cross_entropy + 0.5 * l2 * np.sum(weights[1:] ** 2))
    return weights, np.array(losses)

centers = np.array([[-1.8, -1.0], [1.8, -0.8], [0.0, 1.8]])
X_multi = np.vstack([rng.normal(center, [0.65, 0.65], size=(90, 2)) for center in centers])
t_multi = np.repeat(np.arange(3), 90)
X_multi_train, X_multi_test, t_multi_train, t_multi_test = train_test_split(
    X_multi, t_multi, test_size=0.3, random_state=42, stratify=t_multi
)
multi_mean, multi_std = X_multi_train.mean(0), X_multi_train.std(0)
X_multi_train = (X_multi_train - multi_mean) / multi_std
X_multi_test = (X_multi_test - multi_mean) / multi_std
Phi_multi_train = np.c_[np.ones(len(X_multi_train)), X_multi_train]
Phi_multi_test = np.c_[np.ones(len(X_multi_test)), X_multi_test]
W_softmax, softmax_losses = fit_softmax(Phi_multi_train, t_multi_train, 3)
softmax_pred = np.argmax(Phi_multi_test @ W_softmax, axis=1)
print(f"Softmax test accuracy: {np.mean(softmax_pred == t_multi_test):.3f}")

def plot_multiclass_regions(ax, X_data, labels, weights, title):
    x_min, x_max = X_data[:, 0].min() - 0.6, X_data[:, 0].max() + 0.6
    y_min, y_max = X_data[:, 1].min() - 0.6, X_data[:, 1].max() + 0.6
    xx, yy = np.meshgrid(np.linspace(x_min, x_max, 250), np.linspace(y_min, y_max, 250))
    grid = np.c_[np.ones(xx.size), xx.ravel(), yy.ravel()]
    regions = np.argmax(grid @ weights, axis=1).reshape(xx.shape)
    ax.contourf(xx, yy, regions, levels=np.arange(4) - 0.5, cmap="viridis", alpha=0.2)
    ax.scatter(X_data[:, 0], X_data[:, 1], c=labels, cmap="viridis", edgecolor="white", s=38)
    ax.set(title=title, xlabel="$x_1$", ylabel="$x_2$")

fig, axes = plt.subplots(1, 2, figsize=(12, 4.5))
plot_multiclass_regions(axes[0], X_multi_train, t_multi_train, W_softmax, "Softmax decision regions")
axes[1].plot(softmax_losses)
axes[1].set(title="Multiclass cross-entropy", xlabel="iteration", ylabel="loss")
plt.tight_layout()
plt.show()

for threshold in [0.3, 0.5, 0.7]:
    predictions = (logistic_test_probability >= threshold).astype(int)
    precision, recall, f1, _ = precision_recall_fscore_support(
        t_test, predictions, average="binary", zero_division=0
    )
    print(f"threshold={threshold:.1f}  precision={precision:.3f}  recall={recall:.3f}  F1={f1:.3f}")

matrix = confusion_matrix(t_test, logistic_pred)
fpr, tpr, thresholds = roc_curve(t_test, logistic_test_probability)
auc = roc_auc_score(t_test, logistic_test_probability)

fig, axes = plt.subplots(1, 2, figsize=(10, 4.2))
image = axes[0].imshow(matrix, cmap="Blues")
for row in range(2):
    for column in range(2):
        axes[0].text(column, row, matrix[row, column], ha="center", va="center", fontsize=14)
axes[0].set(title="Confusion matrix (threshold=0.5)", xlabel="predicted", ylabel="true", xticks=[0, 1], yticks=[0, 1])
axes[1].plot(fpr, tpr, linewidth=2, label=f"AUC={auc:.3f}")
axes[1].plot([0, 1], [0, 1], "--", color="gray")
axes[1].set(xlabel="false positive rate", ylabel="true positive rate", title="ROC curve")
axes[1].legend()
plt.tight_layout()
plt.show()

def fit_bayesian_logistic_laplace(Phi, labels, alpha=1.0, max_iterations=50):
    weights = np.zeros(Phi.shape[1])
    prior_precision = alpha * np.eye(Phi.shape[1])
    prior_precision[0, 0] = 0.0
    for _ in range(max_iterations):
        probabilities = stable_sigmoid(Phi @ weights)
        r = np.clip(probabilities * (1 - probabilities), 1e-9, None)
        gradient = Phi.T @ (probabilities - labels) + prior_precision @ weights
        hessian = Phi.T @ (r[:, None] * Phi) + prior_precision + 1e-8 * np.eye(Phi.shape[1])
        step = np.linalg.solve(hessian, gradient)
        weights -= step
        if np.linalg.norm(step) < 1e-9:
            break
    probabilities = stable_sigmoid(Phi @ weights)
    r = np.clip(probabilities * (1 - probabilities), 1e-9, None)
    hessian = Phi.T @ (r[:, None] * Phi) + prior_precision + 1e-8 * np.eye(Phi.shape[1])
    covariance = np.linalg.solve(hessian, np.eye(Phi.shape[1]))
    return weights, covariance


def bayesian_predict(Phi, weights, covariance):
    mean_activation = Phi @ weights
    variance_activation = np.einsum("ij,jk,ik->i", Phi, covariance, Phi)
    kappa = 1.0 / np.sqrt(1.0 + np.pi * variance_activation / 8.0)
    return stable_sigmoid(kappa * mean_activation), variance_activation

w_map, posterior_covariance = fit_bayesian_logistic_laplace(Phi_train, t_train, alpha=1.0)
bayes_test_probability, bayes_test_variance = bayesian_predict(Phi_test, w_map, posterior_covariance)
print(f"Bayesian approximate test accuracy: {np.mean((bayes_test_probability >= 0.5) == t_test):.3f}")
print("posterior standard deviations:", np.round(np.sqrt(np.diag(posterior_covariance)), 3))

grid_axis = np.linspace(-4.5, 4.5, 260)
xx, yy = np.meshgrid(grid_axis, grid_axis)
Phi_grid = np.c_[np.ones(xx.size), xx.ravel(), yy.ravel()]
map_probability = stable_sigmoid(Phi_grid @ w_map).reshape(xx.shape)
bayes_probability, activation_variance = bayesian_predict(Phi_grid, w_map, posterior_covariance)
bayes_probability = bayes_probability.reshape(xx.shape)
activation_variance = activation_variance.reshape(xx.shape)

fig, axes = plt.subplots(1, 3, figsize=(15, 4.3))
for ax, values, title, cmap in [
    (axes[0], map_probability, "MAP plug-in probability", "RdBu_r"),
    (axes[1], bayes_probability, "Laplace predictive probability", "RdBu_r"),
    (axes[2], activation_variance, "Activation uncertainty", "magma"),
]:
    contour = ax.contourf(xx, yy, values, levels=20, cmap=cmap)
    ax.scatter(X_train[:, 0], X_train[:, 1], c=t_train, cmap="RdBu_r", edgecolor="white", s=20)
    ax.set(title=title, xlabel="$x_1$", ylabel="$x_2$")
    fig.colorbar(contour, ax=ax, shrink=0.82)
plt.tight_layout()
plt.show()

X_moons, t_moons = make_moons(n_samples=300, noise=0.18, random_state=42)
moon_mean, moon_std = X_moons.mean(0), X_moons.std(0)
X_moons = (X_moons - moon_mean) / moon_std


def linear_features(X_data):
    return np.c_[np.ones(len(X_data)), X_data]


def quadratic_features(X_data):
    x1, x2 = X_data[:, 0], X_data[:, 1]
    return np.c_[np.ones(len(X_data)), x1, x2, x1**2, x1 * x2, x2**2]

Phi_moon_linear = linear_features(X_moons)
Phi_moon_quadratic = quadratic_features(X_moons)
w_moon_linear, _ = fit_logistic_gd(Phi_moon_linear, t_moons, learning_rate=0.12, iterations=1600, l2=1e-3)
w_moon_quadratic, _ = fit_logistic_gd(Phi_moon_quadratic, t_moons, learning_rate=0.12, iterations=1600, l2=1e-3)

axis = np.linspace(-2.5, 2.5, 260)
xx, yy = np.meshgrid(axis, axis)
grid_points = np.c_[xx.ravel(), yy.ravel()]
fig, axes = plt.subplots(1, 2, figsize=(12, 4.5))
for ax, feature_fn, weights, title in [
    (axes[0], linear_features, w_moon_linear, "Original linear features"),
    (axes[1], quadratic_features, w_moon_quadratic, "Quadratic basis features"),
]:
    probability = stable_sigmoid(feature_fn(grid_points) @ weights).reshape(xx.shape)
    ax.contourf(xx, yy, probability, levels=np.linspace(0, 1, 11), cmap="RdBu_r", alpha=0.3)
    ax.contour(xx, yy, probability, levels=[0.5], colors="black", linewidths=2)
    ax.scatter(X_moons[:, 0], X_moons[:, 1], c=t_moons, cmap="RdBu_r", edgecolor="white", s=28)
    accuracy = np.mean((stable_sigmoid(feature_fn(X_moons) @ weights) >= 0.5) == t_moons)
    ax.set(title=f"{title}\naccuracy={accuracy:.3f}", xlabel="$x_1$", ylabel="$x_2$")
plt.tight_layout()
plt.show()