"""
Faz 8 v1.5b — ORPO (Odds Ratio Preference Optimization) — referans-modelsiz tek-aşama.

v1 faz8_dpo.py (policy+donmuş-reference, DPO) KORUNUR. Bu ORPO: DPO 177M'de NÖTRALDI → ORPO ref-free.
LOSS (tek-aşama):  L = NLL(chosen) − β · log σ(log_odds)
  log_odds = (logp_ch − logp_rj) − (log(1−e^logp_ch) − log(1−e^logp_rj))
  logp_* = yanıt tokenlerinin UZUNLUK-NORMALİZE ortalama log-olasılığı (ORPO şart; odds anlamlı olsun).
Referans-model YOK → DPO'nun 2 ekstra forward'ı kalkar (bellek/hız avantajı). Sadece policy.
prompt = faz8_prep_v15b chat-render ({prompt} '### Asistan:\\n' ile biter); loss YALNIZ yanıt tokenlerinde.

Kaynak veri: faz8_prep_v15b.py → {prompt, chosen, rejected, lang}.
Base: v1.5b RAG-SFT'li model (sft_rag/epoch_*/ckpt.pt). Model tanımı faz6_sft ile birebir (gömülü).

Ortam: Colab/GCP GPU + resmî mamba-ssm. Yerelde test edilebilir: encode_resp/collate/seq_logp/orpo_loss/lr_at.
Kullanım:
  HF_TOKEN=hf_xxx python faz8_orpo.py --data orpo_v15b.jsonl --base sft_rag/epoch_2/ckpt.pt \
      --beta 0.1 --lr 8e-6 --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 + ORPO (saf-python: yerelde test edilebilir) ─────────────
def encode_resp(sp, prompt, resp, max_len, eos_id):
    """(ids, labels) — prompt (zaten chat-render, '### Asistan:\\n' ile biter) -100 maskeli, yanıt+eos öğrenilir.
    Uzunsa SOLDAN kırp (yanıt korunur)."""
    p_ids = sp.encode(prompt, 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):
    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=True):
    """Yanıt-token log-olasılık toplamı (prompt -100 maskeli). ORPO: length_norm=True (avg/token). -> [B]"""
    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 orpo_loss(logp_ch, logp_rj, beta):
    """ORPO (ref-free): NLL(chosen) − β·logσ(log_odds). logp_* = uzunluk-normalize avg log-prob (<0).
    log_odds = (lc−lr) − (log(1−e^lc) − log(1−e^lr)). Döner (loss, tercih_acc, margin, nll)."""
    lc = logp_ch.clamp(max=-1e-4); lr = logp_rj.clamp(max=-1e-4)   # e^logp<1 (log(0) koru; avg-logp zaten <0)
    log_odds = (lc - lr) - (torch.log1p(-torch.exp(lc)) - torch.log1p(-torch.exp(lr)))
    ratio = F.logsigmoid(log_odds)          # ≤ 0
    nll = -logp_ch                          # SFT terimi (chosen mean NLL)
    loss = (nll - beta * ratio).mean()
    acc = (log_odds > 0).float().mean()
    margin = (logp_ch - logp_rj).mean()
    return loss, acc.item(), margin.item(), nll.mean().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 ─────────────
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("prompt") 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="ORPO JSONL: {prompt, chosen, rejected}")
    ap.add_argument("--base", default="sft_rag/epoch_2/ckpt.pt", help="başlangıç ckpt (v1.5b RAG-SFT'li)")
    ap.add_argument("--beta", type=float, default=0.1, help="odds-ratio ağırlığı λ (ORPO ~0.1-0.5)")
    ap.add_argument("--lr", type=float, default=8e-6, help="ORPO SFT-benzeri → DPO'dan yüksek (~5e-6..1e-5)")
    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("--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="orpo")
    ap.add_argument("--save_dir", default="/content/orpo")
    ap.add_argument("--log_every", type=int, default=10)
    args = ap.parse_args()

    assert FORK, "mamba-ssm yok — önce wheel-cache kurulum hücresini çalıştır."
    assert torch.cuda.is_available(), "CUDA yok (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", weights_only=False)
    cfg = st["cfg"]
    policy = HybridLM(cfg, device=dev, dtype=torch.bfloat16)      # ORPO: SADECE policy (ref YOK)
    policy.load_state_dict(st["model"], strict=False); policy.train()
    print(f"[model] policy yüklendi (ORPO ref-free) | {'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["prompt"], r["chosen"], args.max_len, eos_id),
            encode_resp(sp, r["prompt"], 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, nlla = 0.0, 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):
                    logp_ch = seq_logp(policy, ch_ids, ch_lab, length_norm=True)
                    logp_rj = seq_logp(policy, rj_ids, rj_lab, length_norm=True)
                loss, acc, marg, nll = orpo_loss(logp_ch, logp_rj, args.beta)
                (loss / args.grad_accum).backward()
                la += loss.item() / args.grad_accum; acca += acc / args.grad_accum
                marga += marg / args.grad_accum; nlla += nll / 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} | nll {nlla:.3f} | acc {acca:.2f} | "
                      f"margin {marga:+.3f} | gnorm {gn:5.2f} | lr {lr:.2e} | {tps:.1f} pair/s", flush=True)
        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, "orpo": 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] ORPO tamamlandı.", flush=True)


if __name__ == "__main__":
    main()