import re
import numpy as np
from collections import Counter

def sigmoid(z):
    return 1 / (1 + np.exp(-z))

def softmax(z):
    e = np.exp(z)
    return e / e.sum(axis=-1, keepdims=True)

def cross_entropy(probs, Y):
    correct = [probs[t, y] for t, y in enumerate(Y)]
    return -np.log(correct).mean()

class Vocabulary:
    UNK = "<unk>"
    PAD = "<pad>"

    def __init__(self, text, size):
        special_tokens = [self.UNK, self.PAD]
        words = self._tokenize(text)

        counts = Counter(words)
        common_words = [
            word
            for word, _ in counts.most_common(size - len(special_tokens))
        ]

        self.itos = special_tokens + common_words
        self.stoi = {
            word: i
            for i, word in enumerate(self.itos)
        }

    def encode(self, text):
        unk = self.stoi[self.UNK]
        return [
            self.stoi.get(t, unk)
            for t in self._tokenize(text)
        ]

    def decode(self, ids):
        return " ".join(self.itos[i] for i in ids)

    def _tokenize(self, text):
        return re.findall(
            r"[a-z]+(?:'[a-z]+)?|[.,!?;:]", text.lower()
        )

def training_examples(ids, block):
    n = (len(ids) - 1) // block
    X = np.stack([
        ids[i * block:i * block + block]
        for i in range(n)
    ])
    Y = np.stack([
        ids[i * block + 1:i * block + block + 1]
        for i in range(n)
    ])
    return X, Y

class Linear:
    def __init__(self, nin, nout, rng, bias=True):
        s = 1 / np.sqrt(nin)
        self.W = rng.normal(0, s, (nin, nout))
        self.dW = np.zeros_like(self.W)
        self.b = np.zeros(nout) if bias else None
        self.db = np.zeros(nout) if bias else None

    def forward(self, x):
        self.x = x
        y = x @ self.W
        return y + self.b if self.b is not None else y

    def backward(self, dy):
        self.dW += self.x.T @ dy
        if self.b is not None:
            self.db += dy.sum(0)
        return dy @ self.W.T

    def params(self):
        return [(self.W, self.dW)] + ([(self.b, self.db)] if self.b is not None else [])

class Embedding:
    def __init__(self, vocab_size, d, block, rng):
        self.tok = rng.normal(0, 0.05, (vocab_size, d))
        self.dtok = np.zeros_like(self.tok)
        self.pos = rng.normal(0, 0.05, (block, d))
        self.dpos = np.zeros_like(self.pos)

    def forward(self, X):
        self.X = X
        T = len(X)
        return self.tok[X] + self.pos[:T]

    def backward(self, dout):
        np.add.at(self.dtok, self.X, dout)
        self.dpos[:len(self.X)] += dout

    def params(self):
        return [(self.tok, self.dtok), (self.pos, self.dpos)]

class Attention:
    def __init__(self, d, block, rng):
        self.Wq = Linear(d, d, rng, bias=False)
        self.Wk = Linear(d, d, rng, bias=False)
        self.Wv = Linear(d, d, rng, bias=False)
        self.Wo = Linear(d, d, rng, bias=False)
        self.mask = np.triu(np.full((block, block), -1e9), 1)
        self.d = d

    def forward(self, x):
        T = x.shape[0]
        Q = self.Wq.forward(x)
        K = self.Wk.forward(x)
        V = self.Wv.forward(x)
        self.Q, self.K, self.V = Q, K, V
        scores = Q @ K.T / np.sqrt(self.d)
        scores = scores + self.mask[:T, :T]
        self.attn = softmax(scores)
        return self.Wo.forward(self.attn @ V)

    def backward(self, dout):
        dctx = self.Wo.backward(dout)
        dattn = dctx @ self.V.T
        dV = self.attn.T @ dctx
        dscores = self.attn * (dattn - (dattn * self.attn).sum(-1, keepdims=True))
        dscores = dscores / np.sqrt(self.d)
        dQ = dscores @ self.K
        dK = dscores.T @ self.Q
        return self.Wq.backward(dQ) + self.Wk.backward(dK) + self.Wv.backward(dV)

    def params(self):
        return self.Wq.params() + self.Wk.params() + self.Wv.params() + self.Wo.params()

class MLP:
    def __init__(self, d, hidden, rng):
        self.layer1 = Linear(d, hidden, rng)
        self.layer2 = Linear(hidden, d, rng)

    def forward(self, x):
        self.h = sigmoid(self.layer1.forward(x))
        return self.layer2.forward(self.h)

    def backward(self, dout):
        dh = self.layer2.backward(dout)
        return self.layer1.backward(dh * self.h * (1 - self.h))

    def params(self):
        return self.layer1.params() + self.layer2.params()

class Transformer:
    def __init__(self, vocab_size, d, hidden, block, seed=0):
        rng = np.random.default_rng(seed)
        self.block = block
        self.embed = Embedding(vocab_size, d, block, rng)
        self.attn = Attention(d, block, rng)
        self.mlp = MLP(d, hidden, rng)
        self.output = Linear(d, vocab_size, rng)
        self.mods = [self.embed, self.attn, self.mlp, self.output]

    def forward(self, X):
        x = self.embed.forward(X)
        r = x + self.attn.forward(x)
        r2 = r + self.mlp.forward(r)
        return self.output.forward(r2)

    def loss(self, X, Y):
        logits = self.forward(X)
        self.probs = softmax(logits)
        self.Y = Y
        return cross_entropy(self.probs, Y)

    def backward(self):
        dlogits = self.probs.copy()
        dlogits[np.arange(len(self.Y)), self.Y] -= 1
        dlogits /= len(self.Y)
        dr2 = self.output.backward(dlogits)
        dr = dr2 + self.mlp.backward(dr2)
        dx = dr + self.attn.backward(dr)
        self.embed.backward(dx)

    def params(self):
        return [pg for m in self.mods for pg in m.params()]

    def predict(self, context_ids):
        ids = list(context_ids)[-self.block:] or [0]
        logits = self.forward(np.array(ids))
        return softmax(logits[-1])

    def train(self, X, Y, steps, lr):
        n = X.shape[0]
        rng = np.random.default_rng(1)
        for _ in range(steps):
            i = rng.integers(0, n)
            for param, grad in self.params():
                grad[:] = 0
            self.loss(X[i], Y[i])
            self.backward()
            for param, grad in self.params():
                np.clip(grad, -5, 5, out=grad)
                param -= lr * grad
        return self

def generate(model, vocab, prompt, n=40):
    ids = vocab.encode(prompt)
    for _ in range(n):
        probs = model.predict(ids)
        ids.append(int(probs.argmax()))
    return vocab.decode(ids)

d = 16
hidden = 32
block = 16
vocab_size = 50
steps = 800
lr = 0.2
corpus = (
    "The little girl went to the park with her mom. "
    "She saw a big green tree and a red bird singing. "
    "The girl laughed and ran across the soft grass. "
    "Her mom opened a basket and they had a picnic. "
    "The sun was warm and the sky was bright blue. "
    "A small brown dog came and wagged its tail. "
    "They shared some bread with the friendly dog. "
    "When it grew late, they walked home together, happy."
)

vocab = Vocabulary(corpus, vocab_size)
ids = vocab.encode(corpus)

model = Transformer(len(vocab.itos), d=d, hidden=hidden, block=block)
X, Y = training_examples(ids, model.block)

print("Training...")
model.train(X, Y, steps=steps, lr=lr)

print()
print(generate(model, vocab, "The little girl went to the", n=12))