import torch
import torch.nn as nn
import time
from torchtext.datasets import PennTreebank
from torchtext.data.functional import to_map_style_dataset
from torchtext.data.utils import get_tokenizer
from torchtext.vocab import build_vocab_from_iterator
from torch.utils.data import DataLoader

# 1. Define the model
class TransformerModel(nn.Module):
    def __init__(self, vocab_size, d_model, nhead, num_layers, dropout=0.1):
        super(TransformerModel, self).__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.transformer = nn.Transformer(d_model, nhead, num_layers, dropout=dropout)
        self.fc = nn.Linear(d_model, vocab_size)

    def forward(self, src, tgt):
        src = self.embedding(src)
        tgt = self.embedding(tgt)
        x = self.transformer(src, tgt)
        x = self.fc(x)
        return x


# 2. Preparar datos
train_data_raw = to_map_style_dataset(PennTreebank(split='train'))
tokenizer = get_tokenizer("spacy", "en_core_web_sm")

# Agregar tokens especiales (UNK y PAD) al vocabulario
specials = ['<unk>','<pad>']
vocab = build_vocab_from_iterator((tokenizer(y) for y in train_data_raw), specials=specials)
vocab.set_default_index(vocab['<unk>'])

# Dividir los datos en conjuntos de entrenamiento y validación
train_ratio = 0.8
train_data_size = int(len(train_data_raw) * train_ratio)
train_data = train_data_raw[:train_data_size]
valid_data = train_data_raw[train_data_size:]


# 3. Entrenar el modelo
def generate_pairs(tokens, shift=1):
    source = [tokens[i] for i in range(0, len(tokens) - shift)]
    target = [tokens[i] for i in range(shift, len(tokens))]
    return source, target

def tokenize_data(text, max_seq_len=None):
    tokens = tokenizer(text)[:max_seq_len]
    return torch.tensor([vocab[token] if token in vocab else vocab['<unk>'] for token in tokens], dtype=torch.long)

def create_batches(data, batch_size, max_seq_len=None):
    batched_data = []
    for i in range(0, len(data), batch_size):
        batched_text = []
        for text in data[i:i + batch_size]:
            tokens = tokenize_data(text, max_seq_len=max_seq_len)
            if len(tokens) <= 1:
                continue
            batched_text.append(tokens)
        batched_data.append(batched_text)
    return batched_data

def collate_fn(batch):
    batch_srcs = []
    batch_tgts = []
    max_len = max([len(text) for text in batch])
    for text in batch:
        source, target = generate_pairs(text)
        source = torch.tensor(source, dtype=torch.long)
        target = torch.tensor(target, dtype=torch.long)
        source = torch.cat([source, torch.full((max_len - len(source),), vocab['<pad>'], dtype=torch.long)], dim=0)
        target = torch.cat([target, torch.full((max_len - len(target),), vocab['<pad>'], dtype=torch.long)], dim=0)
        batch_srcs.append(source)
        batch_tgts.append(target)
    return torch.stack(batch_srcs).T, torch.stack(batch_tgts).T


vocab_size = len(vocab)
d_model = 512
nhead = 8  # Ajuste de hparámetro
num_layers = 6  # Ajuste de hiperparámetro
num_epochs = 45  # Increase the number of epochs
learning_rate = 1e-3  # Experiment with different learning rates
optimizer_choice = "Adam"  # Choose the optimizer: "Adam", "RMSprop", or "Adagrad"
batch_size = 50  # Ajuste de hiperparámetro
max_seq_len = 50  # Limitar la longitud máxima de la secuencia para evitar problemas de memoria
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

train_batches = create_batches(train_data, batch_size, max_seq_len=max_seq_len)
valid_batches = create_batches(valid_data, batch_size, max_seq_len=max_seq_len)

train_loader = DataLoader(train_batches, batch_size=None, collate_fn=collate_fn)
valid_loader = DataLoader(valid_batches, batch_size=None, collate_fn=collate_fn)

model = TransformerModel(vocab_size=vocab_size, d_model=d_model, nhead=nhead, num_layers=num_layers, dropout=0.1).to(device)
loss_function = nn.CrossEntropyLoss()

# Use different optimizers
if optimizer_choice == "Adam":
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
elif optimizer_choice == "RMSprop":
    optimizer = torch.optim.RMSprop(model.parameters(), lr=learning_rate)
elif optimizer_choice == "Adagrad":
    optimizer = torch.optim.Adagrad(model.parameters(), lr=learning_rate)
else:
    raise ValueError("Invalid optimizer choice")

def train(model, epoch):
    model.train()
    train_loss = 0
    for i, (source, target) in enumerate(train_loader):
        source = source.to(device)
        target = target.to(device)
        optimizer.zero_grad()
        output = model(source, target[:-1])
        loss = loss_function(output.reshape(-1, vocab_size), target[1:].reshape(-1))
        loss.backward()
        optimizer.step()
        train_loss += loss.item()
        if i % 1000 == 0:
            print(f'Epoch: {epoch + 1},\tLoss: {loss.item()}')
    return train_loss / len(train_loader)


for epoch in range(num_epochs):
    start_time = time.time()
    train_loss = train(model, epoch)
    end_time = time.time()
    elapsed_time = end_time - start_time
    print(f"Epoch: {epoch + 1},\tAverage Loss: {train_loss},\tTime taken: {elapsed_time} seconds")

# 4. Evaluar el modelo en el conjunto de validación
def evaluate(model, valid_loader):
    model.eval()
    valid_loss = 0
    with torch.no_grad():
        for source, target in valid_loader:
            source = source.to(device)
            target = target.to(device)
            output = model(source, target[:-1])
            loss = loss_function(output.reshape(-1, vocab_size), target[1:].reshape(-1))
            valid_loss += loss.item()
    return valid_loss / len(valid_loader)


# 5. Entrenar y evaluar el modelo en cada época
for epoch in range(num_epochs):
    start_time = time.time()
    train_loss = train(model, epoch)
    valid_loss = evaluate(model, valid_loader)
    end_time = time.time()
    elapsed_time = end_time - start_time
    print(f"Epoch: {epoch + 1},\tAverage Train Loss: {train_loss},\tAverage Valid Loss: {valid_loss},\tTime taken: {elapsed_time} seconds")


# 6. Guardar el modelo entrenado
torch.save(model.state_dict(), "trained_model.pth")

# 7. Cargar el modelo
def load_model(path, device):
    loaded_model = TransformerModel(vocab_size=vocab_size, d_model=d_model, nhead=nhead, num_layers=num_layers).to(device)
    loaded_model.load_state_dict(torch.load(path))
    loaded_model.eval()
    return loaded_model

loaded_model = load_model("trained_model.pth", device)

# 8. Predecir texto
def decode_output(output):
    probabilities = torch.softmax(output, dim=2)
    predicted_indices = torch.argmax(probabilities, dim=2)
    return " ".join(vocab.lookup_tokens(predicted_indices.cpu().numpy().squeeze().tolist()))

def predict_next_token(input_text, model):
    tokens = tokenize_data(input_text)
    source, target = generate_pairs(tokens, shift=1)
    source = torch.tensor(source, dtype=torch.long).unsqueeze(1).to(device)
    target = torch.tensor(target, dtype=torch.long).unsqueeze(1).to(device)
    with torch.no_grad():
        output = model(source, target[:-1])
        return decode_output(output)

input_text = "My favorite color is"
predicted_output = predict_next_token(input_text, loaded_model)
print(f"Predicted output: {predicted_output}")