import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
import numpy as np

class SimpleT5Tokenizer:
    """Simulates T5's mapping of text strings to integer IDs."""
    def __init__(self, vocab=None):
        self.vocab = vocab or {}
        self.inv_vocab = {v: k for k, v in self.vocab.items()}

    def encode(self, text):
        return [self.vocab.get(word, self.vocab.get(' unbeknownst', 1)) for word in text.split()]

    def decode(self, ids):
        return " ".join([self.inv_vocab.get(i, ' unbeknownst') for i in ids])

class T5Model(nn.Module):
    """Modular T5 Encoder-Decoder implementation."""
    def __init__(self, vocab_size, d_model=128, nhead=4, num_layers=2):
        super(T5Model, self).__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.transformer = nn.Transformer(
            d_model=d_model, nhead=nhead, 
            num_encoder_layers=num_layers, num_decoder_layers=num_layers, 
            batch_first=True
        )
        # The "Head" is just a linear layer mapping back to the vocabulary
        self.lm_head = nn.Linear(d_model, vocab_size)

    def forward(self, src, tgt):
        src_emb = self.embedding(src)
        tgt_emb = self.embedding(tgt)
        
        # Causal mask prevents the decoder from 'cheating' by looking at future tokens
        tgt_mask = self.transformer.generate_square_subsequent_mask(tgt.size(1)).to(src.device)
        
        out = self.transformer(src_emb, tgt_emb, tgt_mask=tgt_mask)
        return self.lm_head(out)

class TextToTextDataset(Dataset):
    """Converts (X, y) into 'classify: [features]' -> '[label]'"""
    def __init__(self, X, y, tokenizer):
        self.X, self.y, self.tokenizer = X, y, tokenizer

    def __len__(self): return len(self.X)

    def __getitem__(self, idx):
        feat_str = "classify: " + " ".join([f"{val:.2f}" for val in self.X[idx]])
        label_str = "positive" if self.y[idx] == 1 else "negative"
        
        src_ids = torch.tensor(self.tokenizer.encode(feat_str), dtype=torch.long)
        tgt_ids_in = torch.tensor(self.tokenizer.encode(label_str), dtype=torch.long)
        tgt_ids_out = torch.tensor(self.tokenizer.encode(label_str + " </s>"), dtype=torch.long)
        
        return src_ids, tgt_ids_in, tgt_ids_out

def collate_fn(batch):
    src_list, tgt_in_list, tgt_out_list = zip(*batch)
    def pad(seqs):
        max_len = max(len(s) for s in seqs)
        padded = torch.full((len(seqs), max_len), 0, dtype=torch.long)
        for i, s in enumerate(seqs): padded[i, :len(s)] = s
        return padded
    return pad(src_list), pad(tgt_in_list), pad(tgt_out_list)

# --- Execution ---
if __name__ == '__main__':
    X, y = make_classification(n_samples=1000, n_features=10, random_state=42)
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)

    vocab = {'<pad>': 0, '<unk>': 1, '</s>': 2, 'classify:': 3, 'positive': 4, 'negative': 5}
    for i in range(100): vocab[f"{i/10:.2f}"] = 10 + i 
    
    tokenizer = SimpleT5Tokenizer(vocab)
    train_loader = DataLoader(TextToTextDataset(X_train, y_train, tokenizer), batch_size=16, collate_fn=collate_fn)
    test_loader = DataLoader(TextToTextDataset(X_test, y_test, tokenizer), batch_size=16, collate_fn=collate_fn)

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = T5Model(vocab_size=len(vocab) + 1000).to(device)
    optimizer = optim.Adam(model.parameters(), lr=1e-4)
    criterion = nn.CrossEntropyLoss(ignore_index=0)

    model.train()
    for epoch in range(5):
        for src, tgt_in, tgt_out in train_loader:
            src, tgt_in, tgt_out = src.to(device), tgt_in.to(device), tgt_out.to(device)
            optimizer.zero_grad()
            logits = model(src, tgt_in)
            loss = criterion(logits.view(-1, logits.size(-1)), tgt_out.view(-1))
            loss.backward()
            optimizer.step()
        print(f"Epoch {epoch+1} complete.")

    model.eval()
    all_preds, all_labels = [], []
    with torch.no_grad():
        for src, tgt_in, tgt_out in test_loader:
            logits = model(src.to(device), tgt_in.to(device))
            preds = torch.argmax(logits[:, 0, :], dim=-1)
            all_preds.extend(preds.cpu().numpy())
            all_labels.extend(tgt_out[:, 0].numpy())

    final_preds = [1 if p == 4 else 0 for p in all_preds]
    final_labels = [1 if l == 4 else 0 for l in all_labels]
    print(f"\nTest Accuracy: {accuracy_score(final_labels, final_preds):.4f}")