# -*- coding: utf-8 -*-
"""Untitled2.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1qAV4U_m8JF3F-icpID2WdjBg1nOfSckE
"""
import os
import sys
# Add the parent directory to sys.path
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# Désactive XLA/TPU pour ne pas interférer avec CUDA
os.environ["PJRT_DEVICE"] = "CUDA"
os.environ["XLA_USE_BF16"] = "0"
os.environ["XLA_DEVICE"] = "CUDA"
os.environ['TOKENIZERS_PARALLELISM'] = 'true'

import torch
import time
import random
import numpy as np
from datasets import load_dataset, get_dataset_split_names, DatasetDict
from transformers import AutoTokenizer, AutoModelForSequenceClassification, DataCollatorWithPadding # Assurez-vous que AdamW est bien ce que vous voulez
from torch.utils.data import DataLoader
from tqdm import tqdm
import os


torch.backends.cudnn.benchmark = True
# Imports pour la précision mixte automatique (AMP)
from torch.amp import autocast
from torch.amp import GradScaler # GradScaler reste ici car spécifique à CUDA
import os

os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"

# J'ajoute des modifications et des nouvelles fonctions ci-dessous.

def set_seed(seed=42):
    os.environ["PYTHONHASHSEED"] = str(seed)
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True
    torch.use_deterministic_algorithms(True, warn_only=True)

def preprocess_data(tokenizer, dataset_dict, max_length=128, num_proc=None): # Ajout de num_proc pour .map()
    """Tokenize et formate le dataset."""
    def tokenize_examples(examples):
        # padding="max_length" ici signifie que tous les exemples sont paddés à max_length globalement.
        # DataCollatorWithPadding s'assurera ensuite de la collation correcte.
        return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=max_length)

    print(f"Tokenization des données (num_proc={num_proc if num_proc else 'auto'})...")
    tokenized_datasets = dataset_dict.map(tokenize_examples, batched=True, num_proc=num_proc)

    if "label" in tokenized_datasets["train"].column_names:
        tokenized_datasets = tokenized_datasets.rename_column("label", "labels")

    tokenized_datasets.set_format("torch", columns=["input_ids", "attention_mask", "labels"])
    return tokenized_datasets

# --- Fonctions d'entraînement et d'évaluation MODIFIÉES ---

def train_epoch(model, dataloader, optimizer, device): # Renommé train en train_epoch
    """Effectue une époque d'entraînement."""
    model.train()
    total_loss, total_correct, total_samples = 0, 0, 0
    for batch in tqdm(dataloader, desc="Training"):
        # Déplacer le batch entier vers le device
        batch = {k: v.to(device) for k, v in batch.items()}

        outputs = model(**batch)
        loss = outputs.loss
        logits = outputs.logits

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_loss += loss.item()
        preds = torch.argmax(logits, dim=-1)
        total_correct += (preds == batch["labels"]).sum().item()
        total_samples += batch["labels"].size(0)

    avg_loss = total_loss / len(dataloader)
    accuracy = total_correct / total_samples
    return avg_loss, accuracy

@torch.no_grad()
def evaluate_model(model, dataloader, device): # Renommé evaluate en evaluate_model
    """Évalue le modèle."""
    model.eval()
    total_loss, total_correct, total_samples = 0, 0, 0
    for batch in dataloader:
        batch = {k: v.to(device) for k, v in batch.items()}
        outputs = model(**batch)
        loss = outputs.loss
        logits = outputs.logits

        total_loss += loss.item()
        preds = torch.argmax(logits, dim=-1)
        total_correct += (preds == batch["labels"]).sum().item()
        total_samples += batch["labels"].size(0)

    avg_loss = total_loss / len(dataloader)
    accuracy = total_correct / total_samples
    return avg_loss, accuracy

import os
import json



from pathlib import Path
import time, math, random
import torch
from torch.utils.data import DataLoader
from datasets import load_dataset, concatenate_datasets, DatasetDict, interleave_datasets
from transformers import (
    BertConfig,
    AutoTokenizer,
    BertForMaskedLM,
    AutoModelForSequenceClassification,
    DataCollatorWithPadding,
    get_scheduler
)
from scipy.stats import pearsonr
from tqdm.auto import tqdm, trange
from utils import save_logs

GLUE_TASKS_INFO = {
    "sst2": {"num_labels": 2, "is_regression": False, "eval_split": "validation"},
    "cola": {"num_labels": 2, "is_regression": False, "eval_split": "validation"},
    "mrpc": {"num_labels": 2, "is_regression": False, "eval_split": "validation"},
    "stsb": {"num_labels": 1, "is_regression": True,  "eval_split": "validation"},
    "qqp":  {"num_labels": 2, "is_regression": False, "eval_split": "validation"},
    "mnli": {"num_labels": 3, "is_regression": False, "eval_split": "validation_matched"},
    "qnli": {"num_labels": 2, "is_regression": False, "eval_split": "validation"},
    "rte":  {"num_labels": 2, "is_regression": False, "eval_split": "validation"},
    "wnli": {"num_labels": 2, "is_regression": False, "eval_split": "validation"},
}
INPUT_COLUMNS = {
    "sst2": ("sentence", None),
    "cola": ("sentence", None),
    "mrpc": ("sentence1", "sentence2"),
    "stsb": ("sentence1", "sentence2"),
    "qqp":  ("question1", "question2"),
    "mnli": ("premise", "hypothesis"),
    "qnli": ("question", "sentence"),
    "rte":  ("sentence1", "sentence2"),
    "wnli": ("sentence1", "sentence2"),
}

def preprocess_glue(tokenizer, raw_ds, task):
    """Tokenise + keep labels as expected by transformers."""
    col1, col2 = INPUT_COLUMNS[task]

    def tok_fn(batch):
        if col2 is None:
            enc = tokenizer(batch[col1], truncation=True, max_length=128)
        else:
            enc = tokenizer(batch[col1], batch[col2], truncation=True, max_length=128)
        # copy label ➜ labels
        enc["labels"] = batch["label"]
        return enc

    keep_cols = [c for c in raw_ds["train"].column_names if c == "label"]
    rm_cols = [c for c in raw_ds["train"].column_names if c not in keep_cols]
    return raw_ds.map(tok_fn, batched=True, remove_columns=rm_cols)


def evaluate_model(model, dataloader, device, is_reg):
    tot_loss, tot = 0.0, 0; correct = 0
    preds_all, labels_all = [], []
    model.eval()
    with torch.no_grad():
        for batch in dataloader:
            batch = {k: v.to(device) for k, v in batch.items()}
            labels = batch["labels"]
            outputs = model(**batch)
            logits = outputs.logits
            loss_fn = torch.nn.MSELoss() if is_reg else torch.nn.CrossEntropyLoss()
            loss = loss_fn(logits.squeeze(-1) if is_reg else logits, labels.float() if is_reg else labels)
            tot_loss += loss.item() * labels.size(0); tot += labels.size(0)
            if is_reg:
                preds_all.append(logits.squeeze(-1).cpu()); labels_all.append(labels.float().cpu())
            else:
                preds = logits.argmax(dim=-1); correct += (preds == labels).sum().item()
    avg_loss = tot_loss / tot
    if is_reg:
        import torch as _t
        from scipy.stats import pearsonr
        preds_all = _t.cat(preds_all).numpy(); labels_all = _t.cat(labels_all).numpy()
        return avg_loss, pearsonr(preds_all, labels_all)[0]
    return avg_loss, correct / tot

# ---------------------------------------------------------------------------
# Fine‑tuning ---------------------------------------------------------------
# ---------------------------------------------------------------------------
def finetune_glue_task(
    task_name: str,
    seed_val: int,
    optimizer_cls: torch.optim.Optimizer,
    optimizer_name: str,
    *,
    num_epochs: int = 3,
    per_device_batch_size: int = 32,
    grad_accum_steps: int = 1,
    lr_scaling_rule: str = "linear",  # "linear" ou "none"
    warmup_ratio: float = 0.1,
    use_small_subset: bool = False,
    optimizer_params: dict[str, any] | None = None,
):
    """Fine‑tune BERT sur une tâche GLUE.

    Args:
        optimizer_cls: classe optimiseur (AdamW, Lion…)
        optimizer_kwargs: kwargs spécifiques transmis à l'optimiseur (ex. betas).
    """

    base_learning_rate = optimizer_params['lr']
    optim_params = optimizer_params.copy()
    del optim_params['lr']
    # ---------------------------------------------------------------------
    # Initialisation – seed et device
    # ---------------------------------------------------------------------
    set_seed(seed_val)
    device = torch.device('cuda')

    info = GLUE_TASKS_INFO[task_name]
    print(f"\n### FINETUNE {task_name} – seed={seed_val} – optimiser={optimizer_name}")

    # ---------------------------------------------------------------------
    # Chargement / pré‑proc des données
    # ---------------------------------------------------------------------
    raw = load_dataset("glue", task_name)
    if use_small_subset:
        raw["train"] = raw["train"].shuffle(seed=seed_val).select(range(5_000))
        raw[info["eval_split"]] = (
            raw[info["eval_split"]].shuffle(seed=seed_val).select(range(200))
        )

    tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
    ds = preprocess_glue(tokenizer, raw, task_name)
    collator = DataCollatorWithPadding(tokenizer)

    train_dl = DataLoader(
        ds["train"],
        batch_size=per_device_batch_size,
        shuffle=True,
        collate_fn=collator,
    )
    eval_dl = DataLoader(
        ds[info["eval_split"]],
        batch_size=per_device_batch_size,
        shuffle=False,
        collate_fn=collator,
    )

    # Model initialisation
    model = AutoModelForSequenceClassification.from_pretrained(
        "bert-base-uncased", num_labels=info["num_labels"]
    ).to(device)
    
    effective_batch = per_device_batch_size * grad_accum_steps
    if lr_scaling_rule == "linear":
        learning_rate = base_learning_rate * (effective_batch / 32)
    else:
        learning_rate = base_learning_rate
        
    if 'weight_decay' not in optim_params:
        raise ValueError(
            "Optimizer parameters must include 'weight_decay' for GLUE fine-tuning."
        )

    # Optimizer initialization
    optim = optimizer_cls(
        model.parameters(), lr=learning_rate, **optim_params
    )

    # Scheduler
    num_update_steps_per_epoch = len(train_dl) // grad_accum_steps
    max_train_steps = num_update_steps_per_epoch * num_epochs
    warmup_steps = int(warmup_ratio * max_train_steps)
    lr_scheduler = get_scheduler(
        "linear",
        optimizer=optim,
        num_warmup_steps=warmup_steps,
        num_training_steps=max_train_steps,
    )

    scaler = GradScaler()
    logs: list[dict[str, float]] = []

    # Training loop
    global_step = 0
    for ep in tqdm(range(1, num_epochs + 1), desc="Epochs", unit="ép"):
        model.train()
        tot_loss, tot_samples = 0.0, 0

        for step, batch in enumerate(
            tqdm(train_dl, desc=f"Train | ép {ep}", leave=False, unit="batch")
        ):
            batch = {k: v.to(device) for k, v in batch.items()}
            with autocast(device_type="cuda"):
                loss = model(**batch).loss / grad_accum_steps 

            scaler.scale(loss).backward()

            if (step + 1) % grad_accum_steps == 0 or (step + 1) == len(train_dl):
                scaler.step(optim)
                scaler.update()
                optim.zero_grad()

                lr_scheduler.step()
                global_step += 1

            # Stats
            bsz = batch["labels"].size(0)
            tot_loss += loss.item() * grad_accum_steps * bsz
            tot_samples += bsz

        tr_loss = tot_loss / tot_samples

        # Evaluation
        ev_loss, ev_metric = evaluate_model(
            model, eval_dl, device, info["is_regression"]
        )
        metric_name = "Acc" if not info["is_regression"] else "Pearson"

        tqdm.write(
            f"Epoch {ep}/{num_epochs} - Train {tr_loss:.3f} | Eval {ev_loss:.3f} | "
            f"{metric_name} {ev_metric*100 if not info['is_regression'] else ev_metric:.2f}"
        )

        logs.append({"epoch": ep, "train": tr_loss, "eval": ev_loss, metric_name: float(ev_metric)})

    # Save logs
    output_dir = f'ablation_study/logs/{task_name}'
    out_dir = Path(output_dir) / f"{optimizer_name}_seed{seed_val}"
    out_dir.mkdir(parents=True, exist_ok=True)
    save_logs(
        optimizer_name,
        seed_val,
        logs,
        folder=output_dir,
    )

if __name__ == "__main__":
    import torch
    from optimizers import *
    
    optimizers_to_run = [#(torch.optim.AdamW, 'Adam', {'lr': 2e-5, 'weight_decay': 1e-2}),
                         #(Yogi, 'Yogi', {'lr': 2e-5, 'weight_decay': 1e-2}),
                         #(Signum, 'Signum', {'lr': 2e-5, 'weight_decay': 1e-2}),
                         #(SignumGrad, 'SignumGrad', {'lr': 2e-5, 'weight_decay': 1e-2}),
                         #(AdamGrad, 'AdamGrad', {'lr': 2e-5, 'weight_decay': 1e-2}),
                         #(Ano, 'Ano', {'lr': 2e-5, 'weight_decay': 1e-2}),
                        #(Anolog, 'Anolog', {'lr': 2e-5, 'weight_decay': 1e-2}),
                        #(Grams, 'Grams', {'lr': 2e-5, 'weight_decay': 1e-2})
                        (Anolog, 'Anolog', {'lr': 2e-5, 'weight_decay': 1e-2}),
                         (Anosqrt, 'Anosqrt', {'lr': 2e-5, 'weight_decay': 1e-2}),
                         (Anoall, 'Anoall', {'lr': 2e-5, 'weight_decay': 1e-2}),
                         #(YogiSignum, 'YogiSignum', {'lr': 2e-5, 'weight_decay': 1e-2})
                  ]
    
    seeds_to_run = [2, 5, 10, 42, 444]
    for seed in seeds_to_run:
        #TASKS = ['stsb', 'cola', 'mrpc', 'rte']
        TASKS = ['mrpc']
        for task in TASKS:
            for opt_class, opt_name, opt_params in optimizers_to_run:
            
                finetune_glue_task(
                    num_epochs=5,
                    task_name=task,
                    seed_val=seed,
                    optimizer_cls=opt_class,
                    optimizer_name=opt_name,
                    optimizer_params=opt_params,
                    use_small_subset=False,
                )
                
                
                
                