#!/usr/bin/env python3
import torch
import torch.nn as nn

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# TOKENIZER CLASS (THIS WAS MISSING!)
class CharTokenizer:
    def __init__(self, text=None):
        if text is not None:
            self.chars = sorted(list(set(text)))
            self.vocab_size = len(self.chars)
            self.char_to_idx = {ch: i for i, ch in enumerate(self.chars)}
            self.idx_to_char = {i: ch for i, ch in enumerate(self.chars)}
        else:
            self.chars = []
            self.vocab_size = 0
            self.char_to_idx = {}
            self.idx_to_char = {}
    
    def encode(self, text):
        return [self.char_to_idx[ch] for ch in text if ch in self.char_to_idx]
    
    def decode(self, indices):
        return ''.join([self.idx_to_char.get(i, '') for i in indices])

# MODEL ARCHITECTURE
class TransformerLanguageModel(nn.Module):
    def __init__(self, vocab_size, d_model=256, nhead=8, num_layers=6, dropout=0.2, seq_length=128):
        super().__init__()
        self.d_model = d_model
        self.seq_length = seq_length
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pos_encoding = nn.Embedding(seq_length, d_model)
        
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=nhead,
            dim_feedforward=d_model * 4,
            dropout=dropout,
            batch_first=True
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        
        self.dropout = nn.Dropout(dropout)
        self.fc_out = nn.Linear(d_model, vocab_size)
    
    def forward(self, x):
        batch_size, seq_len = x.shape
        token_emb = self.embedding(x)
        positions = torch.arange(seq_len, device=x.device).unsqueeze(0).expand(batch_size, -1)
        pos_emb = self.pos_encoding(positions)
        x = self.dropout(token_emb + pos_emb)
        mask = nn.Transformer.generate_square_subsequent_mask(seq_len).to(x.device)
        x = self.transformer(x, mask=mask, is_causal=True)
        logits = self.fc_out(x)
        return logits
    
    def generate(self, tokenizer, prompt="", max_length=500, temperature=0.8):
        import numpy as np
        self.eval()
        
        if prompt == "":
            indices = [np.random.randint(0, tokenizer.vocab_size)]
        else:
            indices = tokenizer.encode(prompt)
        
        with torch.no_grad():
            for _ in range(max_length):
                x = torch.tensor(indices[-self.seq_length:], dtype=torch.long).unsqueeze(0).to(device)
                
                if x.shape[1] < self.seq_length:
                    padding = torch.zeros(1, self.seq_length - x.shape[1], dtype=torch.long).to(device)
                    x = torch.cat([padding, x], dim=1)
                
                logits = self(x)
                logits = logits[0, -1, :] / temperature
                probs = torch.softmax(logits, dim=-1)
                next_idx = torch.multinomial(probs, num_samples=1).item()
                indices.append(next_idx)
        
        return tokenizer.decode(indices)

# LOAD MODEL
print("Loading model...")
checkpoint = torch.load('best_model.pth', map_location=device, weights_only=False)
tokenizer = checkpoint['tokenizer']

model = TransformerLanguageModel(
    vocab_size=tokenizer.vocab_size,
    d_model=256,
    nhead=8,
    num_layers=6,
    dropout=0.2,
    seq_length=128
).to(device)

model.load_state_dict(checkpoint['model_state_dict'])
print("✓ Model loaded!\n")

# INTERACTIVE GENERATION
print("="*70)
print("📝 Shakespeare Text Generation")
print("="*70)
print("Commands:")
print("  - Type a prompt and press Enter")
print("  - Type 'quit' to stop")
print("  - Press Enter for random generation")
print("\nTry: ROMEO:  or  To be or not to be")
print("="*70 + "\n")

while True:
    prompt = input("Enter prompt (or 'quit'): ")
    
    if prompt.lower() in ['quit', 'exit', 'q']:
        print("Goodbye!")
        break
    
    print("\n" + "="*70)
    print("Generating...\n")
    generated = model.generate(tokenizer, prompt=prompt, max_length=300, temperature=0.8)
    print(generated)
    print("="*70 + "\n")