"""
faz9_recur.py — Recurrent-depth (latent reasoning) sarmalayıcı, LAMBA bloğuyla. [v1.5 Faz A]

Hipotez: gizli-uzayda bir çekirdek bloğu R kez tekrar = EŞİT PARAMDA daha çok "düşünme"
(test-zamanı compute). Kaynak: TRM/HRM (2602.12078 — post-norm ŞART) + Geiping recurrent-depth.

Mimari:  embed → PRELUDE (pre-norm) → CORE ×R (post-norm, girdi-enjeksiyonlu, truncated-BPTT) → CODA → lm_head
KRİTİK: çekirdek POST-NORM (h←Norm(h+F(h))). Unrolled recursion'da pre-norm residual'ı ~√t
büyütür → NaN; post-norm sınırlar. Bu dosya bunu bizim Mamba-3+GQA bloklarımızla CPU'da kanıtlar.

Smoke:  python faz9_recur.py --smoke
"""
import os, sys, argparse
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))   # HF symlink-cache: lamba_cpu'yu yan klasörden bul
import torch, torch.nn as nn, torch.nn.functional as F
from lamba_cpu import rms_norm, Mamba3CPU, GQACPU, GatedMLP, Block   # doğrulanmış saf-PyTorch bloklar


def tiny_cfg(**ov):
    cfg = dict(d_model=128, expand=2, head_dim=32, d_state=64, ngroups=1, rope_fraction=0.5,
               d_intermediate=256, vocab_size=512, n_heads=4, n_kv_heads=2)
    cfg.update(ov); return cfg


class CoreBlock(nn.Module):
    """R kez tekrarlanan çekirdek. prenorm=False → POST-NORM (recursion-kararlı)."""
    def __init__(self, cfg, is_attn=False, prenorm=False):
        super().__init__()
        self.prenorm = prenorm
        self.mixer = GQACPU(cfg) if is_attn else Mamba3CPU(cfg)
        self.mlp = GatedMLP(cfg)
        self.n1 = nn.Parameter(torch.ones(cfg["d_model"]))
        self.n2 = nn.Parameter(torch.ones(cfg["d_model"]))

    def forward(self, h, inject=None):
        if inject is not None:                 # prelude çıktısını her adımda enjekte et (girdiyi unutma)
            h = h + inject
        if self.prenorm:                       # PRE-NORM: kıyas için (recursion'da kararsız)
            h = h + self.mixer(rms_norm(h, self.n1))
            h = h + self.mlp(rms_norm(h, self.n2))
        else:                                  # POST-NORM: h←Norm(h+F(h)) — ölçeği sınırlar
            h = rms_norm(h + self.mixer(h), self.n1)
            h = rms_norm(h + self.mlp(h), self.n2)
        return h


class RecurrentDepthLM(nn.Module):
    def __init__(self, cfg, n_prelude=2, n_coda=2, core_attn=False, prenorm_core=False, block_attn=False):
        super().__init__()
        self.embedding = nn.Embedding(cfg["vocab_size"], cfg["d_model"])
        self.prelude = nn.ModuleList([Block(cfg, is_attn=block_attn) for _ in range(n_prelude)])
        self.core = CoreBlock(cfg, is_attn=core_attn, prenorm=prenorm_core)   # TEK paylaşımlı blok
        self.coda = nn.ModuleList([Block(cfg, is_attn=block_attn) for _ in range(n_coda)])
        self.norm_f = nn.Parameter(torch.ones(cfg["d_model"]))
        self.lm_head = nn.Linear(cfg["d_model"], cfg["vocab_size"], bias=False)
        self.lm_head.weight = self.embedding.weight

    def forward(self, ids, R=8, bptt_k=8, want_hidden_norm=False):
        h = self.embedding(ids)
        for l in self.prelude:
            h = l(h)
        prelude_out = h
        no_grad_iters = max(0, R - bptt_k)             # truncated-BPTT: sadece son k iterasyon gradyan taşır
        hn = []
        with torch.no_grad():
            for _ in range(no_grad_iters):
                h = self.core(h, inject=prelude_out)
                if want_hidden_norm: hn.append(h.norm(dim=-1).mean().item())
        for _ in range(R - no_grad_iters):
            h = self.core(h, inject=prelude_out)
            if want_hidden_norm: hn.append(h.norm(dim=-1).mean().item())
        for l in self.coda:
            h = l(h)
        logits = self.lm_head(rms_norm(h, self.norm_f))
        return (logits, hn) if want_hidden_norm else logits


def _params(m): return sum(p.numel() for p in m.parameters())


def smoke():
    torch.manual_seed(0)
    cfg = tiny_cfg()
    B, L = 2, 16
    ids = torch.randint(0, cfg["vocab_size"], (B, L))

    # ── 1) KARARLILIK: post-norm vs pre-norm, R büyürken hidden-norm / NaN ──
    print("=" * 64)
    print("[1] RECURSION KARARLILIĞI — hidden-state L2 normu (R büyüyor)")
    print("    post-norm sınırlı kalmalı; pre-norm ~√R büyüyüp patlamalı")
    for prenorm in (False, True):
        torch.manual_seed(0)
        m = RecurrentDepthLM(cfg, prenorm_core=prenorm).eval()
        tag = "PRE-norm " if prenorm else "POST-norm"
        row = []
        for R in (1, 4, 16, 64):
            with torch.no_grad():
                lg, hn = m(ids, R=R, want_hidden_norm=True)
            nan = torch.isnan(lg).any().item()
            row.append(f"R={R:>3}: |h|={hn[-1]:8.2f}{' NaN!' if nan else ''}")
        print(f"  {tag} | " + " | ".join(row))

    # ── 2) EĞİTİLEBİLİRLİK: post-norm çekirdek, değişken-R, birkaç adım loss düşüyor mu ──
    print("=" * 64)
    print("[2] EĞİTİM smoke (post-norm, değişken-R, truncated-BPTT k=8)")
    torch.manual_seed(0)
    m = RecurrentDepthLM(cfg, prenorm_core=False).train()
    print(f"    model: {_params(m)/1e6:.2f}M param (çekirdek PAYLAŞIMLI → R'den bağımsız)")
    opt = torch.optim.AdamW(m.parameters(), lr=3e-3)
    # sabit küçük "öğrenilebilir" görev: rastgele ama sabit dizi (ezber → loss düşmeli)
    data = torch.randint(0, cfg["vocab_size"], (B, L + 1))
    x, y = data[:, :-1], data[:, 1:]
    for step in range(40):
        R = int(torch.randint(2, 13, (1,)))            # değişken R ~ U[2,12]
        logits = m(x, R=R, bptt_k=8)
        loss = F.cross_entropy(logits.reshape(-1, cfg["vocab_size"]), y.reshape(-1))
        opt.zero_grad(); loss.backward()
        gnorm = torch.nn.utils.clip_grad_norm_(m.parameters(), 1.0)
        opt.step()
        if step % 8 == 0 or step == 39:
            print(f"    step {step:>2} | R={R:>2} | loss {loss.item():.3f} | gnorm {gnorm:.2f}"
                  f"{' NaN!' if torch.isnan(loss) else ''}")

    # ── 3) TEST-ZAMANI ÖLÇEKLEME: eğitilmiş modelde R↑ ile aynı-dizi loss'u ──
    print("=" * 64)
    print("[3] TEST-ZAMANI compute: eğitilmiş model, R↑ → loss (daha çok 'düşünme')")
    m.eval()
    with torch.no_grad():
        for R in (1, 2, 4, 8, 16):
            lg = m(x, R=R)
            l = F.cross_entropy(lg.reshape(-1, cfg["vocab_size"]), y.reshape(-1)).item()
            print(f"    R={R:>2} → loss {l:.3f}")
    print("=" * 64)
    print("SMOKE BİTTİ. Beklenen: [1] post-norm sınırlı/pre-norm patlar, [2] loss düşer NaN yok, "
          "[3] R↑ ile loss ≤ (bozulmaz).")


# ───────────── FAZ B: derinlik-gerektiren görev + fixed-vs-recurrent (GO/NO-GO) ─────────────
def gen_khop(batch, n_keys, k, device="cpu"):
    """k-hop ZİNCİR traversali (döngü-kısayolu YOK). Düğümler tek bir zincir (chain=randperm),
    kenarlar [chain[i], chain[i+1]] KARIŞIK sırada → SEP → start=chain[p]. Hedef=chain[p+k], p∈[0,n-k).
    Zincir döngüsüz → hedef≠start her zaman → 'başlangıcı kopyala' kısayolu yok → gerçek k sıralı hop gerek.
    Vocab 0..n_keys-1 + SEP(n_keys). Loss SADECE son pozisyon. (Eski permütasyon görevi σ^k(s)=s döngü-kısayoluyla kirleniyordu.)"""
    assert n_keys > k, "n_keys > k olmalı"
    SEP = n_keys; seqs = []; tgts = []
    for _ in range(batch):
        chain = torch.randperm(n_keys)
        edges = [[int(chain[i]), int(chain[i + 1])] for i in range(n_keys - 1)]
        toks = []
        for j in torch.randperm(len(edges)).tolist():
            toks += edges[j]
        p = int(torch.randint(0, n_keys - k, (1,)))
        toks += [SEP, int(chain[p])]
        seqs.append(toks); tgts.append(int(chain[p + k]))
    return (torch.tensor(seqs, device=device), torch.tensor(tgts, device=device))


class FixedDepthLM(nn.Module):
    """Param-eşit sabit-derinlik baseline (recursion yok)."""
    def __init__(self, cfg, n_layers, all_attn=False):
        super().__init__()
        self.embedding = nn.Embedding(cfg["vocab_size"], cfg["d_model"])
        self.layers = nn.ModuleList([Block(cfg, is_attn=(all_attn or i % 3 == 1)) for i in range(n_layers)])
        self.norm_f = nn.Parameter(torch.ones(cfg["d_model"]))
        self.lm_head = nn.Linear(cfg["d_model"], cfg["vocab_size"], bias=False)
        self.lm_head.weight = self.embedding.weight

    def forward(self, ids, **kw):
        h = self.embedding(ids)
        for l in self.layers:
            h = l(h)
        return self.lm_head(rms_norm(h, self.norm_f))


def _train_eval(device, n_keys, k, steps, batch, d_model, lr, fast, verbose=True):
    """Tek k için param-eşit fixed+recurrent eğit → (fix_acc, {R:rec_acc}, n_param)."""
    P, C = 1, 1
    cfg = tiny_cfg(d_model=d_model, vocab_size=n_keys + 1)
    rec = RecurrentDepthLM(cfg, n_prelude=P, n_coda=C, core_attn=True, prenorm_core=False, block_attn=fast).to(device)
    fix = FixedDepthLM(cfg, n_layers=P + 1 + C, all_attn=fast).to(device)

    def train(model, recurrent):
        opt = torch.optim.AdamW(model.parameters(), lr=lr); model.train()
        for s in range(steps):
            ids, tgt = gen_khop(batch, n_keys, k, device)
            R = int(torch.randint(2, 13, (1,))) if recurrent else 1
            last = (model(ids, R=R, bptt_k=8) if recurrent else model(ids))[:, -1]
            loss = F.cross_entropy(last, tgt)
            opt.zero_grad(); loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0); opt.step()
            if verbose and s % max(1, steps // 4) == 0:
                print(f"    {'REC' if recurrent else 'FIX'} k={k} step {s:>4} loss {loss.item():.3f}")

    @torch.no_grad()
    def acc(model, recurrent, R=8, n=1000):
        model.eval(); ids, tgt = gen_khop(n, n_keys, k, device)
        last = (model(ids, R=R) if recurrent else model(ids))[:, -1]
        return (last.argmax(-1) == tgt).float().mean().item()

    train(fix, False); train(rec, True)
    rec_by_R = {R: acc(rec, True, R=R) for R in (1, 2, 4, 8, 16)}
    return acc(fix, False), rec_by_R, _params(rec)


def compare(device="cpu", n_keys=12, k=4, steps=2000, batch=64, d_model=128, lr=1e-3, fast=False):
    print(f"[compare] görev: {k}-hop / {n_keys} anahtar | device={device} | steps={steps} | fast={fast}")
    fa, rec, npar = _train_eval(device, n_keys, k, steps, batch, d_model, lr, fast)
    rnd = 1.0 / n_keys
    print(f"\n[SONUÇ] param ~{npar/1e6:.3f}M (eşit) | random={rnd:.3f}")
    print(f"  FIXED (depth 3) acc = {fa:.3f}")
    print(f"  RECURRENT test-zamanı R ölçekleme:")
    for R, a in rec.items():
        print(f"    R={R:>2} → acc {a:.3f}")
    print(f"  GAP (rec_best − fixed) = {max(rec.values()) - fa:+.3f}")
    print("  GO sinyali: recurrent fixed'i geçer VE R↑ ile acc↑ (özellikle R≥k).")


def sweep(device="cpu", n_keys=12, ks=(2, 4, 6), steps=2000, batch=64, d_model=128, lr=1e-3, fast=True):
    """ZORLUK TARAMASI: GAP (rec−fixed) k ile büyüyor mu = recursion derinlikle değer kazanıyor mu."""
    rnd = 1.0 / n_keys
    print(f"[sweep] k={list(ks)} / {n_keys} anahtar | device={device} | steps={steps} | fast={fast} | random={rnd:.3f}\n")
    rows = []
    for k in ks:
        fa, rec, npar = _train_eval(device, n_keys, k, steps, batch, d_model, lr, fast, verbose=False)
        bR = max(rec, key=rec.get); gap = rec[bR] - fa
        rows.append((k, fa, rec, bR, gap))
        print(f"  k={k}: FIX {fa:.3f} | REC R1={rec[1]:.2f} R4={rec[4]:.2f} R8={rec[8]:.2f} R16={rec[16]:.2f} "
              f"| best@R{bR}={rec[bR]:.3f} | GAP {gap:+.3f}")
    print(f"\n[ÖZET] GAP (rec_best − fixed) — k arttıkça büyümeli (random={rnd:.3f}):")
    for k, fa, rec, bR, gap in rows:
        print(f"  k={k}: GAP {gap:+.3f}  {'#' * max(0, int(gap * 60))}")
    gaps = [r[4] for r in rows]
    grow = all(gaps[i] <= gaps[i + 1] + 1e-9 for i in range(len(gaps) - 1)) and gaps[-1] > 0.05
    print(f"\n  → GAP k ile {'✅ ARTIYOR = GO (recursion derinlikle değer kazanıyor)' if grow else '⚠️ ARTMIYOR/karışık = zayıf sinyal'}")
    print("  → Ayrıca zor k'da REC acc ~R≥k civarında sıçramalı (daha çok düşünme = daha çok hop).")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--smoke", action="store_true", help="CPU mekanik smoke (Faz A)")
    ap.add_argument("--compare", action="store_true", help="Faz B: fixed-vs-recurrent k-hop (GO/NO-GO)")
    ap.add_argument("--device", default="cpu", choices=["cpu", "cuda"])
    ap.add_argument("--n_keys", type=int, default=12)
    ap.add_argument("--k", type=int, default=4, help="hop sayısı (derinlik talebi)")
    ap.add_argument("--steps", type=int, default=2000)
    ap.add_argument("--d_model", type=int, default=128)
    ap.add_argument("--fast", action="store_true", help="tüm bloklar attention (Mamba token-loop yok → ~30× hızlı; k-hop için uygun)")
    ap.add_argument("--sweep", action="store_true", help="Faz B zorluk taraması: k=2/4/6 GAP trendi (asıl GO/NO-GO)")
    ap.add_argument("--ks", default="2,4,6", help="sweep k değerleri (virgülle, örn. 2,4,6)")
    args = ap.parse_args()
    if args.smoke:
        smoke()
    elif args.sweep:
        ks = tuple(int(x) for x in args.ks.split(","))
        sweep(device=args.device, n_keys=args.n_keys, ks=ks, steps=args.steps, d_model=args.d_model, fast=args.fast)
    elif args.compare:
        compare(device=args.device, n_keys=args.n_keys, k=args.k, steps=args.steps, d_model=args.d_model, fast=args.fast)
    else:
        print("kullanım: python faz9_recur.py --smoke | --compare | --sweep [--device cuda --n_keys 12 --steps 3000 --fast]")


if __name__ == "__main__":
    main()