import torch
import torch.nn as nn
from torchtext.datasets import WikiText2
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, valid_data_raw, _ = WikiText2(split=('train', 'valid', 'test'))
train_data_raw = to_map_style_dataset(train_data_raw)
valid_data_raw = to_map_style_dataset(valid_data_raw)

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>'])

vocab_size = len(vocab)
d_model = 512      # Reduced to 512
nhead = 8
num_layers = 4     # Reduced to 4
num_epochs = 50    # Increased to 50
learning_rate = 5e-4 # Reduced to 5e-4
optimizer_choice = 'Adam'
batch_size = 128    # Increased to 128
max_seq_len = 100   # Increased to 100
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

train_batches = create_batches(train_data_raw, batch_size, max_seq_len=max_seq_len)
valid_batches = create_batches(valid_data_raw, 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}")