"""
V5-SFT Asama 3: Fine-tune base model on tokenized SFT data with loss masking.

Pipeline:
  1) Base ckpt yukle (runs/tr-200m-v5/best_ckpt.pt veya HF)
  2) sft_{train,val}_{tokens,mask}.bin oku
  3) Random sliding-window sample, mask-aware CE loss
  4) Cosine LR decay, dusuk LR (base/10)
  5) Her N step val + sample generation
  6) Her N step HF push (musabc/nanogpt-tr-v5-sft)

Kullanim:
    python sft_03_train.py --base runs/tr-200m-v5/best_ckpt.pt
    python sft_03_train.py --base runs/tr-200m-v5/best_ckpt.pt --resume
    python sft_03_train.py --base runs/tr-200m-v5/best_ckpt.pt --hf-user musabc --hf-push-every 200
"""

import argparse
import json
import math
import os
import random
import sys
import time
from pathlib import Path

# Liger CE training'de logits'i in-place modifiye ettigi icin SFT'de devre disi
# Bu IMPORT'tan ONCE set edilmeli
os.environ["NANOGPT_NO_LIGER"] = "1"

import numpy as np
import torch
import torch.nn.functional as F


# ===== Hyperparams =====
BATCH_SIZE = 16            # micro-batch per forward
GRAD_ACCUM_STEPS = 4       # effective batch = 64
BLOCK_SIZE = 2048

# Lower LR than base — typical SFT
MUON_LR = 2.2e-3 / 10
ADAMW_LR = 3.5e-4 / 10
MUON_MOMENTUM = 0.95
WEIGHT_DECAY = 0.1

WARMUP_STEPS = 50
NUM_EPOCHS = 3
EVAL_EVERY = 100
HF_PUSH_EVERY = 200
SAMPLE_EVERY = 200
MAX_GRAD_NORM = 1.0

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = torch.bfloat16


def log(msg):
    print(msg, flush=True)


# ===== Data =====
class SFTPackedDataset:
    """Packed bin'den random sliding window sample. Mask-aware."""
    def __init__(self, tokens_path: Path, mask_path: Path, block_size: int,
                  batch_size: int):
        self.tokens = np.memmap(tokens_path, dtype=np.uint16, mode="r")
        self.mask = np.memmap(mask_path, dtype=np.uint8, mode="r")
        assert len(self.tokens) == len(self.mask), \
            f"Token/mask uzunluk uyusmazlik: {len(self.tokens)} vs {len(self.mask)}"
        self.block_size = block_size
        self.batch_size = batch_size
        self.length = len(self.tokens)
        log(f"  Dataset: {self.length:,} token  "
            f"(loss tokens: {int(self.mask.sum()):,})")

    def get_batch(self, rng: np.random.Generator):
        """Random sliding window batches. Tek window'da mask hep 0 olabilir,
        bu durumda loss zero — kabul edilir (random pure-prompt slice)."""
        bs, T = self.batch_size, self.block_size
        # T+1 token gerekli (next-token prediction)
        ix = rng.integers(0, self.length - T - 1, size=bs)
        x = np.zeros((bs, T), dtype=np.int64)
        y = np.zeros((bs, T), dtype=np.int64)
        m = np.zeros((bs, T), dtype=np.float32)
        for i, start in enumerate(ix):
            x[i] = self.tokens[start:start+T].astype(np.int64)
            y[i] = self.tokens[start+1:start+T+1].astype(np.int64)
            # Mask y'ye gore — y[t] = next-token, mask[start+1:start+T+1] uygulanir
            m[i] = self.mask[start+1:start+T+1].astype(np.float32)
        x = torch.from_numpy(x).to(DEVICE, non_blocking=True)
        y = torch.from_numpy(y).to(DEVICE, non_blocking=True)
        m = torch.from_numpy(m).to(DEVICE, non_blocking=True)
        return x, y, m


# ===== Loss =====
def masked_ce_loss(logits: torch.Tensor, targets: torch.Tensor,
                    mask: torch.Tensor, softcap: float = 30.0) -> torch.Tensor:
    """Mask-aware cross-entropy. mask=0 olan tokenlar loss'a katilmaz."""
    # logits: (B, T, V)  targets: (B, T)  mask: (B, T) float
    B, T, V = logits.shape
    if softcap > 0:
        logits = softcap * torch.tanh(logits / softcap)
    # Flatten
    flat_logits = logits.view(-1, V)
    flat_targets = targets.view(-1)
    flat_mask = mask.view(-1)

    # Per-token loss (no reduction)
    losses = F.cross_entropy(flat_logits, flat_targets, reduction="none")
    # Mask uygula
    masked_losses = losses * flat_mask
    # Ortalama (sadece loss var olan tokenlar uzerinden)
    total_mask = flat_mask.sum()
    if total_mask < 1:
        return torch.zeros((), device=logits.device, requires_grad=True)
    return masked_losses.sum() / total_mask


# ===== LR Schedule =====
def get_lr(step: int, total_steps: int, warmup_steps: int,
            base_lr: float, min_lr_ratio: float = 0.01) -> float:
    if step < warmup_steps:
        return base_lr * step / max(warmup_steps, 1)
    progress = (step - warmup_steps) / max(total_steps - warmup_steps, 1)
    progress = min(1.0, progress)
    cosine = 0.5 * (1 + math.cos(math.pi * progress))
    return base_lr * (min_lr_ratio + (1 - min_lr_ratio) * cosine)


# ===== Optimizer setup (Muon + AdamW split) =====
def setup_optimizers(model):
    try:
        from torch.optim import Muon as TorchMuon
        has_native = True
    except ImportError:
        from muon import Muon as TorchMuon
        has_native = False

    muon_params = []
    adam_params = []
    for name, p in model.named_parameters():
        if not p.requires_grad:
            continue
        # Embedding ve lm_head AdamW'da
        if "wte" in name or "lm_head" in name:
            adam_params.append(p)
        # 1D parametreler AdamW'da
        elif p.ndim < 2:
            adam_params.append(p)
        # 2D matrisler (linear/attention) Muon'da
        else:
            muon_params.append(p)

    log(f"  Muon params:  {sum(p.numel() for p in muon_params)/1e6:.2f}M  "
        f"({len(muon_params)} tensor)")
    log(f"  AdamW params: {sum(p.numel() for p in adam_params)/1e6:.2f}M  "
        f"({len(adam_params)} tensor)")

    muon_kwargs = dict(lr=MUON_LR, momentum=MUON_MOMENTUM, nesterov=True, ns_steps=3)
    if has_native:
        muon_kwargs["weight_decay"] = 0.0
    opt_muon = TorchMuon(muon_params, **muon_kwargs)
    opt_adam = torch.optim.AdamW(adam_params, lr=ADAMW_LR,
                                  betas=(0.9, 0.95),
                                  weight_decay=WEIGHT_DECAY)
    return opt_muon, opt_adam


# ===== Sample =====
def sample_test(model, tokenizer, prompt_text: str, max_new: int = 200):
    model.eval()
    from tokenizers import Tokenizer
    if isinstance(tokenizer, str):
        tokenizer = Tokenizer.from_file(tokenizer)
    ids = tokenizer.encode(prompt_text).ids
    x = torch.tensor([ids], dtype=torch.long, device=DEVICE)
    with torch.no_grad():
        with torch.amp.autocast(device_type="cuda", dtype=DTYPE):
            out = model.generate(x, max_new_tokens=max_new,
                                  temperature=0.3, top_k=40,
                                  repetition_penalty=1.0,
                                  no_repeat_ngram_size=0,
                                  eos_token_id=0)  # <|endoftext|>
    text = tokenizer.decode(out[0].tolist())
    model.train()
    return text


# ===== Main =====
def main():
    global BATCH_SIZE, GRAD_ACCUM_STEPS
    parser = argparse.ArgumentParser()
    parser.add_argument("--base", type=str,
                        default="runs/tr-200m-v5/best_ckpt.pt",
                        help="Base ckpt yolu")
    parser.add_argument("--data-dir", type=str, default="data/sft")
    parser.add_argument("--tokenizer", type=str,
                        default="data/tokenizer-tr-v5.json")
    parser.add_argument("--out-dir", type=str,
                        default="runs/tr-200m-v5-sft")
    parser.add_argument("--resume", action="store_true")
    parser.add_argument("--batch", type=int, default=BATCH_SIZE)
    parser.add_argument("--grad-accum", type=int, default=GRAD_ACCUM_STEPS)
    parser.add_argument("--epochs", type=int, default=NUM_EPOCHS)
    parser.add_argument("--hf-user", type=str, default=None,
                        help="HF user, ornek: musabc. Bos -> push yapilmaz")
    parser.add_argument("--hf-push-every", type=int, default=HF_PUSH_EVERY)
    parser.add_argument("--no-compile", action="store_true")
    parser.add_argument("--compile-mode", type=str,
                        default="max-autotune-no-cudagraphs")
    args = parser.parse_args()

    BATCH_SIZE = args.batch
    GRAD_ACCUM_STEPS = args.grad_accum

    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    log(f"Device: {DEVICE}  dtype: {DTYPE}")
    log(f"Batch: {BATCH_SIZE} x grad_accum {GRAD_ACCUM_STEPS} "
        f"= effective {BATCH_SIZE*GRAD_ACCUM_STEPS}")

    # ===== Data =====
    data_dir = Path(args.data_dir)
    log(f"\nData: {data_dir}")
    train_ds = SFTPackedDataset(
        data_dir / "sft_train_tokens.bin",
        data_dir / "sft_train_mask.bin",
        BLOCK_SIZE, BATCH_SIZE
    )
    val_ds = SFTPackedDataset(
        data_dir / "sft_val_tokens.bin",
        data_dir / "sft_val_mask.bin",
        BLOCK_SIZE, BATCH_SIZE
    )

    # Toplam step = (train_tokens / tokens_per_step) * epochs
    tokens_per_step = BATCH_SIZE * GRAD_ACCUM_STEPS * BLOCK_SIZE
    total_loss_tokens = int(train_ds.mask.sum())
    # Bir "epoch" yaklasik loss-token bazli
    steps_per_epoch = max(1, total_loss_tokens // (BATCH_SIZE * GRAD_ACCUM_STEPS *
                                                     BLOCK_SIZE // 10))
    # Daha basit: pack uzunlugu / tokens_per_step
    steps_per_epoch = max(1, train_ds.length // tokens_per_step)
    total_steps = steps_per_epoch * args.epochs
    log(f"  Steps/epoch: {steps_per_epoch}  | total steps: {total_steps}")
    log(f"  Tokens/step: {tokens_per_step:,}")

    # ===== Model =====
    log(f"\nModel yukleniyor: {args.base}")
    from model_v5 import GPTV5, GPTConfigV5
    cfg = GPTConfigV5()
    model = GPTV5(cfg).to(DEVICE)
    ckpt = torch.load(args.base, map_location=DEVICE, weights_only=False)
    if "model" in ckpt:
        sd = ckpt["model"]
        log(f"  Base step: {ckpt.get('step', '?')}  best_val: "
            f"{ckpt.get('best_val', '?')}")
    else:
        sd = ckpt

    # _orig_mod. prefix'i kaldir (compile ile kaydedilmis ise)
    has_prefix = any(k.startswith("_orig_mod.") for k in sd.keys())
    if has_prefix:
        log(f"  ! _orig_mod. prefix tespit edildi, kaldiriliyor")
        sd = {k.replace("_orig_mod.", "", 1): v for k, v in sd.items()}

    missing, unexpected = model.load_state_dict(sd, strict=False)
    log(f"  Yuklendi: {len(sd) - len(unexpected)}/{len(sd)} key")
    if missing:
        log(f"  ! Missing keys ({len(missing)}): {missing[:3]}...")
    if unexpected:
        log(f"  ! Unexpected keys ({len(unexpected)}): {unexpected[:3]}...")
    log(f"  Params: {sum(p.numel() for p in model.parameters())/1e6:.1f}M")

    # Sanity check — random tokens uzerinde forward, ortalama top-1 olasilik
    with torch.no_grad():
        with torch.amp.autocast(device_type="cuda", dtype=DTYPE):
            test_x = torch.randint(0, cfg.vocab_size, (2, 64), device=DEVICE)
            test_logits, _ = model(test_x, test_x)
            probs = F.softmax(test_logits[0, 0].float(), dim=-1)
            top1_p = probs.max().item()
            entropy = -(probs * torch.log(probs + 1e-12)).sum().item()
    log(f"  Sanity: top1_p={top1_p:.4f}, entropy={entropy:.3f} "
        f"(uniform ~= {math.log(cfg.vocab_size):.2f}, trained <~ 5)")

    # Resume?
    start_step = 0
    sft_latest = out_dir / "sft_latest.pt"
    if args.resume and sft_latest.exists():
        log(f"\nSFT resume: {sft_latest}")
        sft_ckpt = torch.load(sft_latest, map_location=DEVICE, weights_only=False)
        model.load_state_dict(sft_ckpt["model"])
        start_step = sft_ckpt.get("step", 0)
        log(f"  Resume step: {start_step}")

    # ===== Optimizers =====
    log(f"\nOptimizer setup:")
    opt_muon, opt_adam = setup_optimizers(model)

    # Compile
    if not args.no_compile:
        log(f"\ntorch.compile (mode={args.compile_mode})...")
        model = torch.compile(model, mode=args.compile_mode, dynamic=False)

    torch.backends.cudnn.benchmark = True
    torch.set_float32_matmul_precision("high")

    # ===== Training =====
    log(f"\nSFT egitim basliyor: {start_step} -> {total_steps}\n")
    rng = np.random.default_rng(42 + start_step)
    t_start = time.time()
    accum_loss = 0.0
    accum_count = 0
    best_val_loss = float("inf")
    log_f = open(out_dir / "sft_train.log", "a", encoding="utf-8")

    model.train()
    for step in range(start_step, total_steps):
        # LR update
        muon_lr = get_lr(step, total_steps, WARMUP_STEPS, MUON_LR)
        adam_lr = get_lr(step, total_steps, WARMUP_STEPS, ADAMW_LR)
        for g in opt_muon.param_groups:
            g["lr"] = muon_lr
        for g in opt_adam.param_groups:
            g["lr"] = adam_lr

        opt_muon.zero_grad(set_to_none=True)
        opt_adam.zero_grad(set_to_none=True)

        step_loss = 0.0
        valid_micro = 0
        for _ in range(GRAD_ACCUM_STEPS):
            x, y, m = train_ds.get_batch(rng)
            with torch.amp.autocast(device_type="cuda", dtype=DTYPE):
                # targets=y geciyoruz → full (B,T,V) logits + softcap uygulanmis
                # Donen loss'u atip kendi mask'li CE'mizi hesapliyoruz
                logits, _ = model(x, y)
                loss = masked_ce_loss(logits, y, m, softcap=0.0)
            if torch.isnan(loss) or loss.item() == 0.0:
                continue
            (loss / GRAD_ACCUM_STEPS).backward()
            step_loss += loss.item()
            valid_micro += 1

        if valid_micro == 0:
            continue
        step_loss /= valid_micro

        # Grad clip
        torch.nn.utils.clip_grad_norm_(
            [p for g in opt_muon.param_groups for p in g["params"]] +
            [p for g in opt_adam.param_groups for p in g["params"]],
            MAX_GRAD_NORM,
        )
        opt_muon.step()
        opt_adam.step()

        accum_loss += step_loss
        accum_count += 1

        # Log
        if (step + 1) % 10 == 0:
            avg = accum_loss / accum_count
            elapsed_min = (time.time() - t_start) / 60
            tok_per_sec = (step - start_step + 1) * tokens_per_step / max(time.time() - t_start, 1)
            line = (f"step {step+1:>5} | loss {avg:.4f} | "
                    f"muon {muon_lr:.2e} adam {adam_lr:.2e} | "
                    f"{tok_per_sec/1e3:.0f}K tok/s | {elapsed_min:.1f}m")
            log(line)
            log_f.write(line + "\n"); log_f.flush()
            accum_loss = 0.0
            accum_count = 0

        # Eval
        if (step + 1) % EVAL_EVERY == 0:
            model.eval()
            val_losses = []
            val_rng = np.random.default_rng(0)
            with torch.no_grad():
                with torch.amp.autocast(device_type="cuda", dtype=DTYPE):
                    for _ in range(20):
                        x, y, m = val_ds.get_batch(val_rng)
                        logits, _ = model(x, y)
                        vl = masked_ce_loss(logits, y, m, softcap=0.0)
                        if not torch.isnan(vl) and vl.item() > 0:
                            val_losses.append(vl.item())
            model.train()
            avg_val = sum(val_losses) / max(len(val_losses), 1)
            line = f"  >>> EVAL: val {avg_val:.4f}  ({len(val_losses)} batch)"
            log(line); log_f.write(line + "\n"); log_f.flush()
            if avg_val < best_val_loss and avg_val > 0:
                best_val_loss = avg_val
                # Save best — model._orig_mod (compile sonra)
                mdl = model._orig_mod if hasattr(model, "_orig_mod") else model
                torch.save({
                    "model": mdl.state_dict(),
                    "step": step + 1,
                    "best_val": best_val_loss,
                }, out_dir / "sft_best.pt")
                line = f"  >>> BEST kaydedildi (val {best_val_loss:.4f})"
                log(line); log_f.write(line + "\n"); log_f.flush()

        # Sample
        if (step + 1) % SAMPLE_EVERY == 0:
            try:
                mdl = model._orig_mod if hasattr(model, "_orig_mod") else model
                prompts = [
                    # Tarif (training'de turkish_recipes 4K)
                    "### Kullanici:\nMercimek corbasi tarifi ver.\n### Asistan:\n",
                    # Factual short (turkish_exam, knowledge)
                    "### Kullanici:\nTurkiye'nin baskenti neresidir?\n### Asistan:\n",
                    # Aciklayici (general)
                    "### Kullanici:\nBir e-mail nasil yazilir? Kisa anlat.\n### Asistan:\n",
                    # Simple math (gsm8k/metamath)
                    "### Kullanici:\nAhmet'in 5 elmasi var, 2 tane yer. Kac kaldi?\n### Asistan:\n",
                ]
                for p in prompts:
                    out = sample_test(mdl, args.tokenizer, p, max_new=120)
                    # Sadece asistan kismini bul
                    if "### Asistan:" in out:
                        asst = out.split("### Asistan:", 1)[1].strip()
                        asst = asst.split("<|endoftext|>", 1)[0].strip()
                    else:
                        asst = out
                    line = f"  [sample] {asst[:220]!r}"
                    log(line); log_f.write(line + "\n"); log_f.flush()
            except Exception as e:
                log(f"  [sample err] {e}")

        # Checkpoint + HF push
        if (step + 1) % 100 == 0:
            mdl = model._orig_mod if hasattr(model, "_orig_mod") else model
            torch.save({
                "model": mdl.state_dict(),
                "step": step + 1,
                "best_val": best_val_loss,
            }, out_dir / "sft_latest.pt")

        if args.hf_user and (step + 1) % args.hf_push_every == 0:
            try:
                from huggingface_hub import HfApi
                api = HfApi()
                repo = f"{args.hf_user}/nanogpt-tr-v5-sft"
                api.create_repo(repo, repo_type="model", exist_ok=True, private=False)
                for fname in ["sft_latest.pt", "sft_best.pt", "sft_train.log"]:
                    p = out_dir / fname
                    if p.exists():
                        api.upload_file(
                            path_or_fileobj=str(p),
                            path_in_repo=fname,
                            repo_id=repo,
                            repo_type="model",
                        )
                log(f"  >>> HF push: {repo} (step {step+1})")
            except Exception as e:
                log(f"  [hf push err] {e}")

    log_f.close()
    log(f"\nSFT tamamlandi. Best val: {best_val_loss:.4f}")


if __name__ == "__main__":
    main()