#!/usr/bin/env python3
# Grokking-oriented full-parameter regression training (DROP-IN, with YOUR save/load style fixed)
#
# Key points:
# - NO freezing hacks
# - NO PEFT
# - Labels unchanged
# - Uses your QwenForRegression wrapper (AutoModel backbone + last-token pooling + 1-layer head)
# - SAVE/LOAD works:
#     * saves backbone weights with correct (stripped) state_dict keys
#     * saves regression head to regression_head.pt
#     * loads from either:
#         - pytorch_model_fsdp.bin (if present) OR
#         - normal HF checkpoint files (safetensors/bin) + regression_head.pt (if present)
# - LR schedule: warmup + FLOORED cosine with RESTARTS (num_cycles), implemented robustly
# - Uses eval_strategy (and falls back to evaluation_strategy if needed)
#
# Run:
#   python train_grokking_fullft.py
#
# TensorBoard:
#   tensorboard/<logDir>/
#
# Checkpoints:
#   ./qwen_regression_ckpt/<logDir>/checkpoint-... (plus regression_head.pt inside each)

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

import torch
from torch import nn
from torch.nn import functional as F
from torch.utils.tensorboard import SummaryWriter

import numpy as np
from datasets import Dataset
from sklearn.metrics import mean_squared_error, r2_score
from scipy.stats import pearsonr

from transformers import (
    AutoTokenizer, AutoConfig, TrainingArguments, Trainer, PreTrainedModel
)
from transformers import AutoModel  # base model (no LM head)


# =========================
# 0) USER SETTINGS (drop-in)
# =========================
base_model_path = "./checkpoint-388560"
tokenizer_path  = "/opt/platform/regression_efficiency/checkpoint-5956"

train_json_path = "evenBetterDataFolded-tr.json"
valid_json_path = "evenBetterDataFolded-vl.json"

# --- grokking-ish defaults ---
SEED = 42

NUM_EPOCHS = 1500                    # grokking requires long training
LR = 3e-5
WEIGHT_DECAY = 0.03                  # grokking dial (try 0.01 / 0.03 / 0.1)
WARMUP_RATIO = 0.02                  # warmup fraction of total steps

# floored cosine w/ restarts:
MIN_LR_RATIO = 0.10                  # floor as fraction of LR
NUM_CYCLES = 4                       # restarts

# turn off dropout-ish knobs (let weight decay be the main regularizer)
HIDDEN_DROPOUT = 0.0
ATTN_DROPOUT   = 0.0
LAYERDROP      = 0.0

PER_DEVICE_TRAIN_BATCH = 24
PER_DEVICE_EVAL_BATCH  = 24
GRAD_ACCUM_STEPS = 1

MAX_GRAD_NORM = 1.0
BF16 = True

# logging / output
logDir = f"grokking_wd{WEIGHT_DECAY}_lr{LR}_floor{MIN_LR_RATIO}_cyc{NUM_CYCLES}"
out_dir = f"./qwen_regression_ckpt/{logDir}"
tb_dir  = f"tensorboard/{logDir}"

os.makedirs(out_dir, exist_ok=True)
os.makedirs(tb_dir, exist_ok=True)

torch.manual_seed(SEED)
random.seed(SEED)
np.random.seed(SEED)


# =========================
# 1) helper: TrainingArguments compat (eval_strategy vs evaluation_strategy)
# =========================
_TA_SIG = inspect.signature(TrainingArguments.__init__)
_HAS_EVAL_STRATEGY = "eval_strategy" in _TA_SIG.parameters
_HAS_EVALUATION_STRATEGY = "evaluation_strategy" in _TA_SIG.parameters

def make_training_args(**kwargs):
    if "eval_strategy" in kwargs and not _HAS_EVAL_STRATEGY and _HAS_EVALUATION_STRATEGY:
        kwargs["evaluation_strategy"] = kwargs.pop("eval_strategy")
    if "evaluation_strategy" in kwargs and _HAS_EVAL_STRATEGY and not _HAS_EVALUATION_STRATEGY:
        kwargs["eval_strategy"] = kwargs.pop("evaluation_strategy")
    return TrainingArguments(**kwargs)


# =========================
# 2) model components
# =========================
class LastTokenPooling(nn.Module):
    """Pool using last non-pad token. Supports left/right padding."""
    def __init__(self):
        super().__init__()

    def forward(self, hidden_states, attention_mask=None):
        if attention_mask is None:
            return hidden_states[:, -1, :]
        B, T, H = hidden_states.size()
        if attention_mask[:, -1].sum().item() == B:
            return hidden_states[:, -1, :]
        seq_lens = attention_mask.sum(dim=1).long() - 1
        idx = seq_lens.view(B, 1, 1).expand(-1, 1, H)
        return hidden_states.gather(1, idx).squeeze(1)

class OneLayerRegressionHead(nn.Module):
    """LayerNorm + Linear(hidden->1)."""
    def __init__(self, hidden_size):
        super().__init__()
        self.net = nn.Sequential(
            nn.LayerNorm(hidden_size),
            nn.Linear(hidden_size, 1),
        )

    def forward(self, x):
        return self.net(x).squeeze(-1)


# =========================
# 3) YOUR QwenForRegression with WORKING save/load (fixed)
# =========================
class QwenForRegression(PreTrainedModel):
    """
    Regression wrapper: AutoModel backbone + last-token pooling + 1-layer head.
    Save/Load:
      - backbone saved with stripped keys (so AutoModel.from_pretrained works)
      - head saved to regression_head.pt
      - can load from FSDP bin or normal HF checkpoint
    """
    config_class = AutoConfig
    base_model_prefix = "backbone"

    def __init__(self, config, writer: SummaryWriter = None):
        super().__init__(config)
        self.backbone = AutoModel.from_config(config)
        if getattr(config, "gradient_checkpointing", False):
            self.backbone.gradient_checkpointing_enable()

        self.pooler = LastTokenPooling()
        self.regression_head = OneLayerRegressionHead(config.hidden_size)

        self.writer = writer
        self.step = 0

    def supports_gradient_checkpointing(self) -> bool:
        return True

    def gradient_checkpointing_enable(self, **kwargs):
        if hasattr(self.backbone, "gradient_checkpointing_enable"):
            self.backbone.gradient_checkpointing_enable(**kwargs)

    def gradient_checkpointing_disable(self, **kwargs):
        if hasattr(self.backbone, "gradient_checkpointing_disable"):
            self.backbone.gradient_checkpointing_disable(**kwargs)

    def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
        outputs = self.backbone(
            input_ids=input_ids,
            attention_mask=attention_mask,
            return_dict=True,
            output_hidden_states=False,
        )
        hidden_states = outputs.last_hidden_state
        pooled = self.pooler(hidden_states, attention_mask)

        if self.writer is not None:
            self.writer.add_scalar("pooled/mean", pooled.mean().item(), self.step)
            self.writer.add_scalar("pooled/std", pooled.std().item(), self.step)
            self.step += 1

        logits = self.regression_head(pooled)

        if labels is not None:
            loss = F.mse_loss(logits, labels)
            return {"loss": loss, "logits": logits}
        return {"logits": logits}

    def save_pretrained(self, save_directory: str, state_dict=None, accelerator=None, **kwargs):
        """
        IMPORTANT FIX vs the buggy pattern:
        - if Trainer passes a full state_dict with 'backbone.' prefix, strip it before saving backbone,
          otherwise AutoModel.from_pretrained won't load.
        """
        model_to_save = self
        if accelerator is not None:
            model_to_save = accelerator.unwrap_model(self)

        os.makedirs(save_directory, exist_ok=True)
        model_to_save.config.save_pretrained(save_directory)

        if state_dict is None:
            state_dict = model_to_save.state_dict()

        # split state_dict
        backbone_sd = {}
        head_sd = {}

        for k, v in state_dict.items():
            if k.startswith("backbone."):
                backbone_sd[k[len("backbone."):]] = v
            elif k.startswith("regression_head."):
                head_sd[k[len("regression_head."):]] = v

        # save backbone in HF format (now keys match)
        model_to_save.backbone.save_pretrained(save_directory, state_dict=backbone_sd, **kwargs)

        # save head
        # also reshape linear weight if it ever comes flattened (paranoia)
        try:
            # prefer head_sd from state_dict (correct for FSDP)
            if not head_sd:
                head_sd = model_to_save.regression_head.state_dict()

            for idx, layer in enumerate(model_to_save.regression_head.net):
                if isinstance(layer, nn.Linear):
                    w_key = f"net.{idx}.weight"
                    if w_key in head_sd:
                        w = head_sd[w_key]
                        out_f, in_f = layer.out_features, layer.in_features
                        if w.dim() == 1 and w.numel() == out_f * in_f:
                            head_sd[w_key] = w.view(out_f, in_f)
        except Exception:
            pass

        torch.save(head_sd, os.path.join(save_directory, "regression_head.pt"))

    @classmethod
    def from_pretrained(cls, model_path, device="cpu", config=None, writer=None):
        model_dir = Path(model_path)
        config = config or AutoConfig.from_pretrained(model_dir)

        fsdp_file = model_dir / "pytorch_model_fsdp.bin"

        # build wrapper
        model = cls(config, writer=writer)

        if fsdp_file.exists():
            fsdp_sd = torch.load(fsdp_file, map_location=device)

            backbone_sd = {}
            head_sd = {}

            for k, v in fsdp_sd.items():
                if k.startswith("backbone."):
                    backbone_sd[k[len("backbone."):]] = v
                elif k.startswith("regression_head."):
                    head_sd[k[len("regression_head."):]] = v

            backbone = AutoModel.from_config(config)
            missing_b, unexpected_b = backbone.load_state_dict(backbone_sd, strict=True)
            if missing_b or unexpected_b:
                raise RuntimeError(f"Backbone load mismatch.\n missing: {missing_b}\n unexpected: {unexpected_b}")

            model.backbone = backbone.to(device)

            # load head if present in fsdp bin
            if head_sd:
                missing_h, unexpected_h = model.regression_head.load_state_dict(head_sd, strict=False)
                if missing_h or unexpected_h:
                    print(f"[warn] head load non-strict. missing={missing_h} unexpected={unexpected_h}")

        else:
            # normal HF checkpoint: load backbone from files
            model.backbone = AutoModel.from_pretrained(model_dir, device_map=None).to(device)

            # load head if file exists; otherwise keep random-init head (useful when starting from LM)
            head_file = model_dir / "regression_head.pt"
            if head_file.exists():
                head_sd = torch.load(head_file, map_location=device)
                missing_h, unexpected_h = model.regression_head.load_state_dict(head_sd, strict=False)
                if missing_h or unexpected_h:
                    print(f"[warn] head load non-strict. missing={missing_h} unexpected={unexpected_h}")
            else:
                print("[info] regression_head.pt not found; using randomly initialized regression head.")

        return model.to(device)


# =========================
# 4) grokking LR schedule: warmup + floored cosine with restarts
# =========================
def build_warmup_cosine_restart_floor_lambda(
    num_warmup_steps: int,
    num_training_steps: int,
    min_lr_ratio: float,
    num_cycles: int,
):
    num_warmup_steps = int(max(0, num_warmup_steps))
    num_training_steps = int(max(1, num_training_steps))
    num_cycles = int(max(1, num_cycles))
    min_lr_ratio = float(min_lr_ratio)

    def lr_lambda(step: int) -> float:
        if step < num_warmup_steps:
            return float(step) / float(max(1, num_warmup_steps))

        remaining = max(1, num_training_steps - num_warmup_steps)
        t = step - num_warmup_steps
        t = min(max(t, 0), remaining)

        # cycle length in steps (integer-ish)
        cycle_len = max(1, remaining // num_cycles)
        cycle_pos = (t % cycle_len) / float(cycle_len)  # [0,1)

        cosine = 0.5 * (1.0 + math.cos(math.pi * cycle_pos))
        return min_lr_ratio + (1.0 - min_lr_ratio) * cosine

    return lr_lambda


class GrokkingTrainer(Trainer):
    def create_scheduler(self, num_training_steps: int, optimizer: Optional[torch.optim.Optimizer] = None):
        if optimizer is None:
            optimizer = self.optimizer

        # prefer warmup_ratio if provided; else warmup_steps
        warmup_steps = 0
        if getattr(self.args, "warmup_ratio", None) is not None and self.args.warmup_ratio > 0:
            warmup_steps = int(self.args.warmup_ratio * num_training_steps)
        elif getattr(self.args, "warmup_steps", 0) and self.args.warmup_steps > 0:
            warmup_steps = int(self.args.warmup_steps)

        lr_lambda = build_warmup_cosine_restart_floor_lambda(
            num_warmup_steps=warmup_steps,
            num_training_steps=num_training_steps,
            min_lr_ratio=MIN_LR_RATIO,
            num_cycles=NUM_CYCLES,
        )
        self.lr_scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
        return self.lr_scheduler


# =========================
# 5) metrics
# =========================
def compute_metrics(eval_pred):
    preds, labels = eval_pred
    if isinstance(preds, (tuple, list)):
        preds = preds[0]
    preds = np.asarray(preds).reshape(-1)
    labels = np.asarray(labels).reshape(-1)

    mse = mean_squared_error(labels, preds)
    r2 = r2_score(labels, preds)

    if np.std(preds) > 1e-8 and np.std(labels) > 1e-8:
        pearson_r, _ = pearsonr(preds, labels)
    else:
        pearson_r = 0.0

    return {"mse": float(mse), "r2": float(r2), "pearson_r": float(pearson_r)}


# =========================
# 6) optional: evaluate train each epoch to visualize grokking gap
# =========================
from transformers import TrainerCallback

class TrainEvalAndGapCallback(TrainerCallback):
    def __init__(self, train_dataset):
        self.train_dataset = train_dataset
        self.trainer_ref = None
        self._guard = False

    def on_evaluate(self, args, state, control, metrics=None, **kwargs):
        if self.trainer_ref is None or self._guard:
            return control
        self._guard = True
        train_metrics = self.trainer_ref.evaluate(self.train_dataset, metric_key_prefix="train")
        self._guard = False

        if metrics:
            gap_logs = {}
            for k_eval, v_eval in metrics.items():
                if not k_eval.startswith("eval_"):
                    continue
                k_train = "train_" + k_eval[len("eval_"):]
                if k_train in train_metrics:
                    gap_logs["gap_" + k_eval[len("eval_"):]] = float(v_eval - train_metrics[k_train])
            if gap_logs:
                self.trainer_ref.log(gap_logs)
        return control


# =========================
# 7) data loading (your style: pad to max_length)
# =========================
tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

writer = SummaryWriter(log_dir=tb_dir)

train_texts, train_labels = [], []
valid_texts, valid_labels = [], []
produceCopyTrain, produceCopyValid = {}, {}
max_length = 0

with open(train_json_path, "r", encoding="utf-8") as f:
    trainData = json.load(f)

for seq, label in trainData.items():
    train_texts.append(seq)
    train_labels.append(label)
    produceCopyTrain[seq] = label
    max_length = max(max_length, len(seq))

with open(valid_json_path, "r", encoding="utf-8") as f:
    validData = json.load(f)

for seq, label in validData.items():
    valid_texts.append(seq)
    valid_labels.append(label)
    produceCopyValid[seq] = label
    max_length = max(max_length, len(seq))

os.makedirs(tb_dir, exist_ok=True)
with open(f"{tb_dir}/train.json", "w", encoding="utf-8") as jf:
    json.dump(produceCopyTrain, jf, indent=2)
with open(f"{tb_dir}/valid.json", "w", encoding="utf-8") as jf:
    json.dump(produceCopyValid, jf, indent=2)

print(train_texts[:2], train_labels[:2])
print(valid_texts[:2], valid_labels[:2])
print("max_length (string-based):", max_length)

def preprocess(example):
    enc = tokenizer(
        example["text"],
        padding="max_length",
        truncation=True,
        max_length=max_length,
        return_tensors="pt",
    )
    return {
        "input_ids": enc["input_ids"][0],
        "attention_mask": enc["attention_mask"][0],
        "labels": torch.tensor(example["label"], dtype=torch.float),
    }

train_raw = Dataset.from_dict({"text": train_texts, "label": train_labels})
valid_raw = Dataset.from_dict({"text": valid_texts, "label": valid_labels})

train_dataset = train_raw.map(preprocess)
valid_dataset = valid_raw.map(preprocess)


# =========================
# 8) config + model init (start from base_model_path; if head missing, random init head)
# =========================
config = AutoConfig.from_pretrained(
    base_model_path,
    hidden_dropout_prob=HIDDEN_DROPOUT,
    attention_probs_dropout_prob=ATTN_DROPOUT,
    layerdrop=LAYERDROP,
)

# This will:
# - load backbone from base_model_path
# - load regression_head.pt if present, else random head (good for LM->regression)
model = QwenForRegression.from_pretrained(
    base_model_path,
    device="cuda" if torch.cuda.is_available() else "cpu",
    writer=writer,
    config=config,
)

# ensure full param training
for p in model.parameters():
    p.requires_grad = True


# =========================
# 9) training args (floored cosine restart via GrokkingTrainer override)
# =========================
training_args = make_training_args(
    output_dir=out_dir,
    per_device_train_batch_size=PER_DEVICE_TRAIN_BATCH,
    per_device_eval_batch_size=PER_DEVICE_EVAL_BATCH,
    gradient_accumulation_steps=GRAD_ACCUM_STEPS,
    num_train_epochs=NUM_EPOCHS,

    logging_dir=tb_dir,
    logging_steps=1,

    eval_strategy="epoch",
    save_strategy="epoch",
    save_total_limit=2,

    report_to="tensorboard",

    learning_rate=LR,
    weight_decay=WEIGHT_DECAY,               # grokking dial
    warmup_ratio=WARMUP_RATIO,

    # We override scheduler anyway, but keep these fields sane.
    lr_scheduler_type="constant",

    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"use_reentrant": False},

    max_grad_norm=MAX_GRAD_NORM,

    bf16=BF16,

    load_best_model_at_end=True,
    metric_for_best_model="pearson_r",       # Trainer will look for eval_pearson_r
    greater_is_better=True,

    remove_unused_columns=False,
    seed=SEED,
    data_seed=SEED,
)

# =========================
# 10) train
# =========================
gap_cb = TrainEvalAndGapCallback(train_dataset)

trainer = GrokkingTrainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=valid_dataset,
    compute_metrics=compute_metrics,
    callbacks=[gap_cb],
)
gap_cb.trainer_ref = trainer

trainer.train()

print("\nDone. Best checkpoint should have working backbone + regression_head.pt.")
print("To reload a saved checkpoint later:")
print(f"  m = QwenForRegression.from_pretrained('{out_dir}/checkpoint-XXXX', device='cuda', config=AutoConfig.from_pretrained('{out_dir}/checkpoint-XXXX'))")