"""
Faz 8 — DPO (Direct Preference Optimization) — SmartCore V1. (Ufuk 1 / Adım 2)

v1-instruct-rag'i (sft_rag3/epoch_2) tercih çiftleriyle hizalar: chosen yanıt > rejected.
⚠️ 177M'de marjinal; MC benchmark'ı değiştirmez, üretim tonu/formatını iyileştirir.

Yöntem: policy (eğitilen) + reference (donmuş kopya). DPO loss:
  L = -log σ( β · [ (logπ_chosen - logπ_rejected) - (logπref_chosen - logπref_rejected) ] )
Loss YALNIZ yanıt tokenlerinde (prompt -100 maskeli). Model tanımı faz6_sft ile birebir.

Ortam: Colab GPU + mamba-og fork (wheel). Yerelde test edilebilir: build_prompt/encode/collate/dpo_loss.
Kullanım:
  HF_TOKEN=hf_xxx python faz8_dpo.py --data dpo.jsonl --base sft_rag3/epoch_2/ckpt.pt \
      --beta 0.1 --lr 5e-7 --epochs 1 --micro_batch 4 --grad_accum 8 --max_len 1024
"""
import os, sys, json, math, time, random, argparse
import torch, torch.nn as nn, torch.nn.functional as F
from functools import partial

try:
    from mamba_ssm.modules.block import Block
    from mamba_ssm.modules.mamba3 import Mamba3
    from mamba_ssm.modules.mlp import GatedMLP
    from mamba_ssm.ops.triton.layer_norm import RMSNorm
    FORK = True
except Exception:
    Block = Mamba3 = GatedMLP = RMSNorm = None
    FORK = False


# ───────────── model (faz6_sft.py ile birebir) ─────────────
def _rms(x, w, eps=1e-5):
    return (x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)) * w


def _rot_half(x):
    a, b = x.chunk(2, -1)
    return torch.cat((-b, a), -1)


class GQAMixer(nn.Module):
    def __init__(self, dim, n_heads=12, n_kv=3, base=10000.0, layer_idx=None, device=None, dtype=None):
        super().__init__()
        self.nh, self.nkv, self.hd = n_heads, n_kv, dim // n_heads
        self.rep = n_heads // n_kv
        fk = {"device": device, "dtype": dtype}
        self.q_proj = nn.Linear(dim, n_heads * self.hd, bias=False, **fk)
        self.k_proj = nn.Linear(dim, n_kv * self.hd, bias=False, **fk)
        self.v_proj = nn.Linear(dim, n_kv * self.hd, bias=False, **fk)
        self.out_proj = nn.Linear(n_heads * self.hd, dim, bias=False, **fk)
        self.qn = nn.Parameter(torch.ones(self.hd, **fk))
        self.kn = nn.Parameter(torch.ones(self.hd, **fk))
        self.register_buffer(
            "inv", 1.0 / (base ** (torch.arange(0, self.hd, 2, device=device).float() / self.hd)),
            persistent=False)

    def _rope(self, x, T):
        f = torch.outer(torch.arange(T, device=x.device, dtype=torch.float32), self.inv)
        e = torch.cat((f, f), -1)
        return (x * e.cos()[None, None] + _rot_half(x) * e.sin()[None, None]).to(x.dtype)

    def forward(self, x, **kw):
        B, T, _ = x.shape
        q = self.q_proj(x).view(B, T, self.nh, self.hd).transpose(1, 2)
        k = self.k_proj(x).view(B, T, self.nkv, self.hd).transpose(1, 2)
        v = self.v_proj(x).view(B, T, self.nkv, self.hd).transpose(1, 2)
        q = _rms(q.float(), self.qn.float()).to(x.dtype)
        k = _rms(k.float(), self.kn.float()).to(x.dtype)
        q, k = self._rope(q, T), self._rope(k, T)
        k = k.repeat_interleave(self.rep, 1)
        v = v.repeat_interleave(self.rep, 1)
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        return self.out_proj(y.transpose(1, 2).contiguous().view(B, T, -1))


class HybridLM(nn.Module):
    def __init__(self, cfg, device=None, dtype=None):
        super().__init__()
        self.cfg = cfg
        self.vocab = cfg["vocab_size"]
        self.scaled_embed = cfg.get("scaled_embed", False)
        d = cfg["d_model"]
        self.embedding = nn.Embedding(self.vocab, d, device=device, dtype=dtype)
        self.layers = nn.ModuleList()
        for i in range(cfg["n_layers"]):
            is_attn = ((i + 1) % cfg["attn_every"] == 0) and i != 0 and i != cfg["n_layers"] - 1
            fk = {"device": device, "dtype": dtype}
            if is_attn:
                mixer_cls = partial(GQAMixer, n_heads=cfg["n_heads"], n_kv=cfg["n_kv_heads"], layer_idx=i, **fk)
            else:
                ssm = dict(d_state=cfg["d_state"], expand=cfg["expand"], headdim=cfg["head_dim"],
                           ngroups=cfg["ngroups"], rope_fraction=cfg["rope_fraction"],
                           is_outproj_norm=False, is_mimo=cfg["is_mimo"], mimo_rank=cfg["mimo_rank"],
                           chunk_size=cfg["chunk_size"])
                mixer_cls = partial(Mamba3, layer_idx=i, **ssm, **fk)
            blk = Block(d, mixer_cls,
                        partial(GatedMLP, hidden_features=cfg["d_intermediate"], out_features=d, **fk),
                        norm_cls=partial(RMSNorm, eps=1e-5, **fk), fused_add_norm=True, residual_in_fp32=True)
            blk.layer_idx = i
            self.layers.append(blk)
        self.norm_f = RMSNorm(d, eps=1e-5, device=device, dtype=dtype)
        self.lm_head = nn.Linear(d, self.vocab, bias=False, device=device, dtype=dtype)
        self.lm_head.weight = self.embedding.weight

    def forward(self, ids):
        h = self.embedding(ids)
        if self.scaled_embed:
            h = h * (self.cfg["d_model"] ** 0.5)
        res = None
        for l in self.layers:
            h, res = l(h, res)
        h = self.norm_f((h + res) if res is not None else h)
        return self.lm_head(h.to(self.lm_head.weight.dtype))


# ───────────── veri + DPO (saf-python: yerelde test edilebilir) ─────────────
def build_prompt(instr, inp=""):
    instr = instr.strip(); inp = (inp or "").strip()
    if inp:
        return f"### Talimat:\n{instr}\n\n### Girdi:\n{inp}\n\n### Yanıt:\n"
    return f"### Talimat:\n{instr}\n\n### Yanıt:\n"


def encode_resp(sp, instr, resp, max_len, eos_id):
    """(ids, labels) — prompt -100 maskeli, yanıt+eos. Uzunsa SOLDAN kırp (yanıt korunur)."""
    p_ids = sp.encode(build_prompt(instr), out_type=int)
    r_ids = sp.encode(resp.strip(), out_type=int) + [eos_id]
    ids = p_ids + r_ids
    labels = [-100] * len(p_ids) + r_ids
    if len(ids) > max_len:
        ids, labels = ids[-max_len:], labels[-max_len:]
    return ids, labels


def collate(batch, pad_id):
    """Sağdan pad. batch = list of (ids, labels)."""
    maxlen = max(len(ids) for ids, _ in batch)
    B = len(batch)
    input_ids = torch.full((B, maxlen), pad_id, dtype=torch.long)
    labels = torch.full((B, maxlen), -100, dtype=torch.long)
    for i, (ids, lab) in enumerate(batch):
        input_ids[i, :len(ids)] = torch.tensor(ids, dtype=torch.long)
        labels[i, :len(lab)] = torch.tensor(lab, dtype=torch.long)
    return input_ids, labels


def seq_logp(model, input_ids, labels, length_norm=False):
    """Dizi başına yanıt-token log-olasılık (prompt -100 maskeli). -> [B]
    length_norm=True → token sayısına böl (uzunluk yanlılığını engeller; margin O(1))."""
    logits = model(input_ids)
    logp = F.log_softmax(logits[:, :-1].float(), dim=-1)
    tgt = labels[:, 1:]
    mask = (tgt != -100)
    tok = torch.gather(logp, -1, tgt.clamp(min=0).unsqueeze(-1)).squeeze(-1)
    s = (tok * mask).sum(-1)
    if length_norm:
        s = s / mask.sum(-1).clamp(min=1)
    return s


def dpo_loss(pol_ch, pol_rj, ref_ch, ref_rj, beta):
    """DPO: -log σ(β·((πc-πr)-(refc-refr))). Döner: (loss, tercih_doğruluğu, ödül_marjı)."""
    logits = beta * ((pol_ch - pol_rj) - (ref_ch - ref_rj))
    loss = -F.logsigmoid(logits).mean()
    acc = (logits > 0).float().mean()
    margin = (pol_ch - pol_rj).mean()
    return loss, acc.item(), margin.item()


def lr_at(step, total, peak, warmup, floor_ratio=0.1):
    if step < warmup:
        return peak * (step + 1) / max(1, warmup)
    prog = (step - warmup) / max(1, total - warmup)
    return floor_ratio * peak + 0.5 * (1 - floor_ratio) * peak * (1 + math.cos(math.pi * prog))


# ───────────── yükleme (Colab) ─────────────
def load_tok(token):
    import sentencepiece as spm
    from huggingface_hub import hf_hub_download
    p = hf_hub_download("kdirgul/smartcore-v1", "tokenizer/tokenizer.model", repo_type="model", token=token)
    return spm.SentencePieceProcessor(model_file=p)


def resolve_ckpt(spec, token):
    if os.path.exists(spec):
        return spec
    from huggingface_hub import hf_hub_download
    print(f"[base] HF: {spec}", flush=True)
    return hf_hub_download("kdirgul/smartcore-v1", spec, repo_type="model", token=token)


def load_pairs(path):
    rows = []
    with open(path, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            ex = json.loads(line)
            if ex.get("instruction") and ex.get("chosen") and ex.get("rejected"):
                rows.append(ex)
    return rows


# ───────────── eğitim ─────────────
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", required=True, help="DPO JSONL: {instruction, chosen, rejected}")
    ap.add_argument("--base", default="sft_rag3/epoch_2/ckpt.pt", help="başlangıç ckpt (policy=reference)")
    ap.add_argument("--beta", type=float, default=0.1)
    ap.add_argument("--lr", type=float, default=5e-7)
    ap.add_argument("--epochs", type=int, default=1)
    ap.add_argument("--micro_batch", type=int, default=4)
    ap.add_argument("--grad_accum", type=int, default=8)
    ap.add_argument("--max_len", type=int, default=1024)
    ap.add_argument("--length_norm", action="store_true",
                    help="seq_logp'i token sayısına böl (uzunluk yanlılığını öldürür; margin O(1) → beta'yı yükselt)")
    ap.add_argument("--warmup_ratio", type=float, default=0.1)
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--save_repo", default="kdirgul/smartcore-v1")
    ap.add_argument("--save_subdir", default="dpo")
    ap.add_argument("--save_dir", default="/content/dpo")
    ap.add_argument("--log_every", type=int, default=10)
    args = ap.parse_args()

    assert FORK, "mamba-og fork yok — önce wheel kurulum hücresini çalıştır."
    assert torch.cuda.is_available(), "CUDA yok (Colab GPU gerekir)."
    dev = "cuda"
    torch.manual_seed(args.seed); random.seed(args.seed)
    torch.set_float32_matmul_precision("high")
    token = os.environ.get("HF_TOKEN")
    try:
        from huggingface_hub import get_token
        token = token or get_token()
    except Exception:
        pass

    sp = load_tok(token); eos_id = sp.eos_id(); pad_id = max(sp.pad_id(), 0)
    st = torch.load(resolve_ckpt(args.base, token), map_location="cpu")
    cfg = st["cfg"]
    policy = HybridLM(cfg, device=dev, dtype=torch.bfloat16)
    policy.load_state_dict(st["model"], strict=False); policy.train()
    reference = HybridLM(cfg, device=dev, dtype=torch.bfloat16)
    reference.load_state_dict(st["model"], strict=False); reference.eval()
    for p in reference.parameters():
        p.requires_grad_(False)
    print(f"[model] policy+reference yüklendi | {'MIMO' if cfg.get('is_mimo') else 'SISO'} | "
          f"base={args.base}", flush=True)

    rows = load_pairs(args.data)
    random.shuffle(rows)
    enc = [(encode_resp(sp, r["instruction"], r["chosen"], args.max_len, eos_id),
            encode_resp(sp, r["instruction"], r["rejected"], args.max_len, eos_id)) for r in rows]
    print(f"[veri] {len(enc)} tercih çifti", flush=True)

    opt = torch.optim.AdamW([p for p in policy.parameters() if p.requires_grad],
                            lr=args.lr, betas=(0.9, 0.95), eps=1e-8, weight_decay=0.0, fused=True)
    eff = args.micro_batch * args.grad_accum
    steps_per_epoch = max(1, len(enc) // eff)
    total_steps = steps_per_epoch * args.epochs
    warmup = max(1, int(total_steps * args.warmup_ratio))
    print(f"[plan] {len(enc)} çift | eff_batch={eff} | {steps_per_epoch} step/epoch | "
          f"{total_steps} step | warmup {warmup} | lr {args.lr} | beta {args.beta}", flush=True)

    os.makedirs(args.save_dir, exist_ok=True)
    gstep = 0; t0 = time.perf_counter()
    for epoch in range(args.epochs):
        random.shuffle(enc)
        for s in range(steps_per_epoch):
            opt.zero_grad(set_to_none=True)
            lr = lr_at(gstep, total_steps, args.lr, warmup)
            for g in opt.param_groups:
                g["lr"] = lr
            la, acca, marga = 0.0, 0.0, 0.0
            base = s * eff
            for a in range(args.grad_accum):
                chunk = enc[base + a * args.micro_batch: base + (a + 1) * args.micro_batch]
                if not chunk:
                    continue
                ch_ids, ch_lab = collate([c for c, _ in chunk], pad_id)
                rj_ids, rj_lab = collate([r for _, r in chunk], pad_id)
                ch_ids, ch_lab = ch_ids.to(dev), ch_lab.to(dev)
                rj_ids, rj_lab = rj_ids.to(dev), rj_lab.to(dev)
                with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
                    pol_ch = seq_logp(policy, ch_ids, ch_lab, args.length_norm)
                    pol_rj = seq_logp(policy, rj_ids, rj_lab, args.length_norm)
                    with torch.no_grad():
                        ref_ch = seq_logp(reference, ch_ids, ch_lab, args.length_norm)
                        ref_rj = seq_logp(reference, rj_ids, rj_lab, args.length_norm)
                loss, acc, marg = dpo_loss(pol_ch, pol_rj, ref_ch, ref_rj, args.beta)
                (loss / args.grad_accum).backward()
                la += loss.item() / args.grad_accum; acca += acc / args.grad_accum; marga += marg / args.grad_accum
            gn = torch.nn.utils.clip_grad_norm_(policy.parameters(), 1.0)
            opt.step(); gstep += 1
            if gstep % args.log_every == 0:
                tps = (gstep * eff) / (time.perf_counter() - t0)
                print(f"e{epoch} step {gstep}/{total_steps} | loss {la:.4f} | acc {acca:.2f} | "
                      f"margin {marga:+.2f} | gnorm {gn:5.2f} | lr {lr:.2e} | {tps:.1f} pair/s", flush=True)
        # epoch sonu kaydet + push
        d = os.path.join(args.save_dir, f"epoch_{epoch}")
        os.makedirs(d, exist_ok=True)
        torch.save({"model": policy.state_dict(), "cfg": cfg, "epoch": epoch, "sft": True, "dpo": True},
                   os.path.join(d, "ckpt.pt"))
        if token and args.save_repo:
            try:
                from huggingface_hub import HfApi
                HfApi(token=token).upload_folder(folder_path=d, repo_id=args.save_repo, repo_type="model",
                                                 path_in_repo=f"{args.save_subdir}/epoch_{epoch}",
                                                 commit_message=f"{args.save_subdir} epoch {epoch}")
                print(f"[ckpt] HF push OK {args.save_subdir}/epoch_{epoch}", flush=True)
            except Exception as e:
                print(f"[ckpt] push HATA: {repr(e)[:160]}", flush=True)
    print("[bitti] DPO tamamlandı.", flush=True)


if __name__ == "__main__":
    main()