import csv
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from torch.nn.utils.rnn import pad_sequence
import math

tokens = list("azertyuiopqsdfghjklmwxcvbnäüöß—– ")
tokensdict = {}

for i in range(len(tokens)):
    tokensdict.update({tokens[i]: [0] * i + [0] * (len(tokens) - (i + 1))})

# Ouvrir le fichier CSV
with open("C:\\Users\\marc2\\Downloads\\7eaaf0e22461b505c749e268c0b72bc4-12ebe211a929f039791dfeaa1a019b64cadddaf1\\7eaaf0e22461b505c749e268c0b72bc4-12ebe211a929f039791dfeaa1a019b64cadddaf1\\top-german-verbs.csv", 'r', encoding="utf-8") as file:
    # Créer un objet lecteur CSV
    reader = [i for i in csv.reader(file)][1:]

class CSVDataset(Dataset):
    def __init__(self, features, labels):
        self.features = features
        self.labels = labels

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

    def __getitem__(self, idx):
        sample = self.features[idx], self.labels[idx]
        return sample

# Supposons que vous ayez vos données sous forme de listes
features = []
labels = []

for i in reader:
    k = []
    for j in i[2]:
        k += [tokens.index(j)+1]
    k += [0] * (25 - len(k))
    features += [torch.Tensor(k)]
    k = []
    for j in i[8]:
        k += [tokens.index(j)+1]
    k += [0] * (25 - len(k))
    labels += [torch.Tensor(k)]

MyDataset = CSVDataset(features=features, labels=labels)

class TransformerModel(nn.Module):
    def __init__(self, vocab_size, emb_dim, nhead, num_encoder_layers, num_decoder_layers, dim_feedforward, dropout=0.1):
        super().__init__()
        self.custom_embedding = nn.Embedding(vocab_size, emb_dim).to("cuda")
        self.pos_encoder = PositionalEncoding(emb_dim, dropout).to("cuda")
        encoder_layer = nn.TransformerEncoderLayer(emb_dim, nhead, dim_feedforward, dropout, batch_first=True).to("cuda")
        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_encoder_layers)
        decoder_layer = nn.TransformerDecoderLayer(emb_dim, nhead, dim_feedforward, dropout, batch_first=True).to("cuda")
        self.transformer_decoder = nn.TransformerDecoder(decoder_layer, num_decoder_layers)
        self.output_layer = nn.Linear(emb_dim, vocab_size).to("cuda")

    def forward(self, src, tgt, src_mask=None, tgt_mask=None, memory_mask=None, src_key_padding_mask=None, tgt_key_padding_mask=None, memory_key_padding_mask=None):
        src_emb = self.custom_embedding(src.long())
        #print("Source Embedding:", src_emb.shape)
        src_emb = self.pos_encoder(src_emb)
        #print("Source Embedding:", src_emb.shape)
        tgt_emb = self.custom_embedding(tgt.long())
        #print("Target Embedding:", tgt_emb.shape)
        tgt_emb = self.pos_encoder(tgt_emb)
        #print("Target Embedding:", tgt_emb.shape)
        encoder_output = self.transformer_encoder(src_emb, src_mask, src_key_padding_mask)
        decoder_output = self.transformer_decoder(tgt_emb, encoder_output, tgt_mask, memory_mask, tgt_key_padding_mask, memory_key_padding_mask)
        output = self.output_layer(decoder_output[:, -1, :])
        return output

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.1, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)

    def forward(self, x):
        x = x + self.pe[:, :x.size(1), :]
        return self.dropout(x)
    
def collate_fn(batch):
    inputs = [item[0].to("cuda") for item in batch]
    targets = [item[1].to("cuda") for item in batch]
    inputs = pad_sequence(inputs, batch_first=True, padding_value=0)
    targets = pad_sequence(targets, batch_first=True, padding_value=0)
    return inputs, targets

train_loader = DataLoader(MyDataset, batch_size=32, shuffle=True, collate_fn=collate_fn)

model = TransformerModel(vocab_size=len(tokens)+1, emb_dim=32, nhead=8, num_encoder_layers=6, num_decoder_layers=6, dim_feedforward=2048)
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

epochs = 1

for epoch in range(epochs):
    total_loss = 0.0

    for batch_idx, (inputs, targets) in enumerate(train_loader):
        for i in range(1, targets.shape[1]):
            optimizer.zero_grad()
            output = model(inputs, targets[:, :i])  # Shifted targets
            #print(output.shape)
            loss = loss_fn(output, targets[:, i].long())  # Reshape targets
            loss.backward()
            optimizer.step()

            total_loss += loss.item()

        print(f"Epoch {epoch + 1}/{epochs}, Batch {batch_idx}/{len(train_loader)}, Loss: {total_loss / (batch_idx + 1)}")

    print(f"Epoch {epoch + 1}/{epochs}, Loss: {total_loss / len(train_loader)}")

torch.save(model, "data/PrateritumGPT.pth")