import torch
from torch.utils.data import Dataset, DataLoader
from sentence_transformers import SentenceTransformer, losses
from tqdm import tqdm
import gc
from plugins.scansite import ScansitePlugin

torch.cuda.empty_cache()

class PreferenceDataset(Dataset):
    def __init__(self, data, tokenizer, max_length=128):
        self.data = data
        self.tokenizer = tokenizer
        self.max_length = max_length

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

    def __getitem__(self, idx):
        url, title, score = self.data[idx]
        encoded = self.tokenizer(title, padding='max_length', truncation=True, max_length=self.max_length, return_tensors="pt")
        return {key: val.squeeze(0) for key, val in encoded.items()}, torch.tensor(score, dtype=torch.float)

def collate_fn(batch):
    input_ids = torch.stack([item[0]['input_ids'] for item in batch])
    attention_mask = torch.stack([item[0]['attention_mask'] for item in batch])
    scores = torch.stack([item[1] for item in batch])
    return {'input_ids': input_ids, 'attention_mask': attention_mask}, scores

def finetune(model_name='nomic-ai/nomic-embed-text-v1', output_model_name="embeddings-ft", num_epochs=2, learning_rate=2e-5, weight_decay=0.01, batch_size=8, num_warmup_steps=0):
    print(f"Fine-tuning parameters:\n"
          f"num_epochs: {num_epochs}\n"
          f"learning rate (lr): {learning_rate}\n"
          f"weight_decay: {weight_decay}\n"
          f"batch_size: {batch_size}\n"
          f"model_name: {model_name}\n"
          f"num_warmup_steps: {num_warmup_steps}")

    scansite_plugin = ScansitePlugin("scansite", None)
    reference_data_valid, reference_data_rejected = scansite_plugin.get_reference_data()

    valid_data_with_scores = [(url, title, (score - 1) / 8 + 0.5) for url, title, score in reference_data_valid]
    rejected_data_with_scores = [(url, title, 0.0) for url, title in reference_data_rejected]

    all_data = valid_data_with_scores + rejected_data_with_scores

    model = SentenceTransformer(model_name, trust_remote_code=True)
    tokenizer = model.tokenizer

    dataset = PreferenceDataset(all_data, tokenizer)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)

    loss_function = torch.nn.MSELoss()

    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=weight_decay)

    total_steps = len(dataloader) * num_epochs
    scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.1, total_iters=total_steps)

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model.to(device)

    for epoch in range(num_epochs):
        model.train()
        for batch in tqdm(dataloader, desc=f"Epoch {epoch + 1}/{num_epochs}"):
            input_data, scores = batch
            input_data = {k: v.to(device) for k, v in input_data.items()}
            scores = scores.to(device)

            optimizer.zero_grad()

            embeddings = model(input_data)['sentence_embedding']

            # Calcul de la similarité cosinus
            embeddings_norm = torch.nn.functional.normalize(embeddings, p=2, dim=1)
            cosine_similarities = torch.sum(embeddings_norm, dim=1)

            # Calcul de la perte
            loss = loss_function(cosine_similarities, scores)

            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()
            scheduler.step()

            del embeddings, cosine_similarities
            torch.cuda.empty_cache()
            gc.collect()

    model.save(output_model_name)

    print("Finetuning terminé et modèle sauvegardé.")

if __name__ == "__main__":
    finetune()