import numpy as np

class MiniGPT:
    """
    Mini-GPT: Decoder-only Transformer Language Model
    Built from scratch — everything we learned in Pages 1-9!
    """
    def __init__(self, vocab_size, d_model=64, n_heads=4,
                 n_layers=2, d_ff=128, max_len=128):
        self.d_model = d_model
        self.vocab_size = vocab_size

        # Token Embedding (Page 6)
        self.token_emb = np.random.randn(vocab_size, d_model) * 0.02

        # Positional Encoding (Page 8)
        self.pos_enc = np.zeros((max_len, d_model))
        pos = np.arange(max_len)[:, np.newaxis]
        div = np.exp(np.arange(0, d_model, 2) * -(np.log(10000) / d_model))
        self.pos_enc[:, 0::2] = np.sin(pos * div)
        self.pos_enc[:, 1::2] = np.cos(pos * div)

        # Transformer Blocks (Page 8)
        self.blocks = []
        for _ in range(n_layers):
            self.blocks.append(TransformerBlock(d_model, n_heads, d_ff))

        # Output Head
        self.W_out = np.random.randn(d_model, vocab_size) * 0.02
        self.b_out = np.zeros((1, vocab_size))

    def forward(self, token_indices):
        """
        token_indices: (batch, seq_len) — integer indices
        returns: logits (batch, seq_len, vocab_size)
        """
        B, S = token_indices.shape

        # Embed tokens + add position
        x = self.token_emb[token_indices] + self.pos_enc[:S]

        # Causal mask: can't look at future tokens!
        mask = np.tril(np.ones((S, S)))[np.newaxis, np.newaxis, :, :]

        # Pass through Transformer blocks
        for block in self.blocks:
            x = block.forward(x, mask)

        # Project to vocabulary
        logits = x @ self.W_out + self.b_out
        return logits

    def generate(self, start_tokens, max_new=50, temperature=0.8):
        """Autoregressive text generation"""
        tokens = list(start_tokens)
        for _ in range(max_new):
            x = np.array([tokens])
            logits = self.forward(x)
            # Take last position, apply temperature
            next_logits = logits[0, -1] / temperature
            probs = np.exp(next_logits - np.max(next_logits))
            probs /= probs.sum()
            # Sample
            next_token = np.random.choice(len(probs), p=probs)
            tokens.append(next_token)
        return tokens

# =====================================================
# TRAINING on Shakespeare text
# =====================================================
print("🏆 Mini-GPT: The Grand Finale!")
print("   Architecture: 2-layer Transformer, d=64, 4 heads")
print("   Training: next-token prediction on Shakespeare")
print("   This is the SAME architecture as GPT — just smaller!")