"""Generate a pre-trained 3-layer MLP classifier and save as model.npz.

Run once: python3 scripts/generate_model.py
Creates model.npz in the project root directory.
"""

import numpy as np

SEED = 42
rng = np.random.default_rng(SEED)

N_FEATURES = 20
HIDDEN1 = 64
HIDDEN2 = 32
N_CLASSES = 5

N_TRAIN = 2000
N_TEST = 500

class_means = rng.normal(0, 1, (N_CLASSES, N_FEATURES))
X_train = np.vstack([
    rng.normal(class_means[c], 0.8, (N_TRAIN // N_CLASSES, N_FEATURES))
    for c in range(N_CLASSES)
])
y_train = np.repeat(np.arange(N_CLASSES), N_TRAIN // N_CLASSES)

X_test = np.vstack([
    rng.normal(class_means[c], 0.8, (N_TEST // N_CLASSES, N_FEATURES))
    for c in range(N_CLASSES)
])
y_test = np.repeat(np.arange(N_CLASSES), N_TEST // N_CLASSES)

def xavier(fan_in, fan_out, rng):
    scale = np.sqrt(2.0 / (fan_in + fan_out))
    return rng.normal(0, scale, (fan_in, fan_out)).astype(np.float32)

W1 = xavier(N_FEATURES, HIDDEN1, rng)
b1 = np.zeros(HIDDEN1, dtype=np.float32)
W2 = xavier(HIDDEN1, HIDDEN2, rng)
b2 = np.zeros(HIDDEN2, dtype=np.float32)
W3 = xavier(HIDDEN2, N_CLASSES, rng)
b3 = np.zeros(N_CLASSES, dtype=np.float32)

def relu(x): return np.maximum(0, x)
def softmax(x):
    e = np.exp(x - x.max(axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)

def forward(X, W1, b1, W2, b2, W3, b3):
    h1 = relu(X @ W1 + b1)
    h2 = relu(h1 @ W2 + b2)
    return softmax(h2 @ W3 + b3)

LR = 0.01
for epoch in range(200):
    idx = rng.permutation(len(X_train))
    for i in range(0, len(X_train), 64):
        batch_idx = idx[i:i+64]
        X_b = X_train[batch_idx].astype(np.float32)
        y_b = y_train[batch_idx]

        h1 = relu(X_b @ W1 + b1)
        h2 = relu(h1 @ W2 + b2)
        logits = h2 @ W3 + b3
        probs = softmax(logits)

        dlogits = probs.copy()
        dlogits[np.arange(len(y_b)), y_b] -= 1
        dlogits /= len(y_b)

        dW3 = h2.T @ dlogits
        db3 = dlogits.sum(0)
        dh2 = dlogits @ W3.T
        dh2[h2 == 0] = 0

        dW2 = h1.T @ dh2
        db2 = dh2.sum(0)
        dh1 = dh2 @ W2.T
        dh1[h1 == 0] = 0

        dW1 = X_b.T @ dh1
        db1 = dh1.sum(0)

        W1 -= LR * dW1; b1 -= LR * db1
        W2 -= LR * dW2; b2 -= LR * db2
        W3 -= LR * dW3; b3 -= LR * db3

probs = forward(X_test.astype(np.float32), W1, b1, W2, b2, W3, b3)
preds = probs.argmax(axis=1)
acc = (preds == y_test).mean()
print(f"Model accuracy: {acc:.4f}")

np.savez("model.npz",
         W1=W1, b1=b1, W2=W2, b2=b2, W3=W3, b3=b3,
         X_test=X_test.astype(np.float32), y_test=y_test,
         class_means=class_means.astype(np.float32))
print("Saved model.npz")
print(f"W1: {W1.shape}, W2: {W2.shape}, W3: {W3.shape}")
