"""
Faz 2 — Smoke test (eğitim döngüsü) + Colab-hazır ana iskelet.

İki mod:
  --synthetic : rastgele token batch'leriyle döngüyü doğrula (yerel CPU, network YOK,
                Faz 1 koşusunu rahatsız etmez). Model+loop+WSD+grad sağlıklı mı.
  (varsayılan): gerçek on-the-fly mixed streaming (EN/TR/kod/math + decontam) — Colab GPU.

Colab (full) örneği:
  python kod/faz2_smoke.py --d_model 768 --n_layer 20 --seq_len 2048 --steps 200 --device cuda
Yerel loop testi:
  python kod/faz2_smoke.py --synthetic --d_model 256 --n_layer 6 --seq_len 256 --steps 8 --device cpu

NOT: minimal (saf-PyTorch) hibrit kullanır → her yerde çalışır. Faz 3 tam-hız için
mamba-og fork (Triton SISO) ile aynı mimari; bu script veri+döngü+mimariyi doğrular.
"""
import os, sys, time, math, argparse, random
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import torch
import torch.nn.functional as F


# ---------------- WSD learning rate ----------------
def wsd_lr(step, total, peak, floor, warmup_frac=0.02, decay_frac=0.25):
    warm = max(1, int(total * warmup_frac))
    dec_start = int(total * (1 - decay_frac))
    if step < warm:
        return peak * (step + 1) / warm
    if step < dec_start:
        return peak
    t = (step - dec_start) / max(1, total - dec_start)
    return peak - (peak - floor) * t


# ---------------- Mixed streaming (Colab) ----------------
SOURCES = {
    "en_fineweb_edu":   ("HuggingFaceFW/fineweb-edu", "sample-10BT", 0.55),
    "tr_fineweb2_hq":   ("epfml/FineWeb2-HQ", "tur_Latn", 0.22),
    "code_codeparrot":  ("codeparrot/codeparrot-clean", None, 0.13),
    "math_openwebmath": ("open-web-math/open-web-math", None, 0.10),
}
TEXT_KEYS = ("text", "content", "code")


def _text(rec):
    for k in TEXT_KEYS:
        v = rec.get(k)
        if isinstance(v, str) and v:
            return v
    for v in rec.values():
        if isinstance(v, str) and v:
            return v
    return ""


class MixedStream:
    """Kaynakları oranlarına göre karıştırıp 2048'lik diziler üretir (decontam'lı)."""
    def __init__(self, tok, seq_len, decontam_path=None):
        os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1"
        from datasets import load_dataset
        self.tok, self.seq_len = tok, seq_len
        self.dec = None
        if decontam_path and os.path.exists(decontam_path):
            from decontam import Decontaminator
            self.dec = Decontaminator.load(decontam_path)
        self.names = list(SOURCES)
        self.weights = [SOURCES[n][2] for n in self.names]
        self.iters, self.bufs = {}, {n: [] for n in self.names}
        for n in self.names:
            repo, cfg, _ = SOURCES[n]
            self.iters[n] = iter(load_dataset(repo, name=cfg, split="train", streaming=True))

    def _reopen(self, n):
        from datasets import load_dataset
        repo, cfg, _ = SOURCES[n]
        self.iters[n] = iter(load_dataset(repo, name=cfg, split="train", streaming=True))

    def _fill(self, n):
        while len(self.bufs[n]) < self.seq_len:
            try:
                rec = next(self.iters[n])
            except StopIteration:      # stream tükendi → baştan aç (uzun koşuda epoch)
                self._reopen(n); rec = next(self.iters[n])
            txt = _text(rec)
            if self.dec is not None and self.dec.is_contaminated(txt):
                continue
            self.bufs[n].extend(self.tok.encode(txt, add_eos=True))

    def next_seq(self):
        n = random.choices(self.names, weights=self.weights, k=1)[0]
        self._fill(n)
        seq = self.bufs[n][:self.seq_len]; self.bufs[n] = self.bufs[n][self.seq_len:]
        return seq

    def batch(self, bsz):
        return torch.tensor([self.next_seq() for _ in range(bsz)], dtype=torch.long)


def synthetic_batch(bsz, seq_len, vocab):
    return torch.randint(0, vocab, (bsz, seq_len), dtype=torch.long)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--synthetic", action="store_true")
    ap.add_argument("--d_model", type=int, default=256)
    ap.add_argument("--n_layer", type=int, default=6)
    ap.add_argument("--seq_len", type=int, default=256)
    ap.add_argument("--micro_batch", type=int, default=4)
    ap.add_argument("--steps", type=int, default=8)
    ap.add_argument("--peak_lr", type=float, default=5e-4)
    ap.add_argument("--device", default="cpu")
    ap.add_argument("--bf16", action="store_true",
                    help="bf16 autocast (cuda). Minimal saf-PyTorch SSD'de test edilmedi → "
                         "smoke'u önce fp32 koş; bf16 asıl fork (Faz 3) içindir.")
    ap.add_argument("--tokenizer", default="kod/tokenizer/tokenizer.model")
    ap.add_argument("--decontam", default="kod/data/decontam_13gram.pkl.gz")
    args = ap.parse_args()

    dev = torch.device(args.device if (args.device != "cuda" or torch.cuda.is_available()) else "cpu")
    torch.manual_seed(0); random.seed(0)

    # tokenizer (vocab → model boyutu)
    from sc_tokenizer import SCTokenizer
    tok = SCTokenizer(args.tokenizer)
    vocab = tok.vocab_size

    # model: minimal hibrit (Mamba-3 SISO + GQA 5:1)
    from hybrid_mamba3 import make_config, build_hybrid, init_weights, n_params
    cfg = make_config(d_model=args.d_model, n_layer=args.n_layer, vocab=vocab,
                      d_mlp_inner=1500 if args.d_model >= 768 else args.d_model * 2,
                      chunk_size=min(64, args.seq_len))
    attn_every = 6 if args.n_layer >= 8 else 3
    model, attn_idx = build_hybrid(cfg, attn_every=attn_every, n_heads=max(2, args.d_model // 64),
                                   n_kv_heads=max(1, args.d_model // 256), device=dev)
    model.to(dev); model.train()
    print(f"model: d={args.d_model} L={args.n_layer} ({args.n_layer-len(attn_idx)} Mamba + "
          f"{len(attn_idx)} GQA) vocab={vocab} | params={n_params(model)/1e6:.1f}M | dev={dev}")

    # optimizer: 2D ağırlıklara wd, diğerlerine yok
    decay = [p for p in model.parameters() if p.ndim >= 2]
    nodecay = [p for p in model.parameters() if p.ndim < 2]
    opt = torch.optim.AdamW([{"params": decay, "weight_decay": 0.1},
                             {"params": nodecay, "weight_decay": 0.0}],
                            lr=args.peak_lr, betas=(0.9, 0.95), eps=1e-8)

    stream = None
    if not args.synthetic:
        print("on-the-fly mixed streaming kuruluyor (EN/TR/kod/math + decontam)...")
        stream = MixedStream(tok, args.seq_len, args.decontam)

    use_bf16 = args.bf16 and dev.type == "cuda"   # default fp32 (minimal SSD bf16'da test edilmedi)
    ln_v = math.log(vocab)
    print(f"smoke başlıyor | {args.steps} adım | başlangıç loss ~ln({vocab})={ln_v:.2f}")
    t0 = time.perf_counter(); seen = 0
    for step in range(args.steps):
        batch = (synthetic_batch(args.micro_batch, args.seq_len, vocab) if args.synthetic
                 else stream.batch(args.micro_batch)).to(dev)
        for g in opt.param_groups:
            g["lr"] = wsd_lr(step, args.steps, args.peak_lr, args.peak_lr * 0.1)
        opt.zero_grad(set_to_none=True)
        ctx = torch.autocast(device_type="cuda", dtype=torch.bfloat16) if use_bf16 else _null()
        with ctx:
            logits, _ = model(batch)
            loss = F.cross_entropy(logits[:, :-1].reshape(-1, logits.size(-1)).float(),
                                   batch[:, 1:].reshape(-1))
        loss.backward()
        gnorm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()
        seen += batch.numel()
        tok_s = seen / (time.perf_counter() - t0)
        print(f"  step {step:3d} | loss {loss.item():6.3f} | grad_norm {gnorm:5.2f} | "
              f"lr {opt.param_groups[0]['lr']:.2e} | {tok_s/1e3:.1f}k tok/s")

    print(f"\nsmoke BİTTİ. NaN yok: {not math.isnan(loss.item())} | "
          f"throughput {seen/(time.perf_counter()-t0)/1e3:.1f}k tok/s")


class _null:
    def __enter__(self): return self
    def __exit__(self, *a): return False


if __name__ == "__main__":
    main()