#!/usr/bin/env python3
# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "torch>=2.4",
#     "transformers>=4.45",
#     "datasets>=3.2",
#     "safetensors>=0.4",
#     "huggingface_hub>=0.25",
#     "zstandard",
# ]
# ///
"""
kiyo_diffusion.py - transforma o DedeProGames/Kiyo-135M (LM autoregressivo, Llama/SmolLM2)
num modelo de linguagem de difusao mascarada (estilo LLaDA / Dream / MetaDiffusion).

Receita (adaptacao AR -> difusao do DiffuLLaMA / Dream):
  1. mesma arquitetura e mesmos pesos do Kiyo, mas atencao BIDIRECIONAL
  2. token novo <|mask|> (vocab 49152 -> 49153), embedding inicializado com a media
  3. shift: o hidden da posicao i preve o token i+1 (igual no AR), entao o lm_head
     pre-treinado continua alinhado com o que o Kiyo ja sabe
  4. perda de difusao mascarada (MDLM/LLaDA): t ~ U(eps,1), mascara cada token com prob. t,
     cross-entropy so nos mascarados com peso 1/t (a "nelbo" do log e limite superior da NLL/token)
  5. annealing da atencao: causal -> bidirecional nos primeiros --anneal-steps steps
  6. sem timestep embedding: a fracao de [MASK] na entrada ja diz o nivel de ruido
  7. geracao: comeca tudo [MASK] e revela os tokens mais confiantes, bloco a bloco (LLaDA)

Uso:
  python kiyo_diffusion.py convert  --src DedeProGames/Kiyo-135M --out kiyo-diff-init
  python kiyo_diffusion.py check    --model kiyo-diff-init --src DedeProGames/Kiyo-135M
  python kiyo_diffusion.py train    --out runs/kiyo-diff     # 1B tokens; usa todas as GPUs (torchrun automatico)
  python kiyo_diffusion.py train    --out /content/drive/MyDrive/kiyo-diff --tokens 5e8   # Colab (retoma sozinho)
  python kiyo_diffusion.py generate --model runs/kiyo-diff/final --prompt "The cat sat on the" --watch
  python kiyo_diffusion.py generate --model runs/kiyo-diff/final --gen-len 0 \\
         --prompt "The capital of France is<|mask|><|mask|>. It is famous for the Eiffel Tower."  # infilling
  python kiyo_diffusion.py score    --model runs/kiyo-diff/final --context "The capital of France is" \\
         --choices " Paris" " London" " Rome"
"""
import argparse
import array
import glob
import json
import math
import os
import platform
import random
import shutil
import sys
import time
from contextlib import nullcontext

os.environ.setdefault("HF_HUB_DOWNLOAD_TIMEOUT", "60")  # o padrao (10 s) da timeout no Colab
os.environ.setdefault("HF_HUB_ETAG_TIMEOUT", "60")

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint
from safetensors.torch import load_file, save_file

MASK_TOKEN = "<|mask|>"
KIYO_MIX = [  # mesmas fontes do Kiyo; o peso e a fracao de TOKENS (ajusta pro mix que tu usou)
    "HuggingFaceFW/fineweb-edu@0.45",
    "mlfoundations/dclm-baseline-1.0@0.35",
    "HuggingFaceTB/finemath:finemath-3plus@0.10",
    "HuggingFaceCode/stack-v3-train@0.10",
]


# ============================================================================ modelo
class RMSNorm(nn.Module):
    def __init__(self, dim, eps):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.eps = eps

    def forward(self, x):
        dt = x.dtype
        x = x.float()
        x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return self.weight * x.to(dt)


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


class Attention(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.nh, self.nkv, self.hd = c["num_attention_heads"], c["num_key_value_heads"], c["head_dim"]
        H = c["hidden_size"]
        self.q_proj = nn.Linear(H, self.nh * self.hd, bias=False)
        self.k_proj = nn.Linear(H, self.nkv * self.hd, bias=False)
        self.v_proj = nn.Linear(H, self.nkv * self.hd, bias=False)
        self.o_proj = nn.Linear(self.nh * self.hd, H, bias=False)

    def forward(self, x, cos, sin, attn_mask=None, causal=False):
        B, L, _ = x.shape
        q = self.q_proj(x).view(B, L, self.nh, self.hd).transpose(1, 2)
        k = self.k_proj(x).view(B, L, self.nkv, self.hd).transpose(1, 2)
        v = self.v_proj(x).view(B, L, self.nkv, self.hd).transpose(1, 2)
        cos, sin = cos.to(q.dtype), sin.to(q.dtype)
        q = q * cos + rotate_half(q) * sin
        k = k * cos + rotate_half(k) * sin
        if self.nkv != self.nh:  # GQA
            r = self.nh // self.nkv
            k, v = k.repeat_interleave(r, dim=1), v.repeat_interleave(r, dim=1)
        # attn_mask=None e causal=False -> atencao bidirecional completa (kernel flash do SDPA)
        o = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask,
                                           is_causal=causal and attn_mask is None)
        return self.o_proj(o.transpose(1, 2).reshape(B, L, -1))


class MLP(nn.Module):
    def __init__(self, H, I):
        super().__init__()
        self.gate_proj = nn.Linear(H, I, bias=False)
        self.up_proj = nn.Linear(H, I, bias=False)
        self.down_proj = nn.Linear(I, H, bias=False)

    def forward(self, x):
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


class Block(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.input_layernorm = RMSNorm(c["hidden_size"], c["rms_norm_eps"])
        self.self_attn = Attention(c)
        self.post_attention_layernorm = RMSNorm(c["hidden_size"], c["rms_norm_eps"])
        self.mlp = MLP(c["hidden_size"], c["intermediate_size"])

    def forward(self, x, cos, sin, attn_mask=None, causal=False):
        x = x + self.self_attn(self.input_layernorm(x), cos, sin, attn_mask, causal)
        return x + self.mlp(self.post_attention_layernorm(x))


class KiyoDiffusion(nn.Module):
    """Llama identico ao Kiyo (mesmos nomes de pesos), so que sem mascara causal."""

    def __init__(self, c):
        super().__init__()
        self.cfg = c
        self.embed_tokens = nn.Embedding(c["vocab_size"], c["hidden_size"])
        self.layers = nn.ModuleList([Block(c) for _ in range(c["num_hidden_layers"])])
        self.norm = RMSNorm(c["hidden_size"], c["rms_norm_eps"])
        tied = c.get("tie_word_embeddings", True)
        self.lm_head = None if tied else nn.Linear(c["hidden_size"], c["vocab_size"], bias=False)
        self.grad_ckpt = False

    def rope(self, L, device):
        # sempre em fp32 (barato) -> o modelo em bf16 nao perde precisao no RoPE
        hd = self.cfg["head_dim"]
        inv = 1.0 / (self.cfg["rope_theta"] ** (torch.arange(0, hd, 2, device=device, dtype=torch.float32) / hd))
        f = torch.outer(torch.arange(L, device=device, dtype=torch.float32), inv)
        emb = torch.cat((f, f), dim=-1)
        return emb.cos()[None, None], emb.sin()[None, None]

    def hidden(self, ids, attn_mask=None, causal=False):
        x = self.embed_tokens(ids)
        cos, sin = self.rope(ids.shape[1], ids.device)
        for blk in self.layers:
            if self.grad_ckpt and self.training:
                x = torch.utils.checkpoint.checkpoint(blk, x, cos, sin, attn_mask, causal, use_reentrant=False)
            else:
                x = blk(x, cos, sin, attn_mask, causal)
        return self.norm(x)

    def forward(self, ids, targets=None, loss_mask=None, p_mask=None, attn_mask=None, causal=False):
        """Sem targets: logits alinhados (logits[:, j] = distribuicao do token da posicao j).
        Com targets: (nelbo, soma da CE nos mascarados, n de mascarados)."""
        h = self.hidden(ids, attn_mask, causal)
        W = self.embed_tokens.weight if self.lm_head is None else self.lm_head.weight
        shift = self.cfg.get("shift_logits", True)
        if targets is None:
            logits = F.linear(h, W)
            if shift:
                logits = torch.cat([logits[:, :1], logits[:, :-1]], dim=1)
            return logits
        if shift:  # hidden i preve o token i+1
            h, targets, loss_mask, p_mask = h[:, :-1], targets[:, 1:], loss_mask[:, 1:], p_mask[:, 1:]
        # lm_head so nas posicoes mascaradas (economiza muita memoria com vocab de 49k)
        logits = F.linear(h[loss_mask], W).float()
        logits[:, self.cfg["mask_token_id"]] = -float("inf")  # [MASK] nunca e uma resposta valida
        ce = F.cross_entropy(logits, targets[loss_mask], reduction="none")
        nelbo = (ce / p_mask[loss_mask]).sum() / targets.numel()
        return nelbo, ce.detach().sum(), loss_mask.sum()


# ============================================================================ IO
def read_safetensors(path):
    idx = os.path.join(path, "model.safetensors.index.json")
    files = sorted(set(json.load(open(idx))["weight_map"].values())) if os.path.exists(idx) else ["model.safetensors"]
    sd = {}
    for f in files:
        sd.update(load_file(os.path.join(path, f)))
    return sd


def local_dir(path):
    if os.path.isdir(path):
        return path
    from huggingface_hub import snapshot_download
    return snapshot_download(path, allow_patterns=["*.json", "*.safetensors", "*.txt", "*.py", "*.jinja"])


def write_card(out, cfg, repo=None):
    base = cfg.get("base_model", "DedeProGames/Kiyo-135M")
    name = (repo or os.path.basename(os.path.normpath(out))).split("/")[-1]
    with open(os.path.join(out, "README.md"), "w") as f:
        f.write(f"""---
license: apache-2.0
language: [en]
pipeline_tag: text-generation
base_model: {base}
tags: [diffusion-language-model, masked-diffusion, autoregressive-conversion, llama, kiyo]
---

# {name}

Masked diffusion language model adapted from [{base}](https://huggingface.co/{base}): same Llama
weights, bidirectional attention, a new `<|mask|>` token, the masked-diffusion (MDLM/LLaDA) objective
and the AR shift operation (DiffuLLaMA/Dream recipe). No timestep embedding. Research use only.

```bash
python kiyo_diffusion.py generate --model {repo or name} --prompt "The capital of France is"
```
""")


def save_pretrained(model, cfg, tok, out, dtype=torch.bfloat16):
    os.makedirs(out, exist_ok=True)
    sd = {(k if k.startswith("lm_head") else "model." + k): v.detach().to("cpu", dtype).contiguous()
          for k, v in model.state_dict().items()}
    save_file(sd, os.path.join(out, "model.safetensors"), metadata={"format": "pt"})
    with open(os.path.join(out, "config.json"), "w") as f:
        json.dump(cfg, f, indent=2)
    tok.save_pretrained(out)
    try:  # deixa o repo autocontido (igual o inference.py do MetaDiffusion)
        shutil.copy(os.path.abspath(__file__), os.path.join(out, "kiyo_diffusion.py"))
    except (shutil.SameFileError, OSError):
        pass


def load_model(path, device="cpu", dtype=torch.float32):
    from transformers import AutoTokenizer
    from transformers.utils import logging as hf_logging
    hf_logging.set_verbosity_error()  # silencia o aviso de model_type desconhecido ao ler o tokenizer
    path = local_dir(path)
    cfg = json.load(open(os.path.join(path, "config.json")))
    model = KiyoDiffusion(cfg)
    sd = {k.removeprefix("model."): v for k, v in read_safetensors(path).items()}
    model.load_state_dict(sd, strict=True)
    return model.to(device=device, dtype=dtype), cfg, AutoTokenizer.from_pretrained(path)


# ============================================================================ convert
def convert(src_id, out, shift=True):
    from transformers import AutoTokenizer
    src = local_dir(src_id)
    hc = json.load(open(os.path.join(src, "config.json")))
    assert hc.get("model_type", "llama") == "llama", "so checkpoints Llama sao suportados"
    assert hc.get("rope_scaling") is None, "rope_scaling nao suportado neste script"
    rope_theta = hc.get("rope_theta") or (hc.get("rope_parameters") or {}).get("rope_theta", 10000.0)
    tok = AutoTokenizer.from_pretrained(src)
    tok.add_special_tokens({"mask_token": MASK_TOKEN})
    mask_id = tok.convert_tokens_to_ids(MASK_TOKEN)
    tied = hc.get("tie_word_embeddings", True)

    sd = {k.removeprefix("model."): v.float() for k, v in read_safetensors(src).items()}
    if tied:
        sd.pop("lm_head.weight", None)
    V_old, H = sd["embed_tokens.weight"].shape
    V = max(V_old, mask_id + 1)
    for key in ["embed_tokens.weight"] + ([] if tied else ["lm_head.weight"]):
        E = sd[key]
        newE = torch.empty(V, H)
        newE[:V_old] = E
        newE[V_old:] = E.mean(0)  # [MASK] comeca como o embedding medio
        sd[key] = newE

    cfg = {k: hc[k] for k in ("hidden_size", "intermediate_size", "num_hidden_layers", "num_attention_heads",
                              "num_key_value_heads", "rms_norm_eps", "max_position_embeddings")}
    cfg["rope_theta"] = float(rope_theta)
    cfg["head_dim"] = hc.get("head_dim") or cfg["hidden_size"] // cfg["num_attention_heads"]
    cfg.update(architectures=["KiyoDiffusion"], model_type="kiyo_diffusion", vocab_size=V,
               mask_token_id=mask_id, eos_token_id=tok.eos_token_id,
               bos_token_id=tok.bos_token_id if tok.bos_token_id is not None else tok.eos_token_id,
               tie_word_embeddings=tied, shift_logits=shift, base_model=src_id)
    model = KiyoDiffusion(cfg)
    model.load_state_dict(sd, strict=True)
    save_pretrained(model, cfg, tok, out)
    n = sum(p.numel() for p in model.parameters())
    print(f"ok -> {out} | {n / 1e6:.2f}M params | vocab {V} | mask_id {mask_id} | "
          f"rope_theta {rope_theta} | shift {shift}", flush=True)


def cmd_convert(a):
    convert(a.src, a.out, shift=not a.no_shift)


# ============================================================================ check
CHECK_TEXT = ("Photosynthesis is the process by which green plants, algae and some bacteria convert light "
              "energy into chemical energy. During this process, carbon dioxide and water are combined to "
              "form glucose, and oxygen is released as a by-product. The reactions take place inside the "
              "chloroplasts, where the pigment chlorophyll absorbs mostly blue and red light. Because almost "
              "every food chain on Earth begins with organisms that photosynthesize, the process is essential "
              "for life as we know it.")


@torch.no_grad()
def cmd_check(a):
    """Modo causal do modelo convertido tem que bater com o Kiyo original (valida pesos/RoPE)."""
    from transformers import AutoModelForCausalLM
    model, cfg, tok = load_model(a.model)
    model.eval()
    ref = AutoModelForCausalLM.from_pretrained(local_dir(a.src), dtype=torch.float32).eval()
    ids = torch.tensor([tok(a.text or CHECK_TEXT, add_special_tokens=False).input_ids[:512]])
    Vref = ref.config.vocab_size
    ours = model(ids, causal=True)
    ours = (ours[:, 1:] if cfg.get("shift_logits", True) else ours[:, :-1])[..., :Vref]
    theirs = ref(ids).logits[:, :-1]
    print(f"max |dif logits| vs Kiyo original: {(ours - theirs).abs().max().item():.2e}  (esperado < 1e-3)")
    print(f"loss AR do Kiyo neste texto: {F.cross_entropy(theirs[0], ids[0, 1:]).item():.3f}")
    torch.manual_seed(0)
    for r in (0.15, 0.5, 0.9):  # CE media nos tokens mascarados, atencao bidirecional, antes de treinar
        m = torch.rand(ids.shape) < r
        m[:, 0] = False
        logits = model(torch.where(m, cfg["mask_token_id"], ids))
        print(f"difusao (sem treino) | {int(r * 100)}% mascarado -> CE {F.cross_entropy(logits[m], ids[m]).item():.3f}")


# ============================================================================ dados
def _texts(ex):
    if ex.get("text"):
        return [ex["text"]]
    if ex.get("content"):
        return [ex["content"]]
    if ex.get("files"):  # The Stack v3: cada linha e um repositorio inteiro
        return [f["content"] for f in ex["files"] if f.get("content") and not f.get("is_vendor")]
    return []


class PackedStream(torch.utils.data.IterableDataset):
    """Streaming do HF -> tokeniza -> empacota em seqs de seq_len (docs separados por EOS).
    Fontes: repo[:config][@peso] (ex: HuggingFaceTB/finemath:finemath-3plus@0.1) ou arquivo local
    .txt/.jsonl/.parquet (aceita glob). O peso e a fracao de TOKENS de cada fonte. Os arquivos de
    cada fonte sao divididos entre GPUs x workers, entao ninguem baixa dado repetido.
    O shuffle e em nivel de DOCUMENTO (pool de textos ja extraidos): um buffer de linhas estoura a RAM
    na Stack v3, onde cada linha e um repositorio inteiro com metadados."""

    def __init__(self, specs, tok_path, seq_len, seed, rank, world, pool_size=2000):
        super().__init__()
        self.specs, self.tok_path, self.seq_len = specs, tok_path, seq_len
        self.seed, self.rank, self.world, self.pool_size = seed, rank, world, pool_size
        self._file_cache = {}

    def _files(self, name):
        """Lista de arquivos de uma fonte: glob local, ou os data_files (hf://) do dataset no Hub."""
        if name not in self._file_cache:
            if any(x in name.lower() for x in (".txt", ".json", ".parquet")):
                files = sorted(glob.glob(name)) or [name]
            else:
                from datasets import load_dataset_builder
                repo, _, conf = name.partition(":")
                files = sorted(load_dataset_builder(repo, conf or None).config.data_files["train"])
            self._file_cache[name] = files
        return self._file_cache[name]

    @staticmethod
    def _rows(path):
        if path.endswith(".parquet"):
            # pyarrow direto, em lotes pequenos: o datasets descomprime row groups inteiros
            # (ate 360 MB na Stack v3) e estoura a RAM
            import fsspec
            import pyarrow.parquet as pq
            with fsspec.open(path, "rb", block_size=8 << 20) as fh:
                pf = pq.ParquetFile(fh, pre_buffer=False, buffer_size=8 << 20)
                cols = [c for c in ("text", "content", "files") if c in pf.schema_arrow.names]
                for batch in pf.iter_batches(batch_size=32, columns=cols, use_threads=False):
                    yield from batch.to_pylist()
        else:  # .jsonl(.zst/.gz) e .txt (DCLM, arquivos locais), lidos direto linha a linha.
            # Nao usar o datasets aqui: dentro de um worker do DataLoader ele "desliga" os workers
            # extras quando o arquivo tem 1 shard, e o worker fica varrendo arquivos vazios.
            import fsspec
            with fsspec.open(path, "rt", compression="infer", encoding="utf-8", errors="ignore",
                             block_size=8 << 20) as fh:
                for line in fh:
                    if ".txt" in path:
                        yield {"text": line}
                    elif line.strip():
                        try:
                            yield json.loads(line)
                        except json.JSONDecodeError:
                            continue

    def _open(self, name, shard, n_shards, epoch):
        files = list(self._files(name))
        random.Random(self.seed + epoch).shuffle(files)  # ordem aleatoria de arquivos a cada epoca
        if len(files) >= n_shards:  # arquivos inteiros por worker: ninguem baixa dado repetido
            for path in files[shard::n_shards]:
                yield from self._rows(path)
        else:  # poucos arquivos (ex: 1 arquivo local): 1 a cada n linhas
            k = 0
            for path in files:
                for row in self._rows(path):
                    if k % n_shards == shard:
                        yield row
                    k += 1

    def __iter__(self):
        from transformers import AutoTokenizer
        from transformers.utils import logging as hf_logging
        hf_logging.set_verbosity_error()
        tok = AutoTokenizer.from_pretrained(self.tok_path)
        eos = tok.eos_token_id
        wi = torch.utils.data.get_worker_info()
        nw, wid = (wi.num_workers, wi.id) if wi else (1, 0)
        shard, n_shards = self.rank * nw + wid, self.world * nw
        srcs = [(s.partition("@")[0], float(s.partition("@")[2] or 1.0)) for s in self.specs]
        its = [self._open(name, shard, n_shards, 0) for name, _ in srcs]
        epoch, rows = [0] * len(srcs), [0] * len(srcs)
        bufs = [array.array("i") for _ in srcs]
        pools = [[] for _ in srcs]
        sent = [0] * len(srcs)
        rng = random.Random(self.seed * 1000 + shard)

        def next_doc(i):
            pool, empty = pools[i], 0
            while len(pool) < self.pool_size:  # reabastece o pool de documentos da fonte i
                try:
                    ex = next(its[i])
                except StopIteration:  # fonte acabou: reembaralha e recomeca
                    if rows[i] == 0:
                        raise RuntimeError(f"fonte vazia neste shard: {srcs[i][0]}")
                    epoch[i], rows[i] = epoch[i] + 1, 0
                    its[i] = self._open(srcs[i][0], shard, n_shards, epoch[i])
                    continue
                rows[i] += 1
                got = _texts(ex)
                empty = 0 if got else empty + 1
                if empty > 1000:
                    raise RuntimeError(f"{srcs[i][0]}: sem coluna 'text', 'content' ou 'files'")
                pool.extend(got)
            j = rng.randrange(len(pool))  # tira um documento aleatorio do pool
            pool[j], pool[-1] = pool[-1], pool[j]
            return pool.pop()

        while True:
            # escolhe a fonte mais atrasada em relacao ao peso -> mix por tokens, nao por documentos
            i = min(range(len(srcs)), key=lambda j: sent[j] / srcs[j][1])
            while len(bufs[i]) < self.seq_len:
                docs = [next_doc(i) for _ in range(16)]
                # split_special_tokens: um "<|mask|>" ou "<|endoftext|>" literal no texto vira texto normal
                for ids in tok(docs, add_special_tokens=False, split_special_tokens=True)["input_ids"]:
                    bufs[i].extend(ids)
                    bufs[i].append(eos)
            block = torch.frombuffer(bufs[i][:self.seq_len], dtype=torch.int32).long()
            del bufs[i][:self.seq_len]
            sent[i] += 1
            yield block


# ============================================================================ treino
def noise(x0, mask_id, eps=1e-3):
    B, L = x0.shape
    # um t por sequencia, estratificado no batch (menos variancia que t sorteado independente)
    u = (torch.rand(1, device=x0.device) + torch.arange(B, device=x0.device) / B) % 1.0
    p = ((1 - eps) * u + eps)[:, None].expand(B, L)
    m = torch.rand(B, L, device=x0.device) < p
    m[:, 0] = False  # ancora: o shift precisa de um 1o token sempre visivel
    return torch.where(m, mask_id, x0), m, p


def anneal_mask(L, frac, device):
    """frac=0 -> causal puro; frac=1 -> bidirecional (None = kernel flash sem mascara)."""
    if frac >= 1.0:
        return None
    causal = torch.ones(L, L, dtype=torch.bool, device=device).tril()
    return causal | (torch.rand(L, L, device=device) < frac)


def lr_at(step, a):
    if step < a.warmup:
        return a.lr * (step + 1) / a.warmup
    prog = min(1.0, (step - a.warmup) / max(1, a.steps - a.warmup))
    return a.min_lr + 0.5 * (a.lr - a.min_lr) * (1 + math.cos(math.pi * prog))


def pick_dtype(name="auto"):
    """bf16 nas GPUs Ampere+ (3060, L4, A100...), fp16 nas mais antigas (T4, sem bf16), fp32 na CPU."""
    if name != "auto":
        return {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[name]
    if not torch.cuda.is_available():
        return torch.float32
    return torch.bfloat16 if torch.cuda.get_device_capability()[0] >= 8 else torch.float16


def maybe_relaunch(a):
    """Se tem varias GPUs e nao veio do torchrun, relanca o proprio script com torchrun."""
    if "LOCAL_RANK" in os.environ or not torch.cuda.is_available():
        return
    n = torch.cuda.device_count() if a.nproc <= 0 else a.nproc
    if n > 1:
        print(f"relancando com torchrun em {n} GPUs", flush=True)
        os.execv(sys.executable, [sys.executable, "-m", "torch.distributed.run", "--standalone",
                                  f"--nproc_per_node={n}", os.path.abspath(__file__), *sys.argv[1:]])


def resolve_init(a, master, ddp):
    """--init pode ser o Kiyo AR (converte na hora pra <out>/init) ou um modelo de difusao ja salvo."""
    target = os.path.join(a.out, "init")

    def is_diffusion():
        with open(os.path.join(local_dir(a.init), "config.json")) as f:
            return json.load(f).get("model_type") == "kiyo_diffusion"

    if master and not is_diffusion() and not os.path.exists(os.path.join(target, "model.safetensors")):
        convert(a.init, target, shift=not a.no_shift)
    if ddp:
        torch.distributed.barrier()
    return local_dir(a.init) if is_diffusion() else target


def cmd_train(a):
    maybe_relaunch(a)
    ddp = int(os.environ.get("WORLD_SIZE", "1")) > 1
    cuda = torch.cuda.is_available()
    if ddp:
        import torch.distributed as dist
        from torch.nn.parallel import DistributedDataParallel as DDP
        if cuda and "microsoft" in platform.uname().release.lower():
            os.environ.setdefault("NCCL_P2P_DISABLE", "1")  # NCCL no WSL2 costuma travar com P2P ligado
        dist.init_process_group("nccl" if cuda else "gloo")
        rank, world, local = dist.get_rank(), dist.get_world_size(), int(os.environ["LOCAL_RANK"])
        device = torch.device("cuda", local) if cuda else torch.device("cpu")
        if cuda:
            torch.cuda.set_device(local)
    else:
        rank, world, local = 0, 1, 0
        device = torch.device("cuda" if cuda else "cpu")
    master = rank == 0
    log = (lambda *x: print(*x, flush=True)) if master else (lambda *x: None)
    torch.manual_seed(a.seed + rank)
    random.seed(a.seed + rank)
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True

    init = resolve_init(a, master, ddp)
    model, cfg, tok = load_model(init, device, torch.float32)
    model.grad_ckpt = a.grad_ckpt
    mask_id = cfg["mask_token_id"]
    if a.seq_len > cfg["max_position_embeddings"]:
        log(f"aviso: --seq-len {a.seq_len} > max_position_embeddings {cfg['max_position_embeddings']}")
    amp_dtype = pick_dtype(a.dtype) if cuda else torch.float32
    amp = (lambda: torch.autocast("cuda", dtype=amp_dtype)) if amp_dtype != torch.float32 else nullcontext
    scaler = torch.amp.GradScaler("cuda", enabled=cuda and amp_dtype == torch.float16)  # so fp16 (T4)
    params = list(model.parameters())
    opt = torch.optim.AdamW([{"params": [p for p in params if p.dim() >= 2], "weight_decay": a.wd},
                            {"params": [p for p in params if p.dim() < 2], "weight_decay": 0.0}],
                           lr=a.lr, betas=(0.9, 0.95), eps=1e-8, fused=cuda)
    step, ck_dir = 0, os.path.join(a.out, "checkpoints")
    cks = sorted(glob.glob(os.path.join(ck_dir, "step_*.pt")))
    if cks:
        s = torch.load(cks[-1], map_location=device, weights_only=False)
        model.load_state_dict(s["model"])
        opt.load_state_dict(s["opt"])
        if scaler.is_enabled() and "scaler" in s:
            scaler.load_state_dict(s["scaler"])
        step = s["step"]
        del s
        log(f"retomando de {cks[-1]} (step {step})")

    if a.compile:
        for blk in model.layers:  # compila bloco a bloco, in-place (nomes do state_dict nao mudam)
            blk.compile()
    ddp_model = DDP(model, device_ids=[local] if cuda else None) if ddp else None
    fwd = ddp_model or model

    per_micro = a.micro_bs * a.seq_len * world
    accum = max(1, round(a.batch_tokens / per_micro))
    tok_step = per_micro * accum
    a.steps = a.steps or math.ceil(a.tokens / tok_step)  # mesmo batch global em qualquer maquina
    anneal_steps = int(a.anneal_frac * a.steps)
    log(f"{sum(p.numel() for p in params) / 1e6:.2f}M params | rope_theta {cfg['rope_theta']} | "
        f"shift {cfg.get('shift_logits', True)} | {str(amp_dtype).removeprefix('torch.')}\n"
        f"{a.micro_bs} seq x {a.seq_len} tok x {world} GPU x {accum} accum = {tok_step:,} tokens/step | "
        f"{a.steps} steps = {tok_step * a.steps / 1e9:.2f}B tokens | annealing {anneal_steps} steps")

    def make_iter(seed):
        ds = PackedStream(a.data, init, a.seq_len, seed, rank, world)
        return iter(torch.utils.data.DataLoader(ds, batch_size=a.micro_bs, num_workers=a.workers,
                                                pin_memory=cuda, prefetch_factor=4 if a.workers > 0 else None))

    data_seed = a.seed + step  # ordem de dados diferente a cada retomada
    it = make_iter(data_seed)

    def next_batch():
        nonlocal it, data_seed
        for tries in range(10):
            try:
                return next(it)
            except Exception as e:  # soluco de rede no streaming: espera e reabre
                print(f"[rank {rank}] erro nos dados: {e!r} -> reabrindo o stream", flush=True)
                time.sleep(min(60, 5 * (tries + 1)))
                data_seed += 1000
                it = make_iter(data_seed)
        raise RuntimeError("o stream de dados falhou 10 vezes seguidas")

    model.train()
    stats = torch.zeros(4, device=device)  # soma nelbo, n micro-batches, soma CE, n mascarados
    t0 = t_start = last_save = time.time()
    step0 = last_log = step
    while step < a.steps:
        lr = lr_at(step, a)
        for g in opt.param_groups:
            g["lr"] = lr
        frac = 1.0 if anneal_steps <= 0 else min(1.0, step / anneal_steps)
        for micro in range(accum):
            x0 = next_batch().to(device, non_blocking=True)
            if random.random() < a.var_len_prob:  # as vezes seqs curtas (LLaDA usa 1%)
                x0 = x0[:, :random.randint(16, a.seq_len)]
            xt, m, p = noise(x0, mask_id, a.eps)
            am = anneal_mask(x0.shape[1], frac, device)
            if ddp_model is not None:  # so sincroniza gradiente no ultimo micro-batch
                ddp_model.require_backward_grad_sync = micro == accum - 1
            with amp():
                nelbo, ce_sum, n_masked = fwd(xt, targets=x0, loss_mask=m, p_mask=p, attn_mask=am)
            scaler.scale(nelbo / accum).backward()
            stats += torch.stack([nelbo.detach().float(), torch.ones((), device=device),
                                  ce_sum.float(), n_masked.float()])
        scaler.unscale_(opt)
        gn = torch.nn.utils.clip_grad_norm_(params, a.clip)
        scaler.step(opt)
        scaler.update()
        opt.zero_grad(set_to_none=True)
        step += 1

        if step % a.log_every == 0 or step == a.steps:
            if ddp:
                dist.all_reduce(stats)
            dt = time.time() - t0
            eta = (time.time() - t_start) / (step - step0) * (a.steps - step)
            log(f"step {step:>7} | nelbo {stats[0] / stats[1]:.4f} | ce_mask {stats[2] / stats[3].clamp(min=1):.4f} | "
                f"lr {lr:.2e} | gnorm {float(gn):.2f} | attn_bidir {frac:.2f} | "
                f"{(step - last_log) * tok_step / dt:,.0f} tok/s | eta {eta / 3600:.1f}h")
            stats.zero_()
            t0, last_log = time.time(), step

        if master and a.sample_every and step % a.sample_every == 0:
            model.eval()
            with amp():
                txt = generate(model, tok, a.sample_prompt, gen_len=48, block_len=16, temperature=0.0,
                               rep_penalty=1.1)
            log(f"[amostra step {step}] {txt!r}")
            model.train()

        want = step % a.save_every == 0 or step == a.steps or \
            (a.save_minutes > 0 and time.time() - last_save > a.save_minutes * 60)
        if ddp:
            flag = torch.tensor([int(want)], device=device)
            dist.broadcast(flag, src=0)
            want = bool(flag.item())
        if want:
            last_save = time.time()
            if master:
                os.makedirs(ck_dir, exist_ok=True)
                path = os.path.join(ck_dir, f"step_{step:08d}.pt")
                torch.save({"model": model.state_dict(), "opt": opt.state_dict(), "step": step,
                            "scaler": scaler.state_dict(), "args": vars(a)}, path + ".tmp")
                os.replace(path + ".tmp", path)
                for old in sorted(glob.glob(os.path.join(ck_dir, "step_*.pt")))[:-a.keep_last]:
                    os.remove(old)
                done = step == a.steps
                out = os.path.join(a.out, "final" if done else "latest")
                save_pretrained(model, cfg, tok, out)
                if done:
                    write_card(out, cfg, a.push_to)
                log(f"checkpoint salvo: {path}")
            if ddp:
                dist.barrier()

    if master and a.push_to:
        from huggingface_hub import HfApi
        api = HfApi()
        api.create_repo(a.push_to, private=not a.public, exist_ok=True)
        api.upload_folder(folder_path=os.path.join(a.out, "final"), repo_id=a.push_to)
        log(f"enviado para https://huggingface.co/{a.push_to}")
    if ddp:
        dist.barrier()
        dist.destroy_process_group()


# ============================================================================ geracao
@torch.no_grad()
def generate_ids(model, ids, gen_len=128, block_len=32, steps=None, temperature=0.6, rep_penalty=1.0,
                 stop_at_eos=True, on_step=None):
    """Sampler do LLaDA: blocos da esquerda p/ direita + revela os tokens mais confiantes.
    Qualquer [MASK] dentro de `ids` tambem e preenchido (infilling)."""
    c = model.cfg
    mask_id, eos, dev = c["mask_token_id"], c["eos_token_id"], next(model.parameters()).device
    x = torch.tensor([list(ids) + [mask_id] * gen_len], dtype=torch.long, device=dev)
    P = len(ids)
    todo = (x[0] == mask_id).nonzero().squeeze(1)
    if todo.numel() == 0:
        return x[0]
    blocks = todo.split(block_len)
    spb = max(1, math.ceil((steps or todo.numel()) / len(blocks)))  # steps=None -> 1 token por step
    for b, blk in enumerate(blocks):
        for i in range(spb):
            left = blk[x[0, blk] == mask_id]
            if left.numel() == 0:
                break
            k = math.ceil(left.numel() / (spb - i))
            logits = model(x)[0, left].float()
            logits[:, mask_id] = -float("inf")
            if rep_penalty != 1.0:
                seen = x[0][x[0] != mask_id].unique()
                sl = logits[:, seen]
                logits[:, seen] = torch.where(sl > 0, sl / rep_penalty, sl * rep_penalty)
            if temperature > 0:
                probs = F.softmax(logits / temperature, dim=-1)
                pred = torch.multinomial(probs, 1).squeeze(-1)
            else:
                probs = F.softmax(logits, dim=-1)
                pred = probs.argmax(-1)
            conf = probs.gather(-1, pred[:, None]).squeeze(-1)
            pick = conf.topk(k).indices  # revela so os k mais confiantes, o resto continua [MASK]
            x[0, left[pick]] = pred[pick]
            if on_step:
                on_step(x[0], b, len(blocks), i, spb)
        if stop_at_eos and gen_len > 0 and (x[0, P:] == eos).any():
            break
    return x[0]


def generate(model, tok, prompt, gen_len=128, block_len=32, steps=None, temperature=0.6, rep_penalty=1.0,
             watch=False, bos=True):
    c = model.cfg
    ids = tok(prompt, add_special_tokens=False).input_ids  # "<|mask|>" no prompt vira [MASK] -> infilling
    start = 0
    if bos or not ids:  # o shift precisa de >= 1 token visivel antes da 1a mascara
        ids, start = [c.get("bos_token_id", c["eos_token_id"])] + ids, 1
    P = len(ids)

    def show(x):
        return tok.decode(x[start:].tolist()).replace(MASK_TOKEN, "_")

    def on_step(x, b, nb, i, spb):
        print(f"[bloco {b + 1}/{nb} step {i + 1}/{spb}] {show(x)}\n", flush=True)

    out = generate_ids(model, ids, gen_len, block_len, steps, temperature, rep_penalty,
                       on_step=on_step if watch else None).tolist()
    end = len(out)
    if gen_len > 0 and c["eos_token_id"] in out[P:]:
        end = P + out[P:].index(c["eos_token_id"])
    return tok.decode(out[start:end]).replace(MASK_TOKEN, "_")


@torch.no_grad()
def score(model, tok, context, continuation, n_mc=32):
    """Log-verossimilhanca aproximada (ELBO por Monte Carlo, como no LLaDA) de `continuation`
    dado `context`. Serve p/ benchmarks de multipla escolha: maior = alternativa mais provavel."""
    c = model.cfg
    mask_id, dev = c["mask_token_id"], next(model.parameters()).device
    ctx = [c.get("bos_token_id", c["eos_token_id"])] + tok(context, add_special_tokens=False).input_ids
    cont = tok(continuation, add_special_tokens=False).input_ids
    P, Lc = len(ctx), len(cont)
    x0 = torch.tensor([ctx + cont], device=dev).repeat(n_mc, 1)
    l = torch.randint(1, Lc + 1, (n_mc,), device=dev)  # quantos tokens da resposta mascarar
    m = torch.zeros_like(x0, dtype=torch.bool)
    m[:, P:] = torch.rand(n_mc, Lc, device=dev).argsort(-1) < l[:, None]
    logp = torch.log_softmax(model(torch.where(m, mask_id, x0)).float(), -1)
    logp = logp.gather(-1, x0[..., None]).squeeze(-1)
    return ((logp * m).sum(-1) * (Lc / l.float())).mean().item()


def _load_for_inference(path):
    cuda = torch.cuda.is_available()
    model, cfg, tok = load_model(path, "cuda" if cuda else "cpu", pick_dtype())
    return model.eval(), tok


def cmd_generate(a):
    model, tok = _load_for_inference(a.model)
    torch.manual_seed(a.seed)
    print(generate(model, tok, a.prompt, a.gen_len, a.block_len, a.steps, a.temperature, a.rep_penalty,
                   a.watch, bos=not a.no_bos))


def cmd_score(a):
    model, tok = _load_for_inference(a.model)
    torch.manual_seed(a.seed)
    res = [(score(model, tok, a.context, ch, a.n_mc), ch) for ch in a.choices]
    for s, ch in sorted(res, reverse=True):
        print(f"{s:10.3f}  {ch!r}")


# ============================================================================ CLI
def main():
    ap = argparse.ArgumentParser(description="Kiyo -> LM de difusao mascarada")
    sub = ap.add_subparsers(dest="cmd", required=True)

    c = sub.add_parser("convert", help="converte o checkpoint AR (o train ja faz isso sozinho)")
    c.add_argument("--src", default="DedeProGames/Kiyo-135M")
    c.add_argument("--out", default="kiyo-diff-init")
    c.add_argument("--no-shift", action="store_true", help="sem shift (estilo MetaDiffusion/LLaDA)")

    k = sub.add_parser("check", help="compara o modo causal com o Kiyo original")
    k.add_argument("--model", default="kiyo-diff-init")
    k.add_argument("--src", default="DedeProGames/Kiyo-135M")
    k.add_argument("--text", default=None)

    t = sub.add_parser("train")
    t.add_argument("--init", default="DedeProGames/Kiyo-135M",
                   help="checkpoint AR (converte sozinho) ou modelo de difusao ja salvo")
    t.add_argument("--out", default="runs/kiyo-diff")
    t.add_argument("--data", nargs="+", default=KIYO_MIX, help="repo[:config][@peso] ou arquivos locais")
    t.add_argument("--seq-len", type=int, default=1024)
    t.add_argument("--micro-bs", type=int, default=4,
                   help="sequencias por GPU por micro-batch (OOM? baixa isso ou usa --grad-ckpt)")
    t.add_argument("--batch-tokens", type=int, default=131_072,
                   help="tokens por step, global (o accum se ajusta ao n de GPUs)")
    t.add_argument("--tokens", type=float, default=1e9, help="orcamento total (1e9 = 1B, 5e8 = 500M)")
    t.add_argument("--steps", type=int, default=None, help="sobrescreve --tokens")
    t.add_argument("--lr", type=float, default=3e-4)
    t.add_argument("--min-lr", type=float, default=3e-5)
    t.add_argument("--warmup", type=int, default=200)
    t.add_argument("--wd", type=float, default=0.1)
    t.add_argument("--clip", type=float, default=1.0)
    t.add_argument("--eps", type=float, default=1e-3, help="taxa minima de mascaramento")
    t.add_argument("--anneal-frac", type=float, default=0.2,
                   help="fracao dos steps com annealing causal->bidirecional (0 = bidirecional direto)")
    t.add_argument("--no-shift", action="store_true", help="so vale na conversao automatica")
    t.add_argument("--var-len-prob", type=float, default=0.01)
    t.add_argument("--workers", type=int, default=2, help="workers do dataloader por GPU")
    t.add_argument("--nproc", type=int, default=0, help="n de GPUs (0 = todas)")
    t.add_argument("--log-every", type=int, default=10)
    t.add_argument("--save-every", type=int, default=250)
    t.add_argument("--save-minutes", type=float, default=30, help="salva tambem a cada N min (0 = off)")
    t.add_argument("--keep-last", type=int, default=3, help="cada checkpoint .pt tem ~1,6 GB")
    t.add_argument("--sample-every", type=int, default=250, help="0 = sem amostras no log")
    t.add_argument("--dtype", default="auto", choices=["auto", "bf16", "fp16", "fp32"],
                   help="auto: bf16 em Ampere+ (3060), fp16 na T4")
    t.add_argument("--sample-prompt", default="The capital of France is")
    t.add_argument("--grad-ckpt", action="store_true")
    t.add_argument("--compile", action="store_true")
    t.add_argument("--seed", type=int, default=0)
    t.add_argument("--push-to", default=None, help="ex: DedeProGames/Kiyo-Diffusion-135M")
    t.add_argument("--public", action="store_true", help="repo publico (padrao: privado)")

    g = sub.add_parser("generate")
    g.add_argument("--model", required=True)
    g.add_argument("--prompt", default="The cat sat on the")
    g.add_argument("--gen-len", type=int, default=128, help="tokens novos (0 = so preencher os [MASK] do prompt)")
    g.add_argument("--block-len", type=int, default=32, help="= gen-len p/ difusao pura, sem blocos")
    g.add_argument("--steps", type=int, default=None, help="padrao = 1 token por step")
    g.add_argument("--temperature", type=float, default=0.6)
    g.add_argument("--rep-penalty", type=float, default=1.1)
    g.add_argument("--no-bos", action="store_true")
    g.add_argument("--seed", type=int, default=0)
    g.add_argument("--watch", action="store_true", help="mostra o texto sendo 'revelado'")

    s = sub.add_parser("score", help="log-verossimilhanca de alternativas (multipla escolha)")
    s.add_argument("--model", required=True)
    s.add_argument("--context", required=True)
    s.add_argument("--choices", nargs="+", required=True)
    s.add_argument("--n-mc", type=int, default=32)
    s.add_argument("--seed", type=int, default=0)

    a = ap.parse_args()
    {"convert": cmd_convert, "check": cmd_check, "train": cmd_train, "generate": cmd_generate,
     "score": cmd_score}[a.cmd](a)


if __name__ == "__main__":
    main()