"""
SmartCore V1 — Mamba-3 (SISO) + GQA 5:1 hibrit, 48K EN+TR hedefi.

Referans iskelet: referans_kod/mamba3_minimal.py (saf-PyTorch Mamba-3, Triton'suz).
Bu dosya o modelin üstüne GQA attention mixer'ı ekler ve 5:1 hibrit kurar:
her (attn_every)'inci katmanın SSM mixer'ı yerine GQAttention konur; ilk ve son
katman Mamba kalır (induction head'ler derinlik ortasında oluşsun diye).

Forward arayüzü minimal modelle birebir: mixer(x, h) -> (y, h). GQA, h'i yok sayar
(eğitim/forward yolu); KV-cache'li decode bu dosyada gerekli değil (param + smoke).

CPU'da çalışır (Ryzen 5500). Amaç: hibrit öğreniyor mu + kesin ~180M param-sayımı.
"""
import sys, time, gc
import torch
import torch.nn as nn
import torch.nn.functional as F

sys.path.insert(0, "referans_kod")
import mamba3_minimal as m3
from mamba3_minimal import Mamba3Config, Mamba3LMHeadModel, RMSNorm

torch.manual_seed(0)
torch.set_num_threads(6)
DEV = torch.device("cpu")


def rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)


class GQAttention(nn.Module):
    """Grouped-Query Attention mixer (QK-norm + RoPE, causal). minimal-Mamba3 ile
    aynı (y, h) arayüzü. n_heads q-head, n_kv_heads kv-head (GQA sıkıştırma)."""

    def __init__(self, d_model, n_heads=12, n_kv_heads=3, rope_base=10000.0, device=None):
        super().__init__()
        assert d_model % n_heads == 0, "d_model n_heads'e bölünmeli"
        assert n_heads % n_kv_heads == 0, "n_heads n_kv_heads'e bölünmeli (GQA)"
        self.n_heads, self.n_kv, self.hd = n_heads, n_kv_heads, d_model // n_heads
        self.rep = n_heads // n_kv_heads
        self.q_proj = nn.Linear(d_model, n_heads * self.hd, bias=False, device=device)
        self.k_proj = nn.Linear(d_model, n_kv_heads * self.hd, bias=False, device=device)
        self.v_proj = nn.Linear(d_model, n_kv_heads * self.hd, bias=False, device=device)
        self.o_proj = nn.Linear(n_heads * self.hd, d_model, bias=False, device=device)
        self.q_norm = RMSNorm(self.hd, device=device)   # QK-norm (RoPE'den önce)
        self.k_norm = RMSNorm(self.hd, device=device)
        inv_freq = 1.0 / (rope_base ** (torch.arange(0, self.hd, 2, device=device).float() / self.hd))
        self.register_buffer("inv_freq", inv_freq, persistent=False)

    def _rope(self, x, T):  # x: (B, nh, T, hd)
        t = torch.arange(T, device=x.device, dtype=torch.float32)
        freqs = torch.outer(t, self.inv_freq)            # (T, hd/2)
        emb = torch.cat((freqs, freqs), dim=-1)          # (T, hd)
        cos, sin = emb.cos()[None, None], emb.sin()[None, None]
        return (x * cos + rotate_half(x) * sin).to(x.dtype)

    def forward(self, x, h=None):
        B, T, _ = x.shape
        q = self.q_proj(x).view(B, T, self.n_heads, self.hd).transpose(1, 2)
        k = self.k_proj(x).view(B, T, self.n_kv, self.hd).transpose(1, 2)
        v = self.v_proj(x).view(B, T, self.n_kv, self.hd).transpose(1, 2)
        q, k = self.q_norm(q), self.k_norm(k)            # QK-norm
        q, k = self._rope(q, T), self._rope(k, T)
        k = k.repeat_interleave(self.rep, dim=1)         # GQA: kv-head'leri çoğalt (CPU-güvenli)
        v = v.repeat_interleave(self.rep, dim=1)
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        y = y.transpose(1, 2).contiguous().view(B, T, -1)
        return self.o_proj(y), h


def init_weights(model):
    """SSM-güvenli init (minimal demo.py'den): A_log, D, dt_bias kritik —
    constructor bunları set etmiyor; atlanırsa SSM NaN üretir. GQA projeksiyonları
    (2D) std=0.02, RMSNorm ağırlıkları (1D) dokunulmaz (default 1.0)."""
    for name, p in model.named_parameters():
        if "A_log" in name:
            torch.nn.init.uniform_(p, -4, -1)
        elif "D" in name and p.dim() == 1:
            torch.nn.init.ones_(p)
        elif "dt_bias" in name:
            torch.nn.init.uniform_(p, 0.001, 0.1)
        elif p.dim() >= 2:
            torch.nn.init.normal_(p, std=0.02)


def build_hybrid(args, attn_every=6, n_heads=12, n_kv_heads=3, device=None):
    """minimal Mamba3LMHeadModel kur, sonra her attn_every'inci katmanın mixer'ını
    GQAttention ile değiştir (ilk/son katman Mamba kalır)."""
    model = Mamba3LMHeadModel(args, device=device)
    attn_idx = []
    for i in range(args.n_layer):
        if (i + 1) % attn_every == 0 and i != 0 and i != args.n_layer - 1:
            model.backbone.layers[i].mixer = GQAttention(
                args.d_model, n_heads, n_kv_heads, device=device)
            attn_idx.append(i)
    init_weights(model)  # GQA swap'ten SONRA — tüm parametreler güvenli init'lensin
    return model, attn_idx


def n_params(model):
    seen, tot = set(), 0
    for p in model.parameters():
        if id(p) in seen:
            continue
        seen.add(id(p)); tot += p.numel()
    return tot


def make_config(d_model, n_layer, vocab, d_mlp_inner=1500, d_state=128, headdim=64, chunk_size=64):
    c = Mamba3Config(d_model=d_model, n_layer=n_layer, d_state=d_state,
                     headdim=headdim, chunk_size=chunk_size, vocab_size=vocab)
    c.d_mlp_inner = d_mlp_inner  # referans eğitim reçetesi (1500), minimal'in 2048'i yerine
    return c


def param_sweep():
    """180M hedef için 48K vocab'ta katman sweep'i (import edilmez, sadece doğrudan çalıştırınca)."""
    print(f"  {'n_layer':>7} {'Mamba':>6} {'GQA':>4} {'embed(M)':>9} {'TOPLAM(M)':>10}")
    for nlayer in [18, 19, 20, 21, 22]:
        c = make_config(d_model=768, n_layer=nlayer, vocab=48000, d_mlp_inner=1500)
        mdl, ai = build_hybrid(c, attn_every=6, n_heads=12, n_kv_heads=3, device=DEV)
        tot = n_params(mdl) / 1e6
        embed_m = c.vocab_size * c.d_model / 1e6
        mark = "  <== 180M bandi" if 175 <= tot <= 185 else ""
        print(f"  {nlayer:>7} {nlayer-len(ai):>6} {len(ai):>4} {embed_m:>9.1f} {tot:>10.1f}{mark}")
        del mdl; gc.collect()


if __name__ == "__main__":
    # Doğrulanan sonuç: n_layer=20 -> 175.5M (17 Mamba+3 GQA), n_layer=21 -> 182.8M.
    param_sweep()