import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
import numpy as np
import matplotlib.pyplot as plt
from typing import List

class UnifiedTransformer(nn.Module):
    """
    A Decoder-only Transformer for mixed-modal next-token prediction.
    Treats all modalities as sequences of discrete tokens.
    """
    def __init__(self, vocab_size: int, d_model: int = 256, nhead: int = 8, 
                 num_layers: int = 4, dim_feedforward: int = 1024, max_seq_len: int = 512):
        super().__init__()
        self.vocab_size = vocab_size
        self.d_model = d_model
        
        # Shared embedding space for all modalities
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pos_embedding = nn.Parameter(torch.zeros(1, max_seq_len, d_model))
        
        decoder_layer = nn.TransformerDecoderLayer(
            d_model=d_model, nhead=nhead, dim_feedforward=dim_feedforward, batch_first=True
        )
        self.transformer = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)
        self.lm_head = nn.Linear(d_model, vocab_size)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        batch_size, seq_len = x.shape
        x = self.embedding(x) + self.pos_embedding[:, :seq_len, :]
        
        # Causal mask to ensure the model only looks at the past
        mask = nn.Transformer.generate_square_subsequent_mask(seq_len).to(x.device)
        
        # In decoder-only mode, memory is the same as the target sequence
        out = self.transformer(x, x, tgt_mask=mask, memory_mask=mask)
        return self.lm_head(out)

class MixedModalDataset(Dataset):
    """
    Simulates mixed-modal data by partitioning the vocabulary.
    Text: 0-999 | Image: 1000-1999 | Speech: 2000-2999
    """
    def __init__(self, num_samples: int = 1000, seq_len: int = 64, 
                 modalities: List[str] = ['text', 'image', 'speech'], 
                 vocab_per_modality: int = 1000):
        self.num_samples = num_samples
        self.seq_len = seq_len
        self.modalities = modalities
        self.vocab_per_modality = vocab_per_modality
        self.modality_offsets = {mod: i * vocab_per_modality for i, mod in enumerate(modalities)}

    def __len__(self):
        return self.num_samples

    def __getitem__(self, idx):
        mod = np.random.choice(self.modalities)
        offset = self.modality_offsets[mod]
        
        # Generate synthetic sequence with local structure (Markov-like)
        seq = np.zeros(self.seq_len + 1, dtype=np.int64)
        current_token = np.random.randint(offset, offset + self.vocab_per_modality)
        
        for i in range(self.seq_len + 1):
            seq[i] = current_token
            if np.random.rand() < 0.8: # Local coherence
                current_token = np.clip(current_token + np.random.randint(-2, 3), 
                                        offset, offset + self.vocab_per_modality - 1)
            else:
                current_token = np.random.randint(offset, offset + self.vocab_per_modality)
        
        return torch.tensor(seq[:-1]), torch.tensor(seq[1:])

# --- Training Execution ---
if __name__ == '__main__':
    # Hyperparameters
    MODALITIES = ['text', 'image', 'speech']
    VOCAB_PER_MOD = 1000
    TOTAL_VOCAB = len(MODALITIES) * VOCAB_PER_MOD
    SEQ_LEN, BATCH_SIZE, D_MODEL, EPOCHS = 64, 32, 128, 10

    dataset = MixedModalDataset(num_samples=2000, seq_len=SEQ_LEN, modalities=MODALITIES)
    dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)

    model = UnifiedTransformer(vocab_size=TOTAL_VOCAB, d_model=D_MODEL, num_layers=3, max_seq_len=SEQ_LEN+1)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.AdamW(model.parameters(), lr=1e-3)

    history = []
    for epoch in range(EPOCHS):
        model.train()
        total_loss = 0
        for batch_x, batch_y in dataloader:
            optimizer.zero_grad()
            logits = model(batch_x)
            loss = criterion(logits.view(-1, TOTAL_VOCAB), batch_y.view(-1))
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
        
        avg_loss = total_loss / len(dataloader)
        history.append(avg_loss)
        print(f"Epoch {epoch+1}/{EPOCHS} | Loss: {avg_loss:.4f}")

    plt.plot(history, marker='o')
    plt.title("Mixed-Modal Convergence"); plt.xlabel("Epoch"); plt.ylabel("Loss"); plt.grid(True); plt.show()