#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Pretrain Veronica-Polymorphic from scratch (clean mixture: FinePDFs / DCLM / FineWeb-Edu).

Basic example:
python veronica-polymorphic/scripts/train_veronica.py \
    --config veronica-polymorphic/configs/veronica-pretrain-12L.json \
    --dataset_paths data/mix_optimal_50_30_20_2048 \
    --output_dir veronica-polymorphic/runs/veronica-pretrain-vMix-2048 \
    --per_device_train_batch_size 4 \
    --gradient_accumulation_steps 4 \
    --learning_rate 2e-4 \
    --label_smoothing 0.01 \
    --rep_alpha 0.0 \
    --max_steps 60000 \
    --max_seq_len 2048

You can use different datasets (e.g., 512 / 1024 / 2048) in separate runs for length curriculum.
"""

import os
import re
import glob
import json
import math
import argparse
import random
from dataclasses import dataclass
from typing import Dict, List, Optional

import torch
import torch.nn.functional as F
from datasets import load_from_disk
from transformers import (
    AutoTokenizer,
    Trainer,
    TrainingArguments,
    TrainerCallback,
    CONFIG_MAPPING,
    MODEL_FOR_CAUSAL_LM_MAPPING,
    LogitsProcessorList,
    NoRepeatNGramLogitsProcessor,
    RepetitionPenaltyLogitsProcessor,
)

# --- Veronica bindings ---
from veronica.configuration_veronica import VeronicaConfig
from veronica.modeling_veronica import VeronicaForCausalLM
from veronica.modeling_components import Fp32LayerNorm

CONFIG_MAPPING.register("veronica", VeronicaConfig)
MODEL_FOR_CAUSAL_LM_MAPPING.register(VeronicaConfig, VeronicaForCausalLM)

# Disable CUDA Graphs (HF Trainer + torch.compile may conflict sometimes)
os.environ.setdefault("TORCH_COMPILE_USE_CUDAGRAPHS", "0")
os.environ.setdefault("TORCHINDUCTOR_DISABLE_CUDAGRAPHS", "1")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")


# ===========================
# Utility
# ===========================

def find_latest_checkpoint(run_dir: str) -> Optional[str]:
    ckpts = glob.glob(os.path.join(run_dir, "checkpoint-*"))
    if not ckpts:
        return None
    ckpts.sort(key=lambda p: int(re.search(r"checkpoint-(\d+)", p).group(1)))
    return ckpts[-1]


def build_tokenizer(candidates: List[str], save_dir: str) -> AutoTokenizer:
    """
    Try to load an existing tokenizer from the provided paths;
    otherwise fallback to gpt2 and add basic special tokens.
    """
    tok = None
    for p in candidates:
        if os.path.exists(p):
            try:
                tok = AutoTokenizer.from_pretrained(p, use_fast=True)
                print(f"[tokenizer] loaded from {p}")
                break
            except Exception:
                pass
    if tok is None:
        print("[tokenizer] fallback: gpt2")
        tok = AutoTokenizer.from_pretrained("gpt2", use_fast=True)

    specials: Dict[str, str] = {}
    if tok.eos_token is None:
        specials["eos_token"] = "<|eos|>"
    if tok.pad_token is None:
        specials["pad_token"] = "<|pad|>"
    if tok.bos_token is None:
        specials["bos_token"] = "<|bos|>"

    if specials:
        tok.add_special_tokens(specials)

    tok.save_pretrained(save_dir)
    tok = AutoTokenizer.from_pretrained(save_dir, use_fast=True)
    base_vocab = tok.vocab_size
    effective_vocab = len(tok)
    print(
        f"[tokenizer] base_vocab={base_vocab} added={effective_vocab - base_vocab} "
        f"effective_vocab={effective_vocab} eos={tok.eos_token_id} "
        f"pad={tok.pad_token_id} bos={tok.bos_token_id}"
    )
    return tok


def load_cfg_with_vocab(cfg_path: str, tok: AutoTokenizer) -> VeronicaConfig:
    """
    Load the config and adapt it to the tokenizer vocabulary.
    Model is designed as UN-TIED (lm_head != wte).
    """
    with open(cfg_path, "r", encoding="utf-8") as f:
        d = json.load(f)
    cfg = VeronicaConfig(**d)
    cfg.model_type = "veronica"
    cfg.vocab_size = int(len(tok))
    # untied model: no tie_word_embeddings
    return cfg


def init_model_from_config(cfg: VeronicaConfig, tok: AutoTokenizer) -> VeronicaForCausalLM:
    model = VeronicaForCausalLM(cfg)
    use_bf16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported()
    dtype = torch.bfloat16 if use_bf16 else (torch.float16 if torch.cuda.is_available() else torch.float32)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(dtype=dtype, device=device)

    effective_vocab = len(tok)
    emb = model.get_input_embeddings().weight
    head = model.lm_head.weight

    # Align embedding/head to the effective vocab
    if emb.shape[0] != effective_vocab or head.shape[0] != effective_vocab:
        old_vocab = emb.shape[0]
        print(f"[model] resize_token_embeddings: {old_vocab} -> {effective_vocab}")
        model.resize_token_embeddings(effective_vocab)
        with torch.no_grad():
            new_emb = model.get_input_embeddings().weight
            new_head = model.lm_head.weight
            mean_emb = new_emb[:old_vocab].mean(dim=0, keepdim=True)
            mean_head = new_head[:old_vocab].mean(dim=0, keepdim=True)
            if effective_vocab > old_vocab:
                new_emb[old_vocab:] = mean_emb
                new_head[old_vocab:] = mean_head

    # Keep LayerNorm params in float32 (after global cast)
    for m in model.modules():
        if isinstance(m, Fp32LayerNorm):
            m.ln.to(dtype=torch.float32)

    model.config.use_cache = False
    n_params = sum(p.numel() for p in model.parameters())
    print(f"[model] params={n_params:,} vocab={effective_vocab}")
    return model


def load_mix_dataset(path: str):
    """
    Load a packed dataset (train/validation) from disk.
    Expected HuggingFace formats: a DatasetDict with 'train' and 'validation',
    or a single Dataset that gets split 99/1.
    """
    ds = load_from_disk(path)
    if isinstance(ds, dict) and "train" in ds and "validation" in ds:
        return ds["train"], ds["validation"]
    split = ds.train_test_split(test_size=0.01, seed=42)
    return split["train"], split["test"]


# ===========================
# Collator
# ===========================

@dataclass
class CausalCollator:
    tokenizer: AutoTokenizer
    mask_runs: bool = False
    run_len: int = 4
    max_seq_len: Optional[int] = None  # target length (e.g., 512/1024/2048)

    def _mask_degenerate_runs(self, labels: torch.Tensor):
        """
        Mask degenerate runs (e.g., '____', '....') with length >= run_len.
        Mostly legacy; can be left off with a clean dataset.
        """
        try:
            id_us = self.tokenizer.encode("_", add_special_tokens=False)[0]
            id_dot = self.tokenizer.encode(".", add_special_tokens=False)[0]
        except Exception:
            return
        B, T = labels.size()
        for b in range(B):
            cnt_u = cnt_d = 0
            for t in range(T):
                tok = int(labels[b, t].item())
                if tok == id_us:
                    cnt_u += 1
                    cnt_d = 0
                elif tok == id_dot:
                    cnt_d += 1
                    cnt_u = 0
                else:
                    cnt_u = cnt_d = 0
                if cnt_u >= self.run_len or cnt_d >= self.run_len:
                    labels[b, t] = -100

    def _crop(self, ids: torch.Tensor) -> torch.Tensor:
        """
        If max_seq_len is set and the sequence is longer,
        crop a random window of length max_seq_len.
        """
        if self.max_seq_len is None:
            return ids
        L = ids.size(0)
        if L <= self.max_seq_len:
            return ids
        start = random.randint(0, L - self.max_seq_len)
        end = start + self.max_seq_len
        return ids[start:end]

    def __call__(self, features):
        ids_list = []
        for f in features:
            ids = torch.tensor(f["input_ids"], dtype=torch.long)
            ids = self._crop(ids)
            ids_list.append(ids)

        pad_id = self.tokenizer.pad_token_id or self.tokenizer.eos_token_id
        ids = torch.nn.utils.rnn.pad_sequence(ids_list, batch_first=True, padding_value=pad_id)
        attn = torch.where(ids == pad_id, 0, 1)

        labels = ids.clone()
        labels[labels == pad_id] = -100
        if self.mask_runs:
            self._mask_degenerate_runs(labels)

        return {"input_ids": ids, "attention_mask": attn, "labels": labels}


# ===========================
# Callback Router + Smoke eval
# ===========================

SMOKE_PROMPTS = [
    "The world we live in today is",
    "Understanding complex ideas requires",
    "Human intelligence differs from artificial intelligence because",
    "A good system design is based on",
    "In the middle of every difficulty lies",
    "Once upon a time, there was a scientist who",
]


class RouterAndSmokeCallback(TrainerCallback):
    def __init__(self, tok: AutoTokenizer):
        self.tok = tok

    def on_log(self, args, state, control, **kwargs):
        model = kwargs.get("model", None)
        if model is None:
            return
        try:
            if hasattr(model, "router_alpha_mean") and model.router_alpha_mean is not None:
                alpha = model.router_alpha_mean.detach().float().cpu()
                p = alpha / alpha.sum()
                ent = -(p * (p.clamp_min(1e-9)).log()).sum()
                ent_norm = float(ent / math.log(len(p)))
                print(f"[router] alpha={alpha.tolist()} entropy_norm={ent_norm:.4f}")
        except Exception:
            pass

    def on_evaluate(self, args, state, control, **kwargs):
        model = kwargs.get("model", None)
        if model is None:
            return
        model.eval()
        dev = next(model.parameters()).device

        prompt = random.choice(SMOKE_PROMPTS)
        ids = self.tok(prompt, return_tensors="pt").to(dev)

        processors = LogitsProcessorList([
            NoRepeatNGramLogitsProcessor(3),
            RepetitionPenaltyLogitsProcessor(1.1),
        ])

        with torch.no_grad():
            out = model.generate(
                **ids,
                max_new_tokens=64,
                do_sample=False,
                logits_processor=processors,
                eos_token_id=self.tok.eos_token_id,
                pad_token_id=(self.tok.pad_token_id or self.tok.eos_token_id),
                use_cache=True,
            )
        txt = self.tok.decode(out[0], skip_special_tokens=True)
        completion = txt[len(prompt):].strip() if txt.startswith(prompt) else txt
        print(f"\n[SMOKE] {prompt} → {completion}\n")
        model.train()


# ===========================
# Callback schedule router_tau / aux_weight
# ===========================

class RouterScheduleCallback(TrainerCallback):
    """
    Linearly schedule router_tau and router_aux_weight between start and end of training.
    """

    def __init__(
        self,
        tau_start: float,
        tau_end: float,
        aux_start: float,
        aux_end: float,
        total_steps: int,
        tau_freeze_steps: int = 0,
        force_prob: float = 0.0,
        force_warmup_steps: int = 0,
    ):
        self.tau_start = float(tau_start)
        self.tau_end = float(tau_end)
        self.aux_start = float(aux_start)
        self.aux_end = float(aux_end)
        self.total_steps = max(int(total_steps), 1)
        self.tau_freeze_steps = max(int(tau_freeze_steps), 0)
        self.force_prob = float(force_prob)
        self.force_warmup_steps = max(int(force_warmup_steps), 0)

    def _interp(self, start: float, end: float, step: int, span: int) -> float:
        t = min(max(step, 0), span)
        alpha = t / float(max(span, 1))
        return (1.0 - alpha) * start + alpha * end

    def on_step_begin(self, args, state, control, **kwargs):
        model = kwargs.get("model", None)
        if model is None:
            return
        step = state.global_step
        # Tau: keep frozen for tau_freeze_steps, then interpolate over the remaining span
        if step < self.tau_freeze_steps:
            new_tau = self.tau_start
        else:
            rem_step = step - self.tau_freeze_steps
            rem_span = max(self.total_steps - self.tau_freeze_steps, 1)
            new_tau = self._interp(self.tau_start, self.tau_end, rem_step, rem_span)

        # Aux always interpolates across total training steps
        new_aux = self._interp(self.aux_start, self.aux_end, step, self.total_steps)

        # update global config
        if hasattr(model, "config"):
            model.config.router_tau = new_tau
            model.config.router_aux_weight = new_aux

        # update all block.mlp (PolymorphicMLP must use router_tau in forward)
        for block in getattr(model, "blocks", []):
            if hasattr(block, "mlp"):
                # default: no forcing unless scheduled below
                block.mlp.router_tau = new_tau
                block.mlp.force_func = -1

        # During early warmup, occasionally force a single branch so all get gradients
        if step < self.force_warmup_steps and self.force_prob > 0.0:
            if random.random() < self.force_prob:
                for block in getattr(model, "blocks", []):
                    if hasattr(block, "mlp") and hasattr(block.mlp, "num_funcs"):
                        k = block.mlp.num_funcs
                        block.mlp.force_func = random.randint(0, max(k - 1, 0))

        if step % 1000 == 0:
            print(
                f"[router-sched] step={step} tau={new_tau:.4f} aux_w={new_aux:.5f} "
                f"freeze<= {self.tau_freeze_steps} force_p={self.force_prob:.3f} warmup<= {self.force_warmup_steps}"
            )


# ===========================
# Custom Trainer with rep_loss
# ===========================

class VeronicaTrainer(Trainer):
    def __init__(self, *args, label_smoothing: float = 0.0, rep_alpha: float = 0.0, **kwargs):
        super().__init__(*args, **kwargs)
        self.label_smoothing = float(label_smoothing)
        self.rep_alpha = float(rep_alpha)

    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
        labels = inputs.get("labels")
        if labels is None:
            raise ValueError("compute_loss called without labels")
        model_inputs = {k: v for k, v in inputs.items() if k != "labels"}

        outputs = model(**model_inputs)
        logits = outputs.logits  # [B, T, V]

        ignore_index = -100
        # SHIFT: predict x_{t+1}
        shift_logits = logits[:, :-1, :].contiguous()
        shift_labels = labels[:, 1:].contiguous()

        valid_mask = (shift_labels != ignore_index)
        safe_labels = shift_labels.clone()
        safe_labels[~valid_mask] = 0

        log_probs = F.log_softmax(shift_logits, dim=-1)  # [B, T-1, V]
        nll_full = -log_probs.gather(-1, safe_labels.unsqueeze(-1)).squeeze(-1)
        nll_loss = nll_full[valid_mask].mean()

        if self.label_smoothing > 0.0:
            smooth_full = -log_probs.mean(dim=-1)
            smooth_loss = smooth_full[valid_mask].mean()
            ce_loss = (1.0 - self.label_smoothing) * nll_loss + self.label_smoothing * smooth_loss
        else:
            ce_loss = nll_loss

        total_loss = ce_loss

        # rep_loss on x_{t+1} when x_{t+1} == x_t
        if self.rep_alpha > 0.0:
            labels_prev = labels[:, :-1]             # x_t
            labels_next = shift_labels               # x_{t+1}
            valid_prev = (labels_prev != ignore_index)
            same_mask = valid_prev & valid_mask & (labels_prev == labels_next)
            if same_mask.any():
                rep_logp = log_probs.gather(-1, safe_labels.unsqueeze(-1)).squeeze(-1)
                rep_p = rep_logp[same_mask].exp()
                total_loss = total_loss + self.rep_alpha * rep_p.mean()

        # aux_loss del router: SUBTRACT to MAXIMIZE entropy (prevent collapse)
        aux_loss = getattr(model, "_last_router_aux", None)
        if aux_loss is not None and hasattr(model, "config"):
            aux_w = float(getattr(model.config, "router_aux_weight", 0.0))
            if aux_w > 0:
                if not torch.is_tensor(aux_loss):
                    aux_loss = torch.as_tensor(aux_loss, device=logits.device, dtype=logits.dtype)
                # Subtract aux (entropy) so that minimizing loss => maximize entropy => soft router
                total_loss = total_loss - aux_w * aux_loss.clamp_min(0.0)

        return (total_loss, outputs) if return_outputs else total_loss


# ===========================
# Main
# ===========================

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", type=str, required=True)
    parser.add_argument("--dataset_paths", type=str, required=True)
    parser.add_argument("--output_dir", type=str, required=True, default="veronica-polymorphic/runs/veronica-pretrain")

    parser.add_argument(
        "--tokenizer_candidates",
        type=str,
        nargs="*",
        default=["veronica-polymorphic/tokenizer", "gpt2"],
    )
    parser.add_argument(
        "--tokenizer_out",
        type=str,
        default="veronica-polymorphic/tokenizer_vmix",
    )

    parser.add_argument("--per_device_train_batch_size", type=int, default=4)
    parser.add_argument("--per_device_eval_batch_size", type=int, default=4)
    parser.add_argument("--gradient_accumulation_steps", type=int, default=4)
    parser.add_argument("--max_steps", type=int, default=60000)
    parser.add_argument("--learning_rate", type=float, default=2e-4)
    parser.add_argument("--warmup_ratio", type=float, default=0.02)
    parser.add_argument("--weight_decay", type=float, default=0.1)
    parser.add_argument("--eval_steps", type=int, default=1000)
    parser.add_argument("--save_steps", type=int, default=1000)
    parser.add_argument("--logging_steps", type=int, default=100)
    parser.add_argument("--label_smoothing", type=float, default=0.01)
    parser.add_argument("--rep_alpha", type=float, default=0.0)
    parser.add_argument("--mask_degenerate_runs", action="store_true")
    parser.add_argument("--seed", type=int, default=42)

    parser.add_argument(
        "--resume_from",
        type=str,
        default=None,
        help="Checkpoint to resume from (e.g., .../checkpoint-22000)",
    )

    parser.add_argument(
        "--max_seq_len",
        type=int,
        default=None,
        help="Maximum window length (e.g., 512, 1024, 2048). If None, uses the full dataset sequence.",
    )

    # Schedule router
    parser.add_argument("--router_tau_start", type=float, default=1.6)
    parser.add_argument("--router_tau_end", type=float, default=1.1)
    parser.add_argument("--router_aux_start", type=float, default=0.005)
    parser.add_argument("--router_aux_end", type=float, default=0.012)
    parser.add_argument("--router_tau_freeze_steps", type=int, default=4000,
                        help="Keep tau constant for the first N steps to avoid early specialization.")
    parser.add_argument("--router_force_prob", type=float, default=0.05,
                        help="Per-step probability to force a single branch during warmup.")
    parser.add_argument("--router_force_warmup_steps", type=int, default=3000,
                        help="Apply random branch forcing only within these initial steps.")

    args = parser.parse_args()

    # Tokenizer
    tok = build_tokenizer(args.tokenizer_candidates, args.tokenizer_out)

    # Config & Model
    cfg = load_cfg_with_vocab(args.config, tok)
    cfg.router_tau = args.router_tau_start
    cfg.router_aux_weight = args.router_aux_start

    model = init_model_from_config(cfg, tok)

    # Diagnostics: verify model forward loss
    model.eval()
    with torch.no_grad():
        dummy = torch.randint(0, model.config.vocab_size, (1, 32), device=model.device)
        out = model(input_ids=dummy, labels=dummy)
        loss_model = out.loss.item()

        logits = out.logits  # [1, 32, V]
        shift_logits = logits[:, :-1, :].contiguous()
        shift_labels = dummy[:, 1:].contiguous()
        loss_manual = F.cross_entropy(
            shift_logits.view(-1, shift_logits.size(-1)),
            shift_labels.view(-1)
        ).item()

    print(f"[diag] loss_model_forward={loss_model:.4f} loss_manual_shift={loss_manual:.4f}")
    model.train()

    # Dataset
    train_ds, val_ds = load_mix_dataset(args.dataset_paths)
    collator = CausalCollator(
        tokenizer=tok,
        mask_runs=args.mask_degenerate_runs,
        max_seq_len=args.max_seq_len,
    )

    # Resume
    resume_ckpt = args.resume_from or find_latest_checkpoint(args.output_dir)
    if resume_ckpt:
        print(f"🟢 Resuming from: {resume_ckpt}")
    else:
        print("⚪ No checkpoint: training from scratch.")

    use_bf16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported()

    train_args = TrainingArguments(
        output_dir=args.output_dir,
        run_name=os.path.basename(args.output_dir.rstrip("/")),
        num_train_epochs=1_000,  # guidato da max_steps
        max_steps=args.max_steps,
        per_device_train_batch_size=args.per_device_train_batch_size,
        per_device_eval_batch_size=args.per_device_eval_batch_size,
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        learning_rate=args.learning_rate,
        warmup_ratio=args.warmup_ratio,
        weight_decay=args.weight_decay,
        lr_scheduler_type="cosine",
        logging_steps=args.logging_steps,
        eval_steps=args.eval_steps,
        save_steps=args.save_steps,
        eval_strategy="steps",  # ✅
        save_total_limit=5,
        bf16=use_bf16,
        fp16=(torch.cuda.is_available() and not use_bf16),
        gradient_checkpointing=True,
        report_to=["tensorboard"],
        dataloader_num_workers=2,
        seed=args.seed,
        label_smoothing_factor=0.0,  # smoothing gestito in compute_loss custom
        max_grad_norm=1.0,
        save_safetensors=False,
    )

    callbacks: List[TrainerCallback] = [
        RouterAndSmokeCallback(tok),
        RouterScheduleCallback(
            tau_start=args.router_tau_start,
            tau_end=args.router_tau_end,
            aux_start=args.router_aux_start,
            aux_end=args.router_aux_end,
            total_steps=args.max_steps,
            tau_freeze_steps=args.router_tau_freeze_steps,
            force_prob=args.router_force_prob,
            force_warmup_steps=args.router_force_warmup_steps,
        ),
    ]

    trainer = VeronicaTrainer(
        model=model,
        args=train_args,
        train_dataset=train_ds,
        eval_dataset=val_ds,
        tokenizer=tok,          # ✅ al posto di processing_class
        data_collator=collator,
        callbacks=callbacks,
        label_smoothing=args.label_smoothing,
        rep_alpha=args.rep_alpha,
    )

    # Sanity check: vocab/emb/head
    effective_vocab = len(tok)
    emb = model.get_input_embeddings().weight
    head = model.lm_head.weight
    assert emb.shape[0] == effective_vocab == head.shape[0], "Mismatch vocab/emb/lm_head"

    # Train
    trainer.train(resume_from_checkpoint=resume_ckpt)
    trainer.save_state()
    trainer.save_model(args.output_dir)
    tok.save_pretrained(args.output_dir)
    with open(os.path.join(args.output_dir, "config.final.json"), "w", encoding="utf-8") as f:
        json.dump(model.config.to_dict(), f, indent=2)
    print("✅ Pretraining completed/saved.")


if __name__ == "__main__":
    main()