# lmr/glue_benchmark.py
"""
Standalone GLUE benchmark runner for your project.

Implements BERT paper fine-tuning protocol + evaluation error bars.

BERT protocol:
- For each task, select the best learning rate on the Dev set among:
    {5e-5, 4e-5, 3e-5, 2e-5}
- For small/unstable datasets (CoLA, MRPC, RTE, STS-B):
    run 5 random restarts per LR (different seed => data shuffle + classifier init),
    i.e. 4 LRs * 5 restarts = 20 runs, pick best on Dev.
- For large datasets:
    run LR sweep only (4 runs), 1 seed each, pick best on Dev.

Evaluation error bars (for the FINAL best model):
- Split the eval split into 5 folds (0..4) by contiguous chunks.
- Compute metric on 5 "leave-one-fold-out" subsets:
    0123, 1234, 0124, 0234, 0134
  (equivalently: drop fold 4,0,3,1,2)
- Report mean/std/stderr across the 5 subset scores.

Saving:
- For EACH run:
    out_dir/checkpoints/<task>/lr_<LR>/restart_<K>/
        best_finetuned.pt
        finetuned.pt
        run_meta.json
- For task summary of all runs:
    out_dir/checkpoints/<task>/all_runs.csv
- Best overall for task:
    out_dir/checkpoints/<task>/best_overall/
        best_finetuned.pt
        best_meta.json
- Evaluation outputs:
    out_dir/<task>/
        <task>_<split>_results.json
        <task>_<split>_preds.csv
        <task>_<split>_errorbar.json   (mean/std/stderr + per-subset scores)
- Final benchmark summary:
    out_dir/glue_summary.csv
"""

import os
import json
import re
import math
import random
import shutil
from pathlib import Path
from typing import Optional, List, Tuple, Dict, Any

import torch
import numpy as np
import pandas as pd
from datasets import load_dataset
from torch.utils.data import DataLoader, TensorDataset
from tqdm import tqdm

import evaluate

from lmr.checkpointing import Checkpointing
from lmr.ddp import unwrap_model

# ---------------------------------------------------------------------
# GLUE config
# ---------------------------------------------------------------------
GLUE_TASKS = {
    "cola": {"type": "classification", "num_labels": 2, "hf_name": "cola"},
    "sst2": {"type": "classification", "num_labels": 2, "hf_name": "sst2"},
    "mrpc": {"type": "classification", "num_labels": 2, "hf_name": "mrpc"},
    "stsb": {"type": "regression",     "num_labels": 1, "hf_name": "stsb"},
    "qqp": {"type": "classification", "num_labels": 2, "hf_name": "qqp"},
    "mnli": {"type": "classification", "num_labels": 3, "hf_name": "mnli"},
    "qnli": {"type": "classification", "num_labels": 2, "hf_name": "qnli"},
    "rte":  {"type": "classification", "num_labels": 2, "hf_name": "rte"},
    "wnli": {"type": "classification", "num_labels": 2, "hf_name": "wnli"},
}

SMALL_TASKS_RANDOM_RESTARTS = {"cola", "mrpc", "rte", "stsb"}
BERT_LR_CANDIDATES = [2e-5]

# Preferred metric key per task (for selecting best run / best epoch)
PREFERRED_METRIC_KEY = {
    "cola": "matthews_correlation",
    "sst2": "accuracy",
    "mrpc": "accuracy",  # also has f1; keep accuracy unless you want change
    "stsb": "pearson",
    "qqp": "accuracy",
    "mnli": "accuracy",
    "qnli": "accuracy",
    "rte": "accuracy",
    "wnli": "accuracy",
}

# ---------------------------------------------------------------------
# Repro helpers
# ---------------------------------------------------------------------
def _set_all_seeds(seed: int):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

def _json_dump(obj: Any, path: Path):
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        json.dump(obj, f, indent=2, ensure_ascii=False)

def _safe_float(x):
    try:
        if isinstance(x, (np.generic,)):
            return float(x.item())
        return float(x)
    except Exception:
        return None

def _metric_to_scalar(task: str, metric_res: Dict[str, Any], fallback_val_loss: Optional[float] = None) -> float:
    if isinstance(metric_res, dict) and metric_res:
        pref = PREFERRED_METRIC_KEY.get(task)
        if pref is not None and pref in metric_res:
            v = _safe_float(metric_res.get(pref))
            if v is not None and not math.isnan(v):
                return float(v)
        for _, v in metric_res.items():
            fv = _safe_float(v)
            if fv is not None and not math.isnan(fv):
                return float(fv)
    if fallback_val_loss is not None:
        try:
            return -float(fallback_val_loss)
        except Exception:
            pass
    return -1e9

# ---------------------------------------------------------------------
# Task-aware example field extraction
# ---------------------------------------------------------------------
def _get_text_pair_from_example(task: str, ex: dict):
    task_field_map = {
        "cola": ("sentence", None),
        "sst2": ("sentence", None),
        "mrpc": ("sentence1", "sentence2"),
        "stsb": ("sentence1", "sentence2"),
        "qqp": ("question1", "question2"),
        "mnli": ("premise", "hypothesis"),
        "qnli": ("question", "sentence"),
        "rte": ("sentence1", "sentence2"),
        "wnli": ("sentence1", "sentence2"),
    }
    f1, f2 = task_field_map.get(task, (None, None))

    def _try_keys(keys):
        for k in keys:
            if k in ex and ex.get(k) is not None:
                return ex.get(k)
        return None

    s1_candidates = []
    s2_candidates = []

    if f1:
        s1_candidates.append(f1)
    s1_candidates += ["sentence1", "premise", "question", "sentence", "text", "question1"]

    if f2:
        s2_candidates.append(f2)
    s2_candidates += ["sentence2", "hypothesis", "question2", "question1", "text2"]

    s1 = _try_keys(s1_candidates)
    s2 = _try_keys(s2_candidates)

    if s1 is None:
        s1 = ex.get("sentence") or ex.get("premise") or ex.get("question") or ex.get("text")
    if s2 is None:
        s2 = ex.get("sentence2") or ex.get("hypothesis") or ex.get("question2")

    s1 = "" if s1 is None else (s1 if isinstance(s1, str) else str(s1))
    s2 = None if s2 is None else (s2 if isinstance(s2, str) else str(s2))
    return s1, s2

# ---------------------------------------------------------------------
# Tokenization helpers
# ---------------------------------------------------------------------
def _pad_and_tensorize(input_ids_list, attention_mask_list, pad_token_id: int):
    max_len = max(len(x) for x in input_ids_list) if input_ids_list else 0
    ids_padded = [x + [pad_token_id] * (max_len - len(x)) for x in input_ids_list]
    mask_padded = [m + [0] * (max_len - len(m)) for m in attention_mask_list]
    input_ids = torch.tensor(ids_padded, dtype=torch.long)
    attention_mask = torch.tensor(mask_padded, dtype=torch.long)
    return input_ids, attention_mask

def _batch_tokenize(tokenizer, texts: List[Tuple[Optional[str], Optional[str]]], max_length: int = 128):
    sanitized = []
    for a, b in texts:
        a_s = "" if a is None else (a if isinstance(a, str) else str(a))
        b_s = None if b is None else (b if isinstance(b, str) else str(b))
        sanitized.append((a_s, b_s))

    try:
        flat = [(a if b is None else (a, b)) for a, b in sanitized]
        enc = tokenizer(flat, truncation=True, padding=False, max_length=max_length)
        if isinstance(enc.get("input_ids", None), torch.Tensor):
            enc["input_ids"] = enc["input_ids"].tolist()
        if isinstance(enc.get("attention_mask", None), torch.Tensor):
            enc["attention_mask"] = enc["attention_mask"].tolist()
        return enc
    except Exception:
        pass

    for method_name in ("batch_encode", "encode_batch", "batch_encode_plus", "encode_batch_pair", "encode_batch_items"):
        fn = getattr(tokenizer, method_name, None)
        if fn is None:
            continue
        try:
            try:
                enc = fn(sanitized, max_length=max_length, truncation=True, padding=False)
            except TypeError:
                enc = fn(sanitized)
            if isinstance(enc.get("input_ids", None), torch.Tensor):
                enc["input_ids"] = enc["input_ids"].tolist()
            if isinstance(enc.get("attention_mask", None), torch.Tensor):
                enc["attention_mask"] = enc["attention_mask"].tolist()
            return enc
        except Exception:
            continue

    input_ids_list = []
    attention_mask_list = []
    for a, b in sanitized:
        try:
            if b is None:
                try:
                    single = tokenizer.encode(a)
                except TypeError:
                    single = tokenizer.encode([a])
            else:
                single = None
                try:
                    single = tokenizer.encode((a, b))
                except Exception:
                    try:
                        single = tokenizer.encode(a, b)
                    except Exception:
                        single = tokenizer(a if b is None else (a, b))

            if isinstance(single, dict):
                ids = single.get("input_ids") or single.get("ids") or []
                mask = single.get("attention_mask") or single.get("mask") or [1] * len(ids)
            elif isinstance(single, torch.Tensor):
                ids = single.tolist()
                mask = [1] * len(ids)
            elif isinstance(single, list):
                ids = single
                mask = [1] * len(ids)
            else:
                tmp = tokenizer(a if b is None else (a, b))
                if isinstance(tmp, dict):
                    ids = tmp.get("input_ids") or tmp.get("ids") or []
                    mask = tmp.get("attention_mask") or tmp.get("mask") or [1] * len(ids)
                elif torch.is_tensor(tmp):
                    ids = tmp.tolist()
                    mask = [1] * len(ids)
                else:
                    ids = list(tmp)
                    mask = [1] * len(ids)

            if len(ids) > max_length:
                ids = ids[:max_length]
                mask = mask[:max_length]

            input_ids_list.append(ids)
            attention_mask_list.append(mask)
        except Exception as e:
            snippet = (a[:80] + "...") if a else "<empty>"
            raise RuntimeError(f"Tokenizer fallback encode failed for example '{snippet}': {e}")

    return {"input_ids": input_ids_list, "attention_mask": attention_mask_list}

# ---------------------------------------------------------------------
# Postprocess preds
# ---------------------------------------------------------------------
def _postprocess_predictions(task: str, logits_np: np.ndarray, cfg_task: dict):
    ttype = cfg_task["type"]

    if logits_np is None or logits_np.size == 0:
        return np.array([])

    if logits_np.ndim == 1:
        if ttype == "classification":
            return (logits_np > 0.5).astype(int)
        return logits_np.astype(float)

    if logits_np.ndim == 2 and logits_np.shape[1] == 1:
        col = logits_np[:, 0]
        if ttype == "classification":
            return (col > 0.5).astype(int)
        return col.astype(float)

    if logits_np.ndim == 2:
        if ttype == "classification":
            return np.argmax(logits_np, axis=-1).astype(int)
        preds = logits_np[:, 0].astype(float) if logits_np.shape[1] == 1 else logits_np.mean(axis=1).astype(float)
        if task == "stsb":
            preds = np.clip(preds, 0.0, 5.0)
        return preds

    return logits_np.ravel()

# ---------------------------------------------------------------------
# Model wrapping helper (robust)
# ---------------------------------------------------------------------
def make_wrapped_model_if_needed(model, hidden_size: Optional[int], num_labels: int, force_num_labels: Optional[int] = None):
    import torch.nn as nn

    base_model = model

    def _detect_head_dim(m):
        try:
            if hasattr(m, "classifier") and isinstance(getattr(m, "classifier"), nn.Linear):
                return getattr(m, "classifier").out_features
            if hasattr(m, "lm_head") and isinstance(getattr(m, "lm_head"), nn.Linear):
                return getattr(m, "lm_head").out_features
            if hasattr(m, "get_output_embeddings"):
                out_emb = m.get_output_embeddings()
                if out_emb is not None:
                    if isinstance(out_emb, nn.Embedding):
                        return out_emb.embedding_dim if hasattr(out_emb, "embedding_dim") else out_emb.num_embeddings
                    if isinstance(out_emb, nn.Linear):
                        return out_emb.out_features
        except Exception:
            pass
        return None

    if force_num_labels is None:
        head_dim = _detect_head_dim(base_model)
        if head_dim is not None and head_dim == num_labels:
            return base_model, False

    inferred_hidden = hidden_size
    if inferred_hidden is None:
        try:
            cand = getattr(base_model, "config", None)
            if cand is not None and hasattr(cand, "hidden_size"):
                inferred_hidden = int(cand.hidden_size)
        except Exception:
            inferred_hidden = None

    if inferred_hidden is None:
        try:
            un = unwrap_model(base_model)
            sd = un.state_dict()
            for k, v in sd.items():
                if re.search(r"embed|embedding|word_embeddings|token_embedding|embed_tokens", k, re.I):
                    if hasattr(v, "shape") and len(v.shape) == 2:
                        inferred_hidden = int(v.shape[1])
                        break
                if re.search(r"q_proj|k_proj|v_proj|o_proj|dense|fc|linear|proj", k, re.I):
                    if hasattr(v, "shape") and len(v.shape) == 2:
                        cand = max(v.shape)
                        if 1 < cand < 1_000_000:
                            inferred_hidden = int(cand)
                            break
        except Exception:
            inferred_hidden = None

    if inferred_hidden is None:
        raise RuntimeError(
            "Cannot infer hidden_size for wrapped classifier head. "
            "Please set `model.config.hidden_size` or pass `hidden_size` explicitly."
        )

    class _WrappedModel(nn.Module):
        def __init__(self, base, hidden_size, num_labels):
            super().__init__()
            self.base = base
            self.classifier = nn.Linear(hidden_size, num_labels)
            self.logits_projector = None

        def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
            try:
                out = self.base(input_ids=input_ids, attention_mask=attention_mask, **kwargs)
            except TypeError:
                out = self.base(input_ids)

            last_hidden = getattr(out, "last_hidden_state", None)
            if last_hidden is not None:
                pooled = last_hidden[:, 0, :]
                logits = self.classifier(pooled)
                return type("Out", (), {"logits": logits, "loss": None})

            if isinstance(out, (tuple, list)) and len(out) > 0:
                cand = out[0]
                if torch.is_tensor(cand):
                    if cand.ndim == 3:
                        pooled = cand[:, 0, :]
                        logits = self.classifier(pooled)
                        return type("Out", (), {"logits": logits, "loss": None})
                    if cand.ndim == 2 and cand.shape[1] == num_labels:
                        return type("Out", (), {"logits": cand, "loss": None})

            logits = getattr(out, "logits", None)
            if logits is not None:
                if logits.ndim == 2 and logits.shape[1] == num_labels:
                    return type("Out", (), {"logits": logits, "loss": getattr(out, "loss", None)})
                exist_dim = logits.shape[1]
                if self.logits_projector is None or self.logits_projector.weight.shape[1] != exist_dim:
                    self.logits_projector = nn.Linear(exist_dim, num_labels).to(logits.device)
                projected = self.logits_projector(logits)
                return type("Out", (), {"logits": projected, "loss": getattr(out, "loss", None)})

            hidden_states = getattr(out, "hidden_states", None)
            if hidden_states is not None:
                last_hidden = hidden_states[-1] if isinstance(hidden_states, (list, tuple)) else hidden_states
                if torch.is_tensor(last_hidden) and last_hidden.ndim == 3:
                    pooled = last_hidden[:, 0, :]
                    logits = self.classifier(pooled)
                    return type("Out", (), {"logits": logits, "loss": None})

            raise RuntimeError("Wrapped base model did not return recognizable hidden states or logits")

    return _WrappedModel(base_model, inferred_hidden, num_labels), True

# ---------------------------------------------------------------------
# Tokenize HF split to tensors
# ---------------------------------------------------------------------
def _tokenize_hf_split_to_tensors(task: str, tokenizer, raw_split, cfg_task, max_length=128, batch_tokenize_size=512):
    texts = []
    labels = []
    empty_s1 = 0
    empty_s2 = 0

    for ex in raw_split:
        s1, s2 = _get_text_pair_from_example(task, ex)
        texts.append((s1, s2))
        labels.append(ex.get("label") if "label" in ex else -100)
        if not s1 or (isinstance(s1, str) and s1.strip() == ""):
            empty_s1 += 1
        if s2 is not None and (not s2 or (isinstance(s2, str) and s2.strip() == "")):
            empty_s2 += 1

    total = len(texts)
    print(
        f"[tokenize] task={task} samples={total} empty_s1={empty_s1} empty_s2={empty_s2} "
        f"({(empty_s1/total if total>0 else 0):.2%}, {(empty_s2/total if total>0 else 0):.2%})"
    )

    input_ids_all = []
    attention_all = []
    for i in range(0, len(texts), batch_tokenize_size):
        enc = _batch_tokenize(tokenizer, texts[i:i+batch_tokenize_size], max_length=max_length)
        ids = enc.get("input_ids")
        masks = enc.get("attention_mask") or enc.get("mask") or enc.get("masks")
        if isinstance(ids, torch.Tensor):
            ids = ids.tolist()
        if isinstance(masks, torch.Tensor):
            masks = masks.tolist()
        input_ids_all.extend(ids)
        attention_all.extend(masks)

    pad_id = getattr(tokenizer, "pad_token_id", None)
    if pad_id is None:
        try:
            pad_id = tokenizer.token_to_id("[PAD]")
        except Exception:
            pad_id = 0

    input_ids_t, attention_mask_t = _pad_and_tensorize(input_ids_all, attention_all, pad_id)
    labels_t = torch.tensor(labels, dtype=torch.long if cfg_task["type"] == "classification" else torch.float)
    return input_ids_t, attention_mask_t, labels_t

# ---------------------------------------------------------------------
# Training
# ---------------------------------------------------------------------
def train_full_finetune(
    task: str,
    tokenizer,
    model,
    raw_train,
    raw_val,
    device: str = "cuda",
    epochs: int = 3,
    batch_size: int = 32,
    lr: float = 2e-5,
    weight_decay: float = 0.01,
    warmup_steps: int = 100,
    max_length: int = 128,
    grad_accum_steps: int = 1,
    out_checkpoint_dir: Optional[str] = None,
    seed: Optional[int] = None,
):
    cfg_task = GLUE_TASKS[task]
    device_t = torch.device(device if torch.cuda.is_available() else "cpu")

    if seed is not None:
        _set_all_seeds(int(seed))

    hidden_size = None
    if hasattr(model, "config") and hasattr(model.config, "hidden_size"):
        try:
            hidden_size = int(model.config.hidden_size)
        except Exception:
            hidden_size = None

    model, wrapped_flag = make_wrapped_model_if_needed(
        model, hidden_size, cfg_task["num_labels"], force_num_labels=cfg_task["num_labels"]
    )
    model.to(device_t)

    train_ids, train_mask, train_labels = _tokenize_hf_split_to_tensors(task, tokenizer, raw_train, cfg_task, max_length=max_length)
    val_ids, val_mask, val_labels = _tokenize_hf_split_to_tensors(task, tokenizer, raw_val, cfg_task, max_length=max_length)

    train_ds = TensorDataset(train_ids, train_mask, train_labels)
    val_ds = TensorDataset(val_ids, val_mask, val_labels)

    g = torch.Generator()
    if seed is not None:
        g.manual_seed(int(seed))

    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, pin_memory=True, generator=g)
    val_loader = DataLoader(val_ds, batch_size=max(64, batch_size), shuffle=False, pin_memory=True)

    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)
    total_steps = max(1, (len(train_loader) // max(1, grad_accum_steps)) * epochs)
    try:
        from transformers import get_cosine_schedule_with_warmup
        scheduler = get_cosine_schedule_with_warmup(optimizer, num_warmup_steps=warmup_steps, num_training_steps=total_steps)
    except Exception:
        scheduler = None

    loss_fn = torch.nn.CrossEntropyLoss() if cfg_task["type"] == "classification" else torch.nn.MSELoss()

    best_metric_res: Dict[str, Any] = {}
    best_score: Optional[float] = None
    best_epoch = -1

    model.train()

    for epoch in range(epochs):
        for step, batch in enumerate(tqdm(train_loader, desc=f"Train {task} epoch {epoch+1} (lr={lr:g})")):
            ids_b, mask_b, labs_b = batch
            ids_b = ids_b.to(device_t)
            mask_b = mask_b.to(device_t)
            labs_b = labs_b.to(device_t)

            out = model(input_ids=ids_b, attention_mask=mask_b, labels=None)
            logits = getattr(out, "logits", None)
            if logits is None:
                if isinstance(out, (tuple, list)):
                    logits = out[0]
                else:
                    raise RuntimeError("Model did not return logits during finetune")

            if cfg_task["type"] == "classification":
                loss = loss_fn(logits, labs_b.long())
            else:
                if logits.ndim == 2 and logits.shape[1] == 1:
                    preds = logits.squeeze(1)
                elif logits.ndim == 2:
                    preds = logits.mean(dim=1)
                else:
                    preds = logits
                loss = loss_fn(preds, labs_b.float())

            loss = loss / max(1, grad_accum_steps)
            loss.backward()

            if (step + 1) % max(1, grad_accum_steps) == 0:
                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
                optimizer.step()
                if scheduler is not None:
                    scheduler.step()
                optimizer.zero_grad()

        # validation
        model.eval()
        tot_val_loss = 0.0
        all_logits = []
        all_labels = []
        with torch.no_grad():
            for ids_b, mask_b, labs_b in tqdm(val_loader, desc=f"Validate {task} epoch {epoch+1}", leave=False):
                ids_b = ids_b.to(device_t)
                mask_b = mask_b.to(device_t)
                labs_b = labs_b.to(device_t)

                out = model(input_ids=ids_b, attention_mask=mask_b, labels=None)
                logits = getattr(out, "logits", None)
                if logits is None:
                    if isinstance(out, (tuple, list)):
                        logits = out[0]
                    else:
                        raise RuntimeError("Model did not return logits during validation")

                if cfg_task["type"] == "classification":
                    l = loss_fn(logits, labs_b.long())
                else:
                    if logits.ndim == 2 and logits.shape[1] == 1:
                        preds = logits.squeeze(1)
                    elif logits.ndim == 2:
                        preds = logits.mean(dim=1)
                    else:
                        preds = logits
                    l = loss_fn(preds, labs_b.float())

                tot_val_loss += l.item() * ids_b.size(0)
                all_logits.append(logits.detach().cpu().numpy())
                all_labels.append(labs_b.detach().cpu().numpy())

        model.train()

        all_logits = np.concatenate(all_logits, axis=0) if all_logits else np.zeros((0, cfg_task["num_labels"]))
        all_labels = np.concatenate(all_labels, axis=0) if all_labels else np.zeros((0,))
        preds = _postprocess_predictions(task, all_logits, cfg_task)

        metric = evaluate.load("glue", cfg_task["hf_name"])
        try:
            metric_res = metric.compute(predictions=preds.tolist(), references=all_labels.tolist())
        except Exception:
            try:
                metric_res = metric.compute(predictions=preds, references=all_labels)
            except Exception as e:
                metric_res = {"error": str(e)}

        avg_val_loss = tot_val_loss / len(val_ds) if len(val_ds) > 0 else float("nan")
        score = _metric_to_scalar(task, metric_res, fallback_val_loss=avg_val_loss)

        print(f"[FT] {task} epoch {epoch+1} lr={lr:g} val_loss={avg_val_loss:.6f} metric={metric_res} score={score:.6f}")

        if best_score is None or float(score) > float(best_score):
            best_score = float(score)
            best_metric_res = metric_res
            best_epoch = epoch + 1

            if out_checkpoint_dir:
                outp = Path(out_checkpoint_dir)
                outp.mkdir(parents=True, exist_ok=True)
                best_fname = outp / "best_finetuned.pt"
                try:
                    sd = unwrap_model(model).state_dict()
                except Exception:
                    sd = model.state_dict()
                torch.save(sd, str(best_fname))
                print(f"[FT] Saved best checkpoint (epoch {best_epoch}) to: {best_fname}")

    if out_checkpoint_dir:
        outp = Path(out_checkpoint_dir)
        outp.mkdir(parents=True, exist_ok=True)
        fname = outp / "finetuned.pt"
        try:
            sd = unwrap_model(model).state_dict()
        except Exception:
            sd = model.state_dict()
        torch.save(sd, str(fname))
        print(f"[FT] Saved finetuned model to: {fname}")

        meta = {
            "task": task,
            "lr": lr,
            "seed": seed,
            "epochs": epochs,
            "batch_size": batch_size,
            "grad_accum_steps": grad_accum_steps,
            "warmup_steps": warmup_steps,
            "weight_decay": weight_decay,
            "max_length": max_length,
            "wrapped_flag": bool(wrapped_flag),
            "best_epoch": int(best_epoch),
            "best_score": float(best_score) if best_score is not None else None,
            "best_metrics": best_metric_res,
        }
        _json_dump(meta, outp / "run_meta.json")

    print(f"[FT] Best validation for task '{task}' (lr={lr:g}, seed={seed}): epoch={best_epoch}, score={best_score}, metrics={best_metric_res}")
    return model, best_metric_res, float(best_score) if best_score is not None else -1e9, best_epoch

# ---------------------------------------------------------------------
# Load finetuned checkpoint for eval (wrapped/unwrapped)
# ---------------------------------------------------------------------
def _load_finetuned_checkpoint_for_task(task: str, base_model, checkpoint_path: str):
    cfg = GLUE_TASKS[task]
    sd = torch.load(checkpoint_path, map_location="cpu")
    keys = list(sd.keys()) if isinstance(sd, dict) else []
    looks_wrapped = any(k.startswith("base.") for k in keys) or any(k.startswith("classifier.") for k in keys)

    hidden_size = None
    if hasattr(base_model, "config") and hasattr(base_model.config, "hidden_size"):
        try:
            hidden_size = int(base_model.config.hidden_size)
        except Exception:
            hidden_size = None

    if looks_wrapped:
        wrapped_model, _ = make_wrapped_model_if_needed(
            base_model, hidden_size, cfg["num_labels"], force_num_labels=cfg["num_labels"]
        )
        try:
            unwrap_model(wrapped_model).load_state_dict(sd, strict=False)
        except Exception:
            wrapped_model.load_state_dict(sd, strict=False)
        return wrapped_model

    try:
        unwrap_model(base_model).load_state_dict(sd, strict=False)
    except Exception:
        base_model.load_state_dict(sd, strict=False)
    return base_model

# ---------------------------------------------------------------------
# Error bar evaluation (5 folds -> 5 leave-one-fold-out subsets)
# ---------------------------------------------------------------------
def _fivefold_indices(n: int):
    """
    Split indices [0..n-1] into 5 contiguous folds as evenly as possible.
    Returns list of 5 lists of indices.
    """
    idx = np.arange(n)
    folds = np.array_split(idx, 5)  # handles remainder nicely
    return [f.tolist() for f in folds]

def _errorbar_subsets_from_folds(folds: List[List[int]]):
    """
    Make the five subsets exactly as requested:
      0123, 1234, 0124, 0234, 0134
    Which correspond to dropping fold: 4,0,3,1,2.
    Returns list of dicts: {"name": "0123", "indices": [...]}
    """
    assert len(folds) == 5
    combos = [
        ("0123", [0,1,2,3]),
        ("1234", [1,2,3,4]),
        ("0124", [0,1,2,4]),
        ("0234", [0,2,3,4]),
        ("0134", [0,1,3,4]),
    ]
    subsets = []
    for name, keep in combos:
        inds = []
        for k in keep:
            inds.extend(folds[k])
        subsets.append({"name": name, "indices": inds})
    return subsets

def _compute_metric_for_indices(task: str, cfg: dict, metric_obj, preds_all: np.ndarray, labels_all: np.ndarray, indices: List[int]):
    if len(indices) == 0:
        return {"error": "empty_indices"}
    p = preds_all[indices]
    y = labels_all[indices]
    if cfg["type"] == "classification":
        preds_out = p.astype(int).tolist()
        refs_out = y.astype(int).tolist()
    else:
        preds_out = p.astype(float).tolist()
        refs_out = y.astype(float).tolist()
    try:
        return metric_obj.compute(predictions=preds_out, references=refs_out)
    except Exception:
        try:
            return metric_obj.compute(predictions=np.array(preds_out), references=np.array(refs_out))
        except Exception as e:
            return {"error": str(e)}

def _compute_errorbar(task: str, cfg: dict, metric_obj, preds_all: np.ndarray, labels_all: np.ndarray):
    """
    Returns:
      {
        "preferred_key": ...,
        "subset_scores": [{"subset":"0123","score":...,"metrics":{...}}, ...],
        "mean": ...,
        "std": ...,
        "stderr": ...
      }
    """
    n = int(len(labels_all))
    folds = _fivefold_indices(n)
    subsets = _errorbar_subsets_from_folds(folds)

    pref = PREFERRED_METRIC_KEY.get(task)
    subset_scores = []
    scores = []

    for s in subsets:
        m = _compute_metric_for_indices(task, cfg, metric_obj, preds_all, labels_all, s["indices"])
        sc = _metric_to_scalar(task, m, fallback_val_loss=None)
        subset_scores.append({"subset": s["name"], "score": float(sc), "metrics": m})
        scores.append(float(sc))

    arr = np.array(scores, dtype=float)
    mean = float(np.mean(arr)) if len(arr) else float("nan")
    std = float(np.std(arr, ddof=1)) if len(arr) > 1 else 0.0
    stderr = float(std / math.sqrt(len(arr))) if len(arr) > 0 else float("nan")

    return {
        "preferred_key": pref,
        "subset_scores": subset_scores,
        "mean": mean,
        "std": std,
        "stderr": stderr,
    }

# ---------------------------------------------------------------------
# Evaluation (returns metrics + also writes preds/results)
# ---------------------------------------------------------------------
def run_glue_task(
    task: str,
    tokenizer,
    model,
    checkpointing: Optional[Checkpointing] = None,
    device: str = "cuda",
    batch_size: int = 64,
    max_length: int = 128,
    output_dir: str = "glue_output",
    compute_errorbar: bool = False,
):
    """
    Run a single GLUE task evaluation. Returns dict per split:
      results_by_split[split] = {
         "metrics": { ... },
         "errorbar": { ... } or None
      }
    """
    assert task in GLUE_TASKS, f"Unknown GLUE task: {task}"
    cfg = GLUE_TASKS[task]

    hf = load_dataset("glue", cfg["hf_name"])
    if task == "mnli":
        val_splits = ["validation_matched", "validation_mismatched"]
    else:
        val_splits = ["validation"]

    results_by_split = {}
    for split in val_splits:
        raw = hf[split]
        print(f"[GLUE] Task={task} split={split} samples={len(raw)}")

        texts = []
        labels = []
        for ex in raw:
            s1, s2 = _get_text_pair_from_example(task, ex)
            texts.append((s1, s2))
            labels.append(ex.get("label") if "label" in ex else -100)

        BATCH = 512
        input_ids_all = []
        attention_all = []
        for i in range(0, len(texts), BATCH):
            enc = _batch_tokenize(tokenizer, texts[i:i+BATCH], max_length=max_length)
            ids = enc.get("input_ids")
            masks = enc.get("attention_mask") or enc.get("mask") or enc.get("masks")
            if isinstance(ids, torch.Tensor):
                ids = ids.tolist()
            if isinstance(masks, torch.Tensor):
                masks = masks.tolist()
            input_ids_all.extend(ids)
            attention_all.extend(masks)
# cp /cwork/jf381/checkpoints/transformer_353M_update_pretrain_2gpu/. -r /work/jf381/checkpoints/transformer_353M_update_pretrain_2gpu_test
        pad_id = getattr(tokenizer, "pad_token_id", None)
        if pad_id is None:
            try:
                pad_id = tokenizer.token_to_id("[PAD]")
            except Exception:
                pad_id = 0

        input_ids, attention_mask = _pad_and_tensorize(input_ids_all, attention_all, pad_id)
        labels_t = torch.tensor(labels, dtype=torch.long if cfg["type"] == "classification" else torch.float)

        ds = TensorDataset(input_ids, attention_mask, labels_t)
        loader = DataLoader(ds, batch_size=batch_size, shuffle=False, pin_memory=True)

        if checkpointing is not None:
            try:
                checkpointing.load_model_states("recent")
            except Exception:
                pass

        device_t = torch.device(device if torch.cuda.is_available() else "cpu")
        model.to(device_t)
        model.eval()

        hidden_size = None
        if hasattr(model, "config") and hasattr(model.config, "hidden_size"):
            try:
                hidden_size = int(model.config.hidden_size)
            except Exception:
                hidden_size = None

        force = 1 if cfg["type"] == "regression" else cfg["num_labels"]
        wrapped_model, _ = make_wrapped_model_if_needed(model, hidden_size, cfg["num_labels"], force_num_labels=force)
        wrapped_model.to(device_t)
        wrapped_model.eval()

        all_logits = []
        all_labels = []
        with torch.no_grad():
            for batch in tqdm(loader, desc=f"Eval {task}:{split}"):
                ids_b, mask_b, labels_b = batch
                ids_b = ids_b.to(device_t)
                mask_b = mask_b.to(device_t)
                out = wrapped_model(input_ids=ids_b, attention_mask=mask_b, labels=None)
                logits = getattr(out, "logits", None)
                if logits is None:
                    if isinstance(out, (tuple, list)):
                        logits = out[0]
                    else:
                        raise RuntimeError("Model forward did not return logits")
                all_logits.append(logits.detach().cpu().numpy())
                all_labels.append(labels_b.detach().cpu().numpy())

        all_logits = np.concatenate(all_logits, axis=0) if all_logits else np.zeros((0, cfg["num_labels"]))
        all_labels = np.concatenate(all_labels, axis=0) if all_labels else np.zeros((0,))

        preds = _postprocess_predictions(task, all_logits, cfg)

        # metric
        metric = evaluate.load("glue", cfg["hf_name"])
        metric_res = _compute_metric_for_indices(task, cfg, metric, preds, all_labels, list(range(len(all_labels))))

        # error bar (only for FINAL best model usually)
        errorbar_res = None
        if compute_errorbar:
            errorbar_res = _compute_errorbar(task, cfg, metric, preds, all_labels)

        os.makedirs(output_dir, exist_ok=True)
        out_json = Path(output_dir) / f"{task}_{split}_results.json"
        with open(out_json, "w", encoding="utf-8") as f:
            json.dump({"task": task, "split": split, "metrics": metric_res}, f, indent=2)

        if errorbar_res is not None:
            out_eb = Path(output_dir) / f"{task}_{split}_errorbar.json"
            with open(out_eb, "w", encoding="utf-8") as f:
                json.dump({"task": task, "split": split, "errorbar": errorbar_res}, f, indent=2)

        csv_p = Path(output_dir) / f"{task}_{split}_preds.csv"
        pd.DataFrame({"pred": preds.tolist(), "label": all_labels.tolist()}).to_csv(csv_p, index=False)

        results_by_split[split] = {"metrics": metric_res, "errorbar": errorbar_res}

    return results_by_split

# ---------------------------------------------------------------------
# Top-level benchmark
# ---------------------------------------------------------------------
def run_glue_benchmark(
    config,
    tokenizer,
    model,
    checkpointing: Optional[Checkpointing] = None,
    out_dir: str = "glue_outputs",
):
    tasks = getattr(config, "glue_tasks", ["mnli"])
    # ["rte","cola"]
    # ["rte", "stsb", "mrpc", "cola", "sst2", "qnli", "qqp", "mnli"]
    batch_size = getattr(config, "batch_size", 64)
    max_length = getattr(config, "max_length", 128)
    device = getattr(config, "device", "cuda")

    auto_train = getattr(config, "auto_train", True)
    train_epochs = getattr(config, "train_epochs", 3)

    train_epochs_per_task = getattr(config, "train_epochs_per_task", {})
    if not train_epochs_per_task:
        train_epochs_per_task = {
            "cola": 5,
            "mrpc": 3,
            "rte": 5,
            "stsb": 3,
            "sst2": 3,
            "qqp": 3,
            "qnli": 3,
            "mnli": 3,
            "wnli": 5,
        }

    train_batch_size = getattr(config, "train_batch_size", 32)
    train_warmup_steps = getattr(config, "train_warmup_steps", 100)
    train_weight_decay = getattr(config, "train_weight_decay", 0.01)
    train_grad_accum_steps = getattr(config, "train_grad_accum_steps", 1)

    lr_candidates = getattr(config, "bert_lr_candidates", BERT_LR_CANDIDATES)
    random_restarts_small = int(getattr(config, "random_restarts_small", 100))
    base_seed = int(getattr(config, "base_seed", 3407))
    # 12345
    # 12345
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    checkpoint_out_root = out_dir / "checkpoints"
    checkpoint_out_root.mkdir(parents=True, exist_ok=True)

    rows = []

    # Save original pretrained state so every run starts identical
    original_state= None
    try:
        original_state = unwrap_model(model).state_dict()
    except Exception:
        try:
            original_state = model.state_dict()
        except Exception:
            original_state = None

    def _reset_model_to_original():
        if original_state is None:
            return
        try:
            unwrap_model(model).load_state_dict(original_state, strict=False)
        except Exception:
            try:
                model.load_state_dict(original_state, strict=False)
            except Exception:
                pass

    for task in tasks:
        assert task in GLUE_TASKS, f"Unknown GLUE task: {task}"
        print(f"\n==== Running GLUE task: {task} ====")

        epochs_this_task = int(train_epochs_per_task.get(task, train_epochs))
        print(f"[GLUE] Epochs for task '{task}': {epochs_this_task}")

        hf = load_dataset("glue", GLUE_TASKS[task]["hf_name"])
        train_raw = hf["train"]
        val_raw = hf["validation_matched"] if task == "mnli" else hf["validation"]

        task_ckpt_root = checkpoint_out_root / task
        task_ckpt_root.mkdir(parents=True, exist_ok=True)

        all_run_records = []

        best_run = {
            "score": None,
            "metrics": None,
            "lr": None,
            "restart": None,
            "seed": None,
            "best_epoch": None,
            "run_dir": None,
            "best_ckpt_path": None,
        }

        if auto_train:
            restarts = random_restarts_small if task in SMALL_TASKS_RANDOM_RESTARTS else 1
            print(f"[GLUE] Auto-training. LRs={lr_candidates}. Restarts/LR={restarts} (small={task in SMALL_TASKS_RANDOM_RESTARTS}).")

            for lr in lr_candidates:
                for restart_idx in range(restarts):
                    # stable, distinct seed per run
                    seed = base_seed + (abs(hash(task)) % 10000) * 1000 + int(restart_idx) * 10 + (int(round(lr * 1e7)) % 1000)

                    run_dir = task_ckpt_root / f"lr_{lr:g}" / f"restart_{restart_idx}"
                    run_dir.mkdir(parents=True, exist_ok=True)

                    print(f"\n[SWEEP] task={task} lr={lr:g} restart={restart_idx}/{restarts-1} seed={seed}")

                    _reset_model_to_original()

                    if checkpointing is not None:
                        try:
                            checkpointing.load_model_states("recent")
                        except Exception:
                            pass

                    _, metric_res, score, best_epoch = train_full_finetune(
                        task=task,
                        tokenizer=tokenizer,
                        model=model,
                        raw_train=train_raw,
                        raw_val=val_raw,
                        device=device,
                        epochs=epochs_this_task,
                        batch_size=train_batch_size,
                        lr=float(lr),
                        weight_decay=train_weight_decay,
                        warmup_steps=train_warmup_steps,
                        max_length=max_length,
                        grad_accum_steps=train_grad_accum_steps,
                        out_checkpoint_dir=str(run_dir),
                        seed=int(seed),
                    )

                    # Update run meta to include restart explicitly
                    meta_path = run_dir / "run_meta.json"
                    meta = {}
                    if meta_path.exists():
                        try:
                            meta = json.loads(meta_path.read_text(encoding="utf-8"))
                        except Exception:
                            meta = {}
                    meta.update({"restart": int(restart_idx)})
                    _json_dump(meta, meta_path)

                    rec = {
                        "task": task,
                        "lr": float(lr),
                        "restart": int(restart_idx),
                        "seed": int(seed),
                        "epochs": int(epochs_this_task),
                        "dev_score": float(score),
                        "dev_metrics": json.dumps(metric_res),
                        "best_epoch": int(best_epoch),
                        "run_dir": str(run_dir),
                        "best_ckpt_path": str(run_dir / "best_finetuned.pt"),
                        "final_ckpt_path": str(run_dir / "finetuned.pt"),
                    }
                    all_run_records.append(rec)

                    if best_run["score"] is None or float(score) > float(best_run["score"]):
                        best_run.update(
                            {
                                "score": float(score),
                                "metrics": metric_res,
                                "lr": float(lr),
                                "restart": int(restart_idx),
                                "seed": int(seed),
                                "best_epoch": int(best_epoch),
                                "run_dir": str(run_dir),
                                "best_ckpt_path": str(run_dir / "best_finetuned.pt"),
                            }
                        )

            # Save all runs CSV
            all_runs_csv = task_ckpt_root / "all_runs.csv"
            pd.DataFrame(all_run_records).to_csv(all_runs_csv, index=False)
            print(f"[GLUE] Saved all runs summary to: {all_runs_csv}")

            print(
                f"\n[GLUE] Best run for task='{task}': "
                f"lr={best_run['lr']:g}, restart={best_run['restart']}, seed={best_run['seed']}, "
                f"dev_score={best_run['score']}, dev_metrics={best_run['metrics']}"
            )

            # Copy best checkpoint to canonical folder
            best_overall_dir = task_ckpt_root / "best_overall"
            best_overall_dir.mkdir(parents=True, exist_ok=True)
            if best_run["best_ckpt_path"] and os.path.exists(best_run["best_ckpt_path"]):
                shutil.copy2(best_run["best_ckpt_path"], best_overall_dir / "best_finetuned.pt")
            _json_dump(best_run, best_overall_dir / "best_meta.json")

            # Load best model for eval
            _reset_model_to_original()
            try:
                model_for_eval = _load_finetuned_checkpoint_for_task(task, model, str(best_overall_dir / "best_finetuned.pt"))
            except Exception as e:
                print(f"[WARN] Failed to load best_overall checkpoint for eval; using current model. err={e}")
                model_for_eval = model

        else:
            if checkpointing is not None:
                try:
                    checkpointing.load_model_states("recent")
                except Exception:
                    pass
            model_for_eval = model

        # Evaluate best model + compute error bars
        task_out_dir = out_dir / task
        task_out_dir.mkdir(parents=True, exist_ok=True)

        res = run_glue_task(
            task=task,
            tokenizer=tokenizer,
            model=model_for_eval,
            checkpointing=None,
            device=device,
            batch_size=batch_size,
            max_length=max_length,
            output_dir=str(task_out_dir),
            compute_errorbar=True,  # <-- key: error bar only for final best model
        )

        # Write summary rows
        for split, pack in res.items():
            metrics = pack["metrics"]
            eb = pack["errorbar"]

            rows.append(
                {
                    "task": task,
                    "split": split,
                    "epochs": epochs_this_task,
                    "selected_lr": best_run["lr"] if best_run["lr"] is not None else None,
                    "selected_restart": best_run["restart"] if best_run["restart"] is not None else None,
                    "selected_seed": best_run["seed"] if best_run["seed"] is not None else None,
                    "selected_dev_score": best_run["score"] if best_run["score"] is not None else None,
                    "eval_metrics": json.dumps(metrics),
                    "errorbar_mean": (eb["mean"] if eb else None),
                    "errorbar_std": (eb["std"] if eb else None),
                    "errorbar_stderr": (eb["stderr"] if eb else None),
                    "errorbar_detail": json.dumps(eb) if eb else None,
                }
            )

    summary_csv = out_dir / "glue_summary.csv"
    pd.DataFrame(rows).to_csv(summary_csv, index=False)
    print(f"\n[GLUE] Summary saved to: {summary_csv}")

    return pd.DataFrame(rows)

# ---------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------
if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--tasks", type=str, default="sst2", help="comma separated glue tasks")
    parser.add_argument("--batch_size", type=int, default=64)
    parser.add_argument("--max_length", type=int, default=128)
    parser.add_argument("--device", type=str, default="cuda")
    parser.add_argument("--out_dir", type=str, default="glue_outputs")
    args = parser.parse_args()

    print("This module is intended to be invoked from your project's main, which provides tokenizer/model/checkpointing.")
    print("Example usage in your main: run_glue_benchmark(config.benchmark, tokenizer, model, checkpointing, out_dir=args.out_dir)")
    print(f"CLI args tasks={args.tasks} batch_size={args.batch_size} max_length={args.max_length} device={args.device} out_dir={args.out_dir}")