import numpy as np

class TinyLM:
    """
    Word-level language model.
    Predicts the next word from the previous CONTEXT_SIZE words.
    Supports batched forward/backward passes for fast training.
    """

    def __init__(self, vocab_size, embed_dim=64, hidden_dim=128, context_size=5):
        self.vocab_size = vocab_size
        self.embed_dim = embed_dim
        self.hidden_dim = hidden_dim
        self.context_size = context_size

        rng = np.random.default_rng(42)
        self.embeddings = rng.standard_normal((vocab_size, embed_dim)).astype(np.float32) * 0.02
        self.W1 = rng.standard_normal((embed_dim * context_size, hidden_dim)).astype(np.float32) * 0.02
        self.b1 = np.zeros((1, hidden_dim), dtype=np.float32)
        self.W2 = rng.standard_normal((hidden_dim, vocab_size)).astype(np.float32) * 0.02
        self.b2 = np.zeros((1, vocab_size), dtype=np.float32)

    def num_params(self):
        return (self.embeddings.size + self.W1.size + self.b1.size
                + self.W2.size + self.b2.size)

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

    def forward(self, token_ids):
        """Single-example forward pass (used at inference time)."""
        embs = [self.embeddings[t] for t in token_ids]
        x = np.concatenate(embs).reshape(1, -1)
        h = np.tanh(np.dot(x, self.W1) + self.b1)
        logits = np.dot(h, self.W2) + self.b2
        probs = self.softmax(logits)
        return x, h, probs

    def forward_batch(self, token_id_batch):
        """
        token_id_batch: (B, context_size) int array
        Returns x (B, embed_dim*context_size), h (B, hidden_dim), probs (B, vocab_size)
        """
        B = token_id_batch.shape[0]
        x = self.embeddings[token_id_batch].reshape(B, -1)
        h = np.tanh(x @ self.W1 + self.b1)
        logits = h @ self.W2 + self.b2
        probs = self.softmax(logits)
        return x, h, probs

    def train_step(self, token_ids, target, lr=0.01):
        """Single-example training step (kept for compatibility)."""
        return self.train_batch(
            np.array([token_ids], dtype=np.int64),
            np.array([target], dtype=np.int64),
            lr=lr,
        )

    def train_batch(self, token_id_batch, target_batch, lr=0.01):
        """
        Vectorized batched training step.
        token_id_batch: (B, context_size) int array
        target_batch: (B,) int array
        Returns mean loss over the batch.
        """
        B = token_id_batch.shape[0]
        x, h, probs = self.forward_batch(token_id_batch)

        target_probs = probs[np.arange(B), target_batch]
        loss = -np.mean(np.log(target_probs + 1e-9))

        dout = probs.copy()
        dout[np.arange(B), target_batch] -= 1.0
        dout /= B

        dW2 = h.T @ dout
        db2 = dout.sum(axis=0, keepdims=True)
        dh = dout @ self.W2.T
        dtanh = (1.0 - h * h) * dh
        dW1 = x.T @ dtanh
        db1 = dtanh.sum(axis=0, keepdims=True)
        dx = dtanh @ self.W1.T

        self.W2 -= lr * dW2
        self.b2 -= lr * db2
        self.W1 -= lr * dW1
        self.b1 -= lr * db1

        dx = dx.reshape(B, self.context_size, self.embed_dim)
        np.add.at(self.embeddings, token_id_batch, -lr * dx)

        return loss

    def sample(self, probs, temperature=0.6, top_k=10):
        logits = np.log(np.clip(probs[0], 1e-9, 1.0)) / max(temperature, 0.05)
        if top_k and top_k < len(logits):
            idx = np.argpartition(logits, -top_k)[-top_k:]
            mask = np.full_like(logits, -1e9)
            mask[idx] = logits[idx]
            logits = mask
        exp = np.exp(logits - np.max(logits))
        p = exp / np.sum(exp)
        return int(np.random.choice(len(p), p=p))