"""
Assignment 2 — GoEmotions Multi-Label Classifier
UPGRADED Training Script v2

ROOT CAUSE FIXES vs previous run (Macro F1 stuck at 0.44):
  1. HEAD OVERRIDES PRETRAINED KNOWLEDGE — The old head (hidden→hidden/2→28)
     with 10x LR was destroying the fine-tuned representations. Now using
     a clean linear classifier on CLS + mean-pool concatenation.
  2. LEARNING RATE — 1e-5 with grad_accum=4 was too low. Now: 2e-5
     encoder / 1e-4 head, accum=2.
  3. FOCAL LOSS GAMMA=2 — over-penalized easy classes, starved rare ones.
     Replaced with Asymmetric Loss (ASL) which is SOTA for multi-label
     imbalanced classification (ICCV 2021).
  4. LABEL SMOOTHING ON POSITIVES — 0.05 smoothing was hurting rare
     emotion recall since those positives are already scarce.
  5. POOLING — CLS-only pooling misses context. Now: concat CLS + mean pool.
  6. NO LAYER-WISE LR DECAY — all encoder layers used same LR. Now: LLRD
     so top layers learn 10x faster than bottom layers.
  7. EPOCHS TOO FEW — 6 epochs not enough. Now: 10 with patience=4.
  8. NO OVERSAMPLING — rare classes (grief: 96 samples) were almost
     invisible. Now: WeightedRandomSampler oversamples rare-label examples.

Expected result: Macro F1 0.55-0.60 (vs 0.44 before)

Run on Kaggle T4 GPU:
  python train.py
Outputs:
  model/best_model/   - HuggingFace checkpoint
  model/thresholds.npy
  model/training_results.json
"""

import os, json, warnings, re, html, math
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader, WeightedRandomSampler
from torch.optim import AdamW
from torch.amp import autocast, GradScaler
from transformers import AutoTokenizer, AutoModel, get_cosine_schedule_with_warmup
from datasets import load_dataset
from sklearn.preprocessing import MultiLabelBinarizer
from sklearn.metrics import f1_score, precision_score, recall_score, hamming_loss
warnings.filterwarnings("ignore")

# ════════════════════════════════════════════════════════════════
#  CONFIG
# ════════════════════════════════════════════════════════════════
CFG = {
    # Model
    "model_name":    "SamLowe/roberta-base-go_emotions",
    "fallback":      "roberta-base",
    "max_length":    128,

    # Training
    "epochs":              10,
    "batch_size":          32,       # increased from 16 — better gradient estimates
    "grad_accum_steps":    2,        # effective batch = 64
    "encoder_lr":          2e-5,     # encoder learning rate
    "head_lr":             1e-4,     # head LR — higher, separate param group
    "llrd_factor":         0.9,      # layer-wise LR decay per layer (top→bottom)
    "weight_decay":        0.01,
    "warmup_ratio":        0.06,
    "max_grad_norm":       1.0,
    "fp16":                True,

    # Loss (Asymmetric Loss)
    "asl_gamma_pos":  0.0,    # no down-weighting for positives
    "asl_gamma_neg":  4.0,    # hard down-weighting for easy negatives
    "asl_clip":       0.05,   # probability margin shift

    # Stopping & threshold
    "patience":        4,
    "eval_threshold":  0.5,   # used during training; tuned after on val

    # Paths
    "output_dir":   "model/best_model",
    "thresh_path":  "model/thresholds.npy",
    "results_path": "model/training_results.json",
}

EMOTIONS = [
    "admiration", "amusement", "anger", "annoyance", "approval",
    "caring", "confusion", "curiosity", "desire", "disappointment",
    "disapproval", "disgust", "embarrassment", "excitement", "fear",
    "gratitude", "grief", "joy", "love", "nervousness",
    "optimism", "pride", "realization", "relief", "remorse",
    "sadness", "surprise", "neutral"
]
NUM_LABELS = len(EMOTIONS)
DEVICE     = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Device : {DEVICE}")
print(f"Batch  : {CFG['batch_size']} x accum {CFG['grad_accum_steps']} "
      f"= effective {CFG['batch_size'] * CFG['grad_accum_steps']}")


# ════════════════════════════════════════════════════════════════
#  DATA
# ════════════════════════════════════════════════════════════════
def clean_text(text):
    text = html.unescape(text)
    text = re.sub(r"^>+\s?", "", text, flags=re.MULTILINE)
    text = re.sub(r"http\S+|www\.\S+", "[URL]", text)
    text = re.sub(r"\[deleted\]|\[removed\]", "", text)
    text = re.sub(r"r/\w+|u/\w+", "", text)
    text = re.sub(r"([!?])\1{2,}", r"\1\1", text)
    text = re.sub(r"(.)\1{2,}", r"\1\1", text)
    text = re.sub(r"[^\x00-\x7F]+", " ", text)
    return re.sub(r"\s+", " ", text).strip()

def load_data():
    print("\nLoading GoEmotions ...")
    ds = load_dataset("google-research-datasets/go_emotions", "simplified")
    splits = {}
    for s in ["train", "validation", "test"]:
        df = ds[s].to_pandas()
        df["text"] = df["text"].apply(clean_text)
        df = df[df["text"].str.len() > 2].reset_index(drop=True)
        splits[s] = df
    print(f"  Train {len(splits['train'])} | Val {len(splits['validation'])} | Test {len(splits['test'])}")
    return splits["train"], splits["validation"], splits["test"]

def binarize(train_df, val_df, test_df):
    mlb   = MultiLabelBinarizer(classes=list(range(NUM_LABELS)))
    mlb.fit(train_df["labels"])
    y_tr  = mlb.transform(train_df["labels"]).astype(np.float32)
    y_va  = mlb.transform(val_df["labels"]).astype(np.float32)
    y_te  = mlb.transform(test_df["labels"]).astype(np.float32)
    return y_tr, y_va, y_te, mlb


# ════════════════════════════════════════════════════════════════
#  DATASET — with optional word-dropout augmentation
# ════════════════════════════════════════════════════════════════
class EmotionDataset(Dataset):
    def __init__(self, texts, labels, tokenizer, max_len, augment=False):
        self.texts     = list(texts)
        self.labels    = torch.tensor(labels, dtype=torch.float32)
        self.tokenizer = tokenizer
        self.max_len   = max_len
        self.augment   = augment

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

    def _word_dropout(self, text, p=0.10):
        words = text.split()
        if len(words) <= 3:
            return text
        keep = [w for w in words if np.random.rand() > p]
        return " ".join(keep) if keep else text

    def __getitem__(self, idx):
        text = self.texts[idx]
        if self.augment and np.random.rand() < 0.3:
            text = self._word_dropout(text)
        enc = self.tokenizer(
            text, max_length=self.max_len,
            padding="max_length", truncation=True, return_tensors="pt"
        )
        return {
            "input_ids":      enc["input_ids"].squeeze(0),
            "attention_mask": enc["attention_mask"].squeeze(0),
            "labels":         self.labels[idx],
        }


# ════════════════════════════════════════════════════════════════
#  MODEL — CLS + Mean pooling, clean linear head
# ════════════════════════════════════════════════════════════════
class MeanPooling(nn.Module):
    def forward(self, last_hidden, attention_mask):
        mask   = attention_mask.unsqueeze(-1).float()
        summed = (last_hidden * mask).sum(1)
        count  = mask.sum(1).clamp(min=1e-9)
        return summed / count

class EmotionClassifier(nn.Module):
    def __init__(self, model_name, num_labels=28, dropout=0.1):
        super().__init__()
        self.encoder    = AutoModel.from_pretrained(model_name)
        hidden          = self.encoder.config.hidden_size
        self.pool       = MeanPooling()
        # CLS + mean concatenated → richer representation
        self.norm       = nn.LayerNorm(hidden * 2)
        self.drop       = nn.Dropout(dropout)
        self.classifier = nn.Linear(hidden * 2, num_labels)
        nn.init.normal_(self.classifier.weight, std=0.02)
        nn.init.zeros_(self.classifier.bias)

    def forward(self, input_ids, attention_mask):
        out  = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
        cls  = out.last_hidden_state[:, 0, :]
        mean = self.pool(out.last_hidden_state, attention_mask)
        feat = torch.cat([cls, mean], dim=-1)
        feat = self.norm(feat)
        return self.classifier(self.drop(feat))


# ════════════════════════════════════════════════════════════════
#  LOSS — Asymmetric Loss (ASL)
#  Paper: "Asymmetric Loss For Multi-Label Classification" ICCV 2021
#  Down-weights easy negatives heavily (gamma_neg=4) while treating
#  positives gently (gamma_pos=0) — ideal for 185x imbalance.
# ════════════════════════════════════════════════════════════════
class AsymmetricLoss(nn.Module):
    def __init__(self, gamma_neg=4, gamma_pos=0, clip=0.05, eps=1e-8):
        super().__init__()
        self.gamma_neg = gamma_neg
        self.gamma_pos = gamma_pos
        self.clip      = clip
        self.eps       = eps

    def forward(self, logits, targets):
        probs     = torch.sigmoid(logits)
        probs_neg = (probs - self.clip).clamp(min=0)   # margin shift

        loss_pos = targets       * torch.log(probs.clamp(min=self.eps))
        loss_neg = (1 - targets) * torch.log((1 - probs_neg).clamp(min=self.eps))

        p_t       = probs * targets + probs_neg * (1 - targets)
        focal_pos = (1 - p_t) ** self.gamma_pos
        focal_neg = p_t       ** self.gamma_neg

        return -(focal_pos * loss_pos + focal_neg * loss_neg).mean()


# ════════════════════════════════════════════════════════════════
#  OPTIMIZER — Layer-wise LR Decay (LLRD)
#  Layer 11 (top) → encoder_lr; Layer 0 (bottom) → encoder_lr * 0.9^11
# ════════════════════════════════════════════════════════════════
def build_optimizer(model, encoder_lr, head_lr, llrd, weight_decay):
    no_decay     = ["bias", "LayerNorm.weight", "layer_norm.weight"]
    named_params = list(model.encoder.named_parameters())
    param_groups = []

    # Embeddings — lowest LR
    emb_lr = encoder_lr * (llrd ** 13)
    for wd in (False, True):
        p = [v for n, v in named_params
             if "embeddings" in n and (any(nd in n for nd in no_decay) == wd)]
        if p:
            param_groups.append({
                "params": p,
                "lr": emb_lr,
                "weight_decay": 0.0 if wd else weight_decay
            })

    # Transformer layers
    num_layers = model.encoder.config.num_hidden_layers
    for layer_idx in range(num_layers):
        layer_lr = encoder_lr * (llrd ** (num_layers - layer_idx - 1))
        for wd in (False, True):
            p = [v for n, v in named_params
                 if (f"layer.{layer_idx}." in n or f"layers.{layer_idx}." in n)
                 and (any(nd in n for nd in no_decay) == wd)]
            if p:
                param_groups.append({
                    "params": p,
                    "lr": layer_lr,
                    "weight_decay": 0.0 if wd else weight_decay
                })

    # Pooler
    pooler = [v for n, v in named_params if "pooler" in n]
    if pooler:
        param_groups.append({"params": pooler, "lr": encoder_lr, "weight_decay": weight_decay})

    # Head (norm + dropout + classifier + mean pooler)
    head_params = (list(model.norm.parameters()) +
                   list(model.drop.parameters()) +
                   list(model.classifier.parameters()) +
                   list(model.pool.parameters()))
    param_groups.append({"params": head_params, "lr": head_lr, "weight_decay": weight_decay})

    return AdamW(param_groups, eps=1e-8)


# ════════════════════════════════════════════════════════════════
#  THRESHOLD TUNING — coarse-to-fine per class
# ════════════════════════════════════════════════════════════════
def tune_thresholds(y_true, y_probs):
    best = np.full(NUM_LABELS, 0.5)
    for i in range(NUM_LABELS):
        col = y_probs[:, i]
        best_f1, best_t = 0.0, 0.5
        # Coarse pass
        for t in np.arange(0.05, 0.95, 0.05):
            f1 = f1_score(y_true[:, i], (col >= t).astype(int), zero_division=0)
            if f1 > best_f1:
                best_f1, best_t = f1, t
        # Fine pass
        for t in np.arange(max(0.01, best_t - 0.06), min(0.99, best_t + 0.06), 0.01):
            f1 = f1_score(y_true[:, i], (col >= t).astype(int), zero_division=0)
            if f1 > best_f1:
                best_f1, best_t = f1, t
        best[i] = best_t
    return best


# ════════════════════════════════════════════════════════════════
#  EVALUATION
# ════════════════════════════════════════════════════════════════
def evaluate(y_true, y_probs, thresholds=None, name=""):
    if thresholds is None:
        thresholds = np.full(NUM_LABELS, 0.5)
    y_pred      = (y_probs >= thresholds).astype(int)
    per_class   = f1_score(y_true, y_pred, average=None, zero_division=0)
    bottom5     = [(EMOTIONS[i], round(per_class[i], 3)) for i in np.argsort(per_class)[:5]]
    top5        = [(EMOTIONS[i], round(per_class[i], 3)) for i in np.argsort(per_class)[-5:]]
    results = {
        "model":        name,
        "macro_f1":     round(f1_score(y_true, y_pred, average="macro",    zero_division=0), 4),
        "micro_f1":     round(f1_score(y_true, y_pred, average="micro",    zero_division=0), 4),
        "weighted_f1":  round(f1_score(y_true, y_pred, average="weighted", zero_division=0), 4),
        "precision":    round(precision_score(y_true, y_pred, average="macro", zero_division=0), 4),
        "recall":       round(recall_score(y_true, y_pred, average="macro",    zero_division=0), 4),
        "hamming_loss": round(hamming_loss(y_true, y_pred), 4),
        "per_class_f1": {e: round(float(per_class[i]), 4) for i, e in enumerate(EMOTIONS)},
    }
    print(f"\n{'='*55}\nResults — {name}\n{'='*55}")
    for k, v in results.items():
        if k not in ("model", "per_class_f1"):
            print(f"  {k:15s}: {v}")
    print(f"\n  Bottom-5: {bottom5}")
    print(f"  Top-5   : {top5}")
    return results


# ════════════════════════════════════════════════════════════════
#  MAIN
# ════════════════════════════════════════════════════════════════
def train():
    os.makedirs(CFG["output_dir"], exist_ok=True)

    train_df, val_df, test_df   = load_data()
    y_train, y_val, y_test, mlb = binarize(train_df, val_df, test_df)

    # Load model
    try:
        tokenizer = AutoTokenizer.from_pretrained(CFG["model_name"])
        model     = EmotionClassifier(CFG["model_name"], NUM_LABELS).to(DEVICE).float()
        used      = CFG["model_name"]
    except Exception as e:
        print(f"Primary model failed ({e}), using fallback.")
        tokenizer = AutoTokenizer.from_pretrained(CFG["fallback"])
        model     = EmotionClassifier(CFG["fallback"], NUM_LABELS).to(DEVICE).float()
        used      = CFG["fallback"]

    print(f"\nModel      : {used}")
    print(f"Parameters : {sum(p.numel() for p in model.parameters()):,}")

    # WeightedRandomSampler — oversample rare-label examples
    label_counts  = y_train.sum(axis=0)
    sample_weight = np.ones(len(y_train))
    for i in range(len(y_train)):
        active = np.where(y_train[i] == 1)[0]
        if len(active) > 0:
            rarest = label_counts[active].min()
            sample_weight[i] = label_counts.max() / max(rarest, 1)
    sampler = WeightedRandomSampler(
        torch.tensor(sample_weight, dtype=torch.float),
        num_samples=len(y_train), replacement=True
    )

    train_ds = EmotionDataset(train_df["text"], y_train, tokenizer, CFG["max_length"], augment=True)
    val_ds   = EmotionDataset(val_df["text"],   y_val,   tokenizer, CFG["max_length"], augment=False)
    test_ds  = EmotionDataset(test_df["text"],  y_test,  tokenizer, CFG["max_length"], augment=False)

    train_loader = DataLoader(train_ds, batch_size=CFG["batch_size"],
                              sampler=sampler, num_workers=2, pin_memory=True)
    val_loader   = DataLoader(val_ds,   batch_size=CFG["batch_size"] * 2,
                              shuffle=False, num_workers=2, pin_memory=True)
    test_loader  = DataLoader(test_ds,  batch_size=CFG["batch_size"] * 2,
                              shuffle=False, num_workers=2, pin_memory=True)

    criterion = AsymmetricLoss(
        gamma_neg=CFG["asl_gamma_neg"],
        gamma_pos=CFG["asl_gamma_pos"],
        clip=CFG["asl_clip"]
    )

    optimizer = build_optimizer(
        model, CFG["encoder_lr"], CFG["head_lr"],
        CFG["llrd_factor"], CFG["weight_decay"]
    )

    steps_per_epoch = math.ceil(len(train_loader) / CFG["grad_accum_steps"])
    total_steps     = steps_per_epoch * CFG["epochs"]
    warmup_steps    = int(total_steps * CFG["warmup_ratio"])
    scheduler = get_cosine_schedule_with_warmup(optimizer, warmup_steps, total_steps)
    scaler    = GradScaler(device="cuda", enabled=(CFG["fp16"] and DEVICE.type == "cuda"))

    print(f"\nSteps/epoch : {steps_per_epoch}")
    print(f"Total steps : {total_steps}  |  Warmup: {warmup_steps}")
    print(f"Encoder LR  : {CFG['encoder_lr']}  |  Head LR: {CFG['head_lr']}")
    print(f"Loss        : AsymmetricLoss gamma_neg={CFG['asl_gamma_neg']}")
    print(f"Sampler     : WeightedRandomSampler (rare-class boost)")

    best_f1, patience_count  = 0.0, 0
    best_val_probs = best_val_labels = None
    history = []

    print(f"\n{'─'*55}")
    print(f"Training — {CFG['epochs']} epochs max  |  patience={CFG['patience']}")
    print(f"{'─'*55}")

    for epoch in range(CFG["epochs"]):

        # ── Train ─────────────────────────────────────────────
        model.train()
        train_loss = 0.0
        optimizer.zero_grad()
        fp16_on = CFG["fp16"] and DEVICE.type == "cuda"

        for step, batch in enumerate(train_loader):
            ids  = batch["input_ids"].to(DEVICE)
            mask = batch["attention_mask"].to(DEVICE)
            lbls = batch["labels"].to(DEVICE)

            with autocast(device_type=DEVICE.type, enabled=fp16_on):
                loss = criterion(model(ids, mask), lbls) / CFG["grad_accum_steps"]

            scaler.scale(loss).backward()

            if (step + 1) % CFG["grad_accum_steps"] == 0 or (step + 1) == len(train_loader):
                scaler.unscale_(optimizer)
                nn.utils.clip_grad_norm_(model.parameters(), CFG["max_grad_norm"])
                scaler.step(optimizer)
                scaler.update()
                scheduler.step()
                optimizer.zero_grad()

            train_loss += loss.item() * CFG["grad_accum_steps"]

        # ── Validate ──────────────────────────────────────────
        model.eval()
        vp_list, vl_list, val_loss = [], [], 0.0
        with torch.no_grad():
            for batch in val_loader:
                ids  = batch["input_ids"].to(DEVICE)
                mask = batch["attention_mask"].to(DEVICE)
                lbls = batch["labels"].to(DEVICE)
                with autocast(device_type=DEVICE.type, enabled=fp16_on):
                    logits    = model(ids, mask)
                    val_loss += criterion(logits, lbls).item()
                vp_list.append(torch.sigmoid(logits).cpu().numpy())
                vl_list.append(lbls.cpu().numpy())

        vp = np.vstack(vp_list)
        vl = np.vstack(vl_list)
        macro_f1  = f1_score(vl, (vp >= CFG["eval_threshold"]).astype(int),
                             average="macro", zero_division=0)
        head_lr   = optimizer.param_groups[-1]["lr"]

        print(f"  Epoch {epoch+1:02d}/{CFG['epochs']} | "
              f"TLoss {train_loss/len(train_loader):.4f} | "
              f"VLoss {val_loss/len(val_loader):.4f} | "
              f"MacroF1 {macro_f1:.4f} | LR {head_lr:.2e}")

        history.append({
            "epoch": epoch + 1,
            "train_loss": round(train_loss / len(train_loader), 4),
            "val_loss":   round(val_loss   / len(val_loader),   4),
            "macro_f1":   round(macro_f1, 4),
        })

        if macro_f1 > best_f1:
            best_f1, patience_count = macro_f1, 0
            model.encoder.save_pretrained(CFG["output_dir"])
            tokenizer.save_pretrained(CFG["output_dir"])
            torch.save({
                "norm":       model.norm.state_dict(),
                "drop":       model.drop.state_dict(),
                "classifier": model.classifier.state_dict(),
            }, os.path.join(CFG["output_dir"], "head.pt"))
            best_val_probs  = vp.copy()
            best_val_labels = vl.copy()
            print(f"  ✅ New best  F1={best_f1:.4f}  [saved]")
        else:
            patience_count += 1
            print(f"  No improvement ({patience_count}/{CFG['patience']})")
            if patience_count >= CFG["patience"]:
                print(f"  Early stopping at epoch {epoch+1}")
                break

    # ── Threshold tuning ──────────────────────────────────────
    print("\nTuning per-class thresholds (coarse→fine) ...")
    thresholds = tune_thresholds(best_val_labels, best_val_probs)
    np.save(CFG["thresh_path"], thresholds)
    print(f"  Mean  : {thresholds.mean():.3f}")
    print(f"  Range : [{thresholds.min():.2f}, {thresholds.max():.2f}]")
    rare = ["grief", "pride", "nervousness", "relief", "embarrassment"]
    for e in rare:
        i = EMOTIONS.index(e)
        print(f"  [{e}] threshold = {thresholds[i]:.2f}")

    # ── Test evaluation ────────────────────────────────────────
    print("\nLoading best checkpoint ...")
    model.encoder = AutoModel.from_pretrained(CFG["output_dir"]).to(DEVICE)
    ckpt = torch.load(os.path.join(CFG["output_dir"], "head.pt"), map_location=DEVICE)
    model.norm.load_state_dict(ckpt["norm"])
    model.drop.load_state_dict(ckpt["drop"])
    model.classifier.load_state_dict(ckpt["classifier"])
    model.eval()

    test_probs = []
    with torch.no_grad():
        for batch in test_loader:
            ids  = batch["input_ids"].to(DEVICE)
            mask = batch["attention_mask"].to(DEVICE)
            test_probs.append(torch.sigmoid(model(ids, mask)).cpu().numpy())

    test_results = evaluate(
        y_test, np.vstack(test_probs), thresholds,
        name="RoBERTa-GoEmotions-ASL-LLRD-v2"
    )

    a1_best = 0.5245
    output  = {
        "model_used":       used,
        "config":           CFG,
        "history":          history,
        "test_results":     test_results,
        "assignment1_best": {"model": "roberta-base", "macro_f1": a1_best},
        "improvement":      round(test_results["macro_f1"] - a1_best, 4),
    }
    with open(CFG["results_path"], "w") as f:
        json.dump(output, f, indent=2)

    print(f"\n{'='*55}")
    print(f"  Macro F1 (test) : {test_results['macro_f1']}")
    print(f"  Micro F1 (test) : {test_results['micro_f1']}")
    print(f"  A1 best         : {a1_best}")
    print(f"  Improvement     : +{output['improvement']:.4f}")
    print(f"  Saved to        : {CFG['output_dir']}/")
    print(f"{'='*55}")


if __name__ == "__main__":
    train()