#!/usr/bin/env python3
import os
import json
import random
from pathlib import Path
from typing import Dict

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

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

from transformers import (
    AutoTokenizer,
    AutoConfig,
    AutoModel,
    TrainingArguments,
    Trainer,
    PreTrainedModel,
    set_seed,
)
from transformers.data.data_collator import DataCollatorWithPadding

# Drop-in replacement (cleaned + "what the model sees" showcase BEFORE training):
# - keeps your model class definitions unchanged
# - aligns training behavior to the sweep script:
#   * dynamic padding (pad_to_multiple_of=8)
#   * tokenizer.pad_token handling
#   * robust dropout setting + post-load dropout patching
#   * fused AdamW
#   * scheduler: cosine + num_cycles=4
#
# -----------------------------
# User knobs
# -----------------------------
logDir = "clean_cosine_restart_besthp_preview_fixed-wd-0.9_reproduce"
base_model_path = "./checkpoint-388560"
tokenizer_path  = "../regression_efficiency/checkpoint-5956"
train_json = "evenBetterDataFolded-tr.json"
valid_json = "evenBetterDataFolded-vl.json"

seed = 42
preview_n_texts = 4      # how many raw examples to preview
preview_tok_trunc = 200  # how many tokens to print per example (for readability)
preview_batch_size = 4   # for collator preview

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

writer = SummaryWriter(log_dir=f"tensorboard/{logDir}")

class LastTokenPooling(nn.Module):
    """
    Pool using the last non-padded token. Supports left- or right-padding.
    """
    def __init__(self):
        super().__init__()

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



class OneLayerRegressionHead(nn.Module):
    """
    Exactly one layernorm+linear, no residual-GELU block.
    """
    def __init__(self, hidden_size):
        super().__init__()
        self.net = nn.Sequential(
            nn.LayerNorm(hidden_size),
            nn.Linear(hidden_size, 1),
        )

    def forward(self, x):
        # x: [B, hidden_size]
        return self.net(x).squeeze(-1)   # [B]



class QwenForRegression(PreTrainedModel):
    """
    A regression model that uses only the base transformer (no LM head) and last-token pooling.
    """
    config_class = AutoConfig
    base_model_prefix = "backbone"

    def __init__(self, config, writer: SummaryWriter = None):
        super().__init__(config)
        # use base model without LM head to avoid unused lm_head parameters
        self.backbone = AutoModel.from_config(config)
        if getattr(config, "gradient_checkpointing", False):
            self.backbone.gradient_checkpointing_enable()

        hidden_size = config.hidden_size
        self.pooler = LastTokenPooling()
        self.regression_head = OneLayerRegressionHead(hidden_size)
        self.writer = writer
        self.step = 0

    def supports_gradient_checkpointing(self) -> bool:
        return True

    def gradient_checkpointing_enable(self, **kwargs):
        self.backbone.gradient_checkpointing_enable(**kwargs)

    def gradient_checkpointing_disable(self, **kwargs):
        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  # [B,T,H]
        pooled = self.pooler(hidden_states, attention_mask)  # [B,H]
        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
    ):
        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)
        model_to_save.backbone.save_pretrained(
            save_directory, state_dict=state_dict, **kwargs
        )

        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)
        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"
        backbone_sd, head_sd = {}, {}
    
        if fsdp_file.exists():
            fsdp_sd = torch.load(fsdp_file, map_location=device)
    
            # — strip the exact "backbone." prefix —
            for k, v in fsdp_sd.items():
                if k.startswith("backbone."):
                    new_k = k[len("backbone."):]
                    backbone_sd[new_k] = v
                elif k.startswith("regression_head."):
                    new_k = k[len("regression_head."):]
                    head_sd[new_k] = v
    
            # instantiate backbone from config
            backbone = AutoModel.from_config(config)
    
            # strict load: will now match
            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}"
                )
    
        else:
            # fallback to HF sharded .safetensors
            backbone = AutoModel.from_pretrained(model_dir, device_map=None,config=config)
            #head_sd  = torch.load(model_dir / "regression_head.pt", map_location=device)
    
        # build your full model
        model = cls(config, writer=writer)
        model.backbone = backbone.to(device)
        if hasattr(model.config, "use_cache"):
            model.config.use_cache = False
        if hasattr(model.backbone, "config") and hasattr(model.backbone.config, "use_cache"):
            model.backbone.config.use_cache = False
        # load regression head strictly, too
        missing_h, unexpected_h = model.regression_head.load_state_dict(head_sd, strict=False)
        # if missing_h or unexpected_h:
        #     raise RuntimeError(
        #         f"Head load mismatch.\n missing: {missing_h}\n unexpected: {unexpected_h}"
        #     )
    
        return model.to(device).eval()


# =========================
# Sweep-alignment utilities (APPLY CHANGES HERE)
# =========================
def robust_set_dropout(config, p_hidden: float, p_attn: float, layerdrop: float):
    """
    Mirror the sweep script: set all plausible dropout fields if present.
    """
    hidden_fields = [
        "hidden_dropout_prob", "hidden_dropout", "dropout",
        "emb_dropout", "resid_pdrop", "classifier_dropout",
    ]
    attn_fields = [
        "attention_probs_dropout_prob", "attention_dropout",
        "attn_dropout", "attn_pdrop",
    ]
    for f in hidden_fields:
        if hasattr(config, f):
            setattr(config, f, float(p_hidden))
    for f in attn_fields:
        if hasattr(config, f):
            setattr(config, f, float(p_attn))
    if hasattr(config, "layerdrop"):
        setattr(config, "layerdrop", float(layerdrop))


def patch_all_dropout_modules(model: nn.Module, p_hidden: float):
    """
    Post-load patch: in your from_pretrained else-branch,
    backbone is loaded via AutoModel.from_pretrained(model_dir, device_map=None)
    which may ignore our modified config.
    We therefore patch nn.Dropout modules in-place to match hidden dropout.
    """
    for m in model.modules():
        if isinstance(m, nn.Dropout):
            m.p = float(p_hidden)


def install_head_dropout(model: QwenForRegression, head_dropout: float):
    """
    head_dropout without changing class definition:
    swap net to LN -> Dropout -> Linear
    """
    hs = model.config.hidden_size
    model.regression_head.net = nn.Sequential(
        nn.LayerNorm(hs),
        nn.Dropout(p=float(head_dropout)),
        nn.Linear(hs, 1),
    )


# =========================
# Load tokenizer
# =========================
tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, use_fast=True)
tokenizer.pad_token = tokenizer.eos_token


# =========================
# Load data
# =========================
with open(train_json, "r", encoding="utf-8") as f:
    trainData = json.load(f)
with open(valid_json, "r", encoding="utf-8") as f:
    validData = json.load(f)

train_texts  = list(trainData.keys())
train_labels = [float(trainData[k]) for k in train_texts]
valid_texts  = list(validData.keys())
valid_labels = [float(validData[k]) for k in valid_texts]

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

print("Train examples:", train_texts[:2], train_labels[:2])
print("Valid examples:", valid_texts[:2], valid_labels[:2])


# =========================
# Tokenize (dynamic padding later)
# =========================
train_raw = Dataset.from_dict({"text": train_texts, "label": train_labels})
valid_raw = Dataset.from_dict({"text": valid_texts, "label": valid_labels})

def tok_fn(batch):
    return tokenizer(batch["text"], truncation=True, add_special_tokens=True)

train_dataset = train_raw.map(tok_fn, batched=True, remove_columns=["text"])
valid_dataset = valid_raw.map(tok_fn, batched=True, remove_columns=["text"])

train_dataset = train_dataset.rename_column("label", "labels")
valid_dataset = valid_dataset.rename_column("label", "labels")

train_dataset.set_format(type="torch")
valid_dataset.set_format(type="torch")

data_collator = DataCollatorWithPadding(
    tokenizer=tokenizer,
    pad_to_multiple_of=8,
    return_tensors="pt",
)


# =========================
# "What the model sees" preview (unchanged)
# =========================
def preview_tokenization_examples(texts, labels, tok, n=3, tok_trunc=200):
    print("\n==============================")
    print("PREVIEW: what the model sees")
    print("==============================")
    print("Tokenizer special_tokens_map:", tok.special_tokens_map)
    if getattr(tok, "additional_special_tokens", None):
        print("Tokenizer additional_special_tokens (count):", len(tok.additional_special_tokens))
        print("First few additional specials:", tok.additional_special_tokens[:10])

    idxs = list(range(min(n, len(texts))))
    for i in idxs:
        text = texts[i]
        y = labels[i]
        enc = tok(text, add_special_tokens=True)
        ids = enc["input_ids"]
        toks = tok.convert_ids_to_tokens(ids)

        print("\n--- Example", i, "---")
        print("Label:", y)
        print("Raw text (first 300 chars):")
        print(text[:300] + ("..." if len(text) > 300 else ""))

        print("\nToken IDs (truncated):")
        print(ids[:tok_trunc], "...(len=%d)" % len(ids) if len(ids) > tok_trunc else "(len=%d)" % len(ids))

        print("\nTokens (truncated):")
        print(toks[:tok_trunc], "...(len=%d)" % len(toks) if len(toks) > tok_trunc else "(len=%d)" % len(toks))

        decoded = tok.decode(ids, skip_special_tokens=False)
        print("\nDecoded (skip_special_tokens=False) first 400 chars:")
        print(decoded[:400] + ("..." if len(decoded) > 400 else ""))

def preview_collated_batch(ds, tok, collator, batch_size=4, tok_trunc=120):
    print("\n==============================")
    print("PREVIEW: collated batch (dynamic padding)")
    print("==============================")
    batch_items = [ds[i] for i in range(min(batch_size, len(ds)))]
    batch = collator(batch_items)

    input_ids = batch["input_ids"]
    attn = batch["attention_mask"]
    labels = batch["labels"]

    print("Batch shapes:",
          "input_ids", tuple(input_ids.shape),
          "attention_mask", tuple(attn.shape),
          "labels", tuple(labels.shape))

    pad_id = tok.pad_token_id

    for r in range(min(2, input_ids.shape[0])):
        ids = input_ids[r].tolist()
        toks = tok.convert_ids_to_tokens(ids)

        visible_len = int((np.array(ids) != pad_id).sum()) if pad_id is not None else int(attn[r].sum().item())

        print(f"\n--- Batch row {r} ---")
        print("Label:", float(labels[r].item()))
        print("Non-pad token length:", visible_len)

        print("IDs (truncated):")
        print(ids[:tok_trunc], "...")

        print("Tokens (truncated):")
        print(toks[:tok_trunc], "...")

        decoded = tok.decode(ids, skip_special_tokens=False)
        print("Decoded (skip_special_tokens=False) first 400 chars:")
        print(decoded[:400] + ("..." if len(decoded) > 400 else ""))

preview_tokenization_examples(train_texts, train_labels, tokenizer, n=preview_n_texts, tok_trunc=preview_tok_trunc)
preview_collated_batch(train_dataset, tokenizer, data_collator, batch_size=preview_batch_size, tok_trunc=preview_tok_trunc)


# =========================
# Metrics (keep)
# =========================
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:
        pr, _ = pearsonr(preds, labels)
    else:
        pr = 0.0

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


# =========================
# Best hparams you provided
# =========================
HP = dict(
    learning_rate=5e-5,
    weight_decay=0.9,
    hidden_dropout=0.3,
    attn_dropout=0.3,
    layerdrop=0.05,
    head_dropout=0.2,
    max_grad_norm=2,
    warmup_ratio=0.1,
    lr_scheduler_type="cosine",  # <-- IMPORTANT: match your best hp
    num_cycles=40,
)


# =========================
# Build config (robust) + load model
# =========================
config = AutoConfig.from_pretrained(base_model_path)
robust_set_dropout(config, HP["hidden_dropout"], HP["attn_dropout"], HP["layerdrop"])

model = QwenForRegression.from_pretrained(
    base_model_path,
    device="cuda",
    writer=writer,
    config=config,
)

# Patch dropout after load to avoid config-ignored behavior
patch_all_dropout_modules(model, p_hidden=HP["hidden_dropout"])

# Head dropout without changing model class
install_head_dropout(model, head_dropout=HP["head_dropout"])

# Disable cache (saves VRAM) - keep
if hasattr(model.config, "use_cache"):
    model.config.use_cache = False
if hasattr(model.backbone, "config") and hasattr(model.backbone.config, "use_cache"):
    model.backbone.config.use_cache = False


def print_trainable_summary(m):
    total = 0
    trainable = 0
    for _, p in m.named_parameters():
        n = p.numel()
        total += n
        if p.requires_grad:
            trainable += n
    print(f"\nTotal parameters:     {total:,}")
    print(f"Trainable parameters: {trainable:,}")
    print(f"Frozen parameters:    {total-trainable:,}")

print_trainable_summary(model)


# =========================
# TrainingArguments (apply changes)
# =========================
USE_BF16 = torch.cuda.is_available()

training_args = TrainingArguments(
    output_dir=f"./qwen_regression_ckpt/{logDir}",

    per_device_train_batch_size=24,   # match sweep default unless you intentionally want 16
    per_device_eval_batch_size=24,
    gradient_accumulation_steps=1,
    num_train_epochs=1000,

    learning_rate=HP["learning_rate"],
    weight_decay=HP["weight_decay"],
    max_grad_norm=HP["max_grad_norm"],

    warmup_steps=21600,

    lr_scheduler_type=HP["lr_scheduler_type"],          # cosine
    lr_scheduler_kwargs={"num_cycles": HP["num_cycles"]},

    bf16=USE_BF16,

    logging_dir=f"tensorboard/{logDir}",
    logging_steps=10,
    report_to="tensorboard",

    optim="adamw_torch_fused",  # match sweep

    eval_strategy="epoch",  # HF-standard spelling
    save_strategy="epoch",
    save_total_limit=2,
    load_best_model_at_end=True,
    metric_for_best_model="mse",
    greater_is_better=False,

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

    dataloader_pin_memory=True,
    remove_unused_columns=False,

    seed=seed,
    data_seed=seed,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=valid_dataset,
    data_collator=data_collator,
    compute_metrics=compute_metrics,
    tokenizer=tokenizer,
)
print('dropout')
# show a few dropout modules
cnt = 0
for n, m in model.named_modules():
    if isinstance(m, nn.Dropout):
        print("dropout:", n, "p=", m.p)
        cnt += 1
        if cnt >= 8:
            break
print(model.backbone.config)
trainer.train()