#!/usr/bin/env python3
"""M0 milestone: prove the core training mechanic in pure numpy.

Trains a tiny char-level MLP language model twice, once with fp32 weights
and once with BitNet-style ternary quantization-aware training:

  forward:  W_t = clip(round(W / s), -1, 1) * s,  s = mean(|W|) (absmean)
  backward: straight-through estimator, gradient flows to latent fp W as-is

If the ternary run tracks the fp32 run within a modest gap, the core bet of
Mind-1B (ternary weights trained with STE) works at toy scale and the same
loop ports to torch/Triton at 100M+ scale. Trains on its own source code.

Run: python3 toy_ternary_lm.py [--steps 3000]
"""

import argparse
import os

# matrices here are tiny; multithreaded BLAS thrashes and runs 100x slower
for _v in ("OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", "MKL_NUM_THREADS",
           "VECLIB_MAXIMUM_THREADS"):
    os.environ.setdefault(_v, "1")

import numpy as np


def ternary(W):
    s = np.mean(np.abs(W)) + 1e-8
    return np.clip(np.round(W / s), -1, 1) * s


class ToyLM:
    """context-window one-hot MLP: (ctx*V) -> hidden -> V"""

    def __init__(self, vocab, ctx=16, hidden=256, seed=0, quantize=False):
        rng = np.random.default_rng(seed)
        self.V, self.ctx, self.q = vocab, ctx, quantize
        self.W1 = rng.normal(0, 0.02, (ctx * vocab, hidden)).astype(np.float32)
        self.b1 = np.zeros(hidden, dtype=np.float32)
        self.W2 = rng.normal(0, 0.02, (hidden, vocab)).astype(np.float32)
        self.b2 = np.zeros(vocab, dtype=np.float32)
        self.adam = {k: [np.zeros_like(v), np.zeros_like(v)]
                     for k, v in self.params().items()}
        self.t = 0

    def params(self):
        return {"W1": self.W1, "b1": self.b1, "W2": self.W2, "b2": self.b2}

    def forward(self, X):
        W1 = ternary(self.W1) if self.q else self.W1
        W2 = ternary(self.W2) if self.q else self.W2
        h_pre = X @ W1 + self.b1
        h = np.maximum(h_pre, 0)
        logits = h @ W2 + self.b2
        return logits, (X, h_pre, h, W1, W2)

    def loss_and_grads(self, X, y):
        logits, (Xc, h_pre, h, W1, W2) = self.forward(X)
        logits -= logits.max(axis=1, keepdims=True)
        p = np.exp(logits)
        p /= p.sum(axis=1, keepdims=True)
        n = len(y)
        loss = -np.log(p[np.arange(n), y] + 1e-9).mean()

        dlogits = p
        dlogits[np.arange(n), y] -= 1
        dlogits /= n
        # STE: grads computed wrt quantized W flow straight to latent fp W
        gW2 = h.T @ dlogits
        gb2 = dlogits.sum(0)
        dh = dlogits @ W2.T
        dh[h_pre <= 0] = 0
        gW1 = Xc.T @ dh
        gb1 = dh.sum(0)
        return loss, {"W1": gW1, "b1": gb1, "W2": gW2, "b2": gb2}

    def step(self, grads, lr=1e-3, b1=0.9, b2=0.999, eps=1e-8):
        self.t += 1
        for k, v in self.params().items():
            m, s = self.adam[k]
            g = grads[k]
            m[:] = b1 * m + (1 - b1) * g
            s[:] = b2 * s + (1 - b2) * g * g
            mh = m / (1 - b1 ** self.t)
            sh = s / (1 - b2 ** self.t)
            v -= lr * mh / (np.sqrt(sh) + eps)

    def sample(self, itos, stoi, seed_text, n=200, temp=0.8, rng=None):
        rng = rng or np.random.default_rng(1)
        ctx = [stoi.get(c, 0) for c in seed_text[-self.ctx:]]
        ctx = [0] * (self.ctx - len(ctx)) + ctx
        out = []
        for _ in range(n):
            X = np.zeros((1, self.ctx * self.V), dtype=np.float32)
            for i, c in enumerate(ctx):
                X[0, i * self.V + c] = 1
            logits, _ = self.forward(X)
            p = np.exp(logits[0] / temp - np.max(logits[0] / temp))
            p /= p.sum()
            c = rng.choice(self.V, p=p)
            out.append(itos[c])
            ctx = ctx[1:] + [int(c)]
        return "".join(out)


def batches(data, V, ctx, batch, rng):
    while True:
        idx = rng.integers(0, len(data) - ctx - 1, size=batch)
        X = np.zeros((batch, ctx * V), dtype=np.float32)
        y = np.empty(batch, dtype=np.int64)
        for b, i in enumerate(idx):
            for j in range(ctx):
                X[b, j * V + data[i + j]] = 1
            y[b] = data[i + ctx]
        yield X, y


def train(quantize, text, steps, ctx=16):
    chars = sorted(set(text))
    stoi = {c: i for i, c in enumerate(chars)}
    itos = {i: c for i, c in enumerate(chars)}
    data = np.array([stoi[c] for c in text], dtype=np.int64)
    V = len(chars)
    model = ToyLM(V, ctx=ctx, quantize=quantize)
    rng = np.random.default_rng(42)
    gen = batches(data, V, ctx, 64, rng)
    label = "ternary-QAT" if quantize else "fp32       "
    for step in range(1, steps + 1):
        X, y = next(gen)
        loss, grads = model.loss_and_grads(X, y)
        model.step(grads)
        if step % (steps // 6) == 0:
            print("  [%s] step %5d  loss %.3f" % (label, step, loss))
    return model, itos, stoi, loss


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--steps", type=int, default=3000)
    a = p.parse_args()

    text = open(__file__).read()
    print("corpus: this file, %d chars, %d unique" % (len(text), len(set(text))))

    print("\nbaseline fp32:")
    _, _, _, fp_loss = train(False, text, a.steps)
    print("\nternary QAT (absmean + STE):")
    m, itos, stoi, q_loss = train(True, text, a.steps)

    w = np.concatenate([m.W1.ravel(), m.W2.ravel()])
    s = np.mean(np.abs(w))
    tw = np.clip(np.round(w / s), -1, 1)
    zeros = float(np.mean(tw == 0))
    print("\nresults")
    print("  fp32 final loss:    %.3f" % fp_loss)
    print("  ternary final loss: %.3f (gap %.3f)" % (q_loss, q_loss - fp_loss))
    print("  ternary sparsity:   %.0f%% zeros (free activation sparsity)" % (zeros * 100))
    print("\nsample from the ternary model:")
    print("  " + m.sample(itos, stoi, "def ", n=160).replace("\n", "\n  "))


if __name__ == "__main__":
    main()