"""Adaptación a difusión por enmascaramiento del Qwen3.5-0.8B abliterado.

Basado en summerMC/dqwen (recipe LLaDA/Dream). Entrena LoRA + rev_gates + mask_embed
sobre los pesos abliterados y guarda un checkpoint listo para publicar.
"""
import modal, sys, os

app = modal.App("qwen35-diffusion")
vol_hf = modal.Volume.from_name("qwen35-hf-cache", create_if_missing=True)
vol_out = modal.Volume.from_name("qwen35-diffusion-out", create_if_missing=True)

image = (
    modal.Image.debian_slim(python_version="3.11")
    .pip_install("torch", "transformers", "peft", "datasets", "accelerate",
                 "safetensors", "sentencepiece", "huggingface_hub", "numpy",
                 "flash-linear-attention", "triton")
    .add_local_dir("/root/alblit/diffusion/dqwen", "/root/dqwen")
)

BASE = "OpceanAI/Qwen3.5-0.8B-abliterated"
OUT = "/out"


@app.function(image=image, gpu="H200", timeout=7200,
              volumes={"/cache": vol_hf, "/out": vol_out},
              env={"HF_HOME": "/cache", "TOKENIZERS_PARALLELISM": "false"})
def train(smoke: bool = False, steps: int = 4000, batch: int = 8, accum: int = 4,
          max_len: int = 512, max_prompt: int = 256, max_steps_override: int = 0):
    sys.path.insert(0, "/root")
    import torch
    from transformers import AutoModelForImageTextToText, AutoTokenizer
    from dqwen.core import DiffusionLM, patch_bidirectional, rev_gate_state, diffusion_generate
    from dqwen.data import iter_rows, build_dataset
    from dqwen.train import Stage1Cfg, train_stage1, LORA_TARGETS
    from peft import LoraConfig, get_peft_model

    torch.set_grad_enabled(True)
    dev = "cuda"
    print("GPU:", torch.cuda.get_device_name(0), flush=True)

    tok = AutoTokenizer.from_pretrained(BASE)
    lm = AutoModelForImageTextToText.from_pretrained(BASE, dtype=torch.bfloat16, device_map=dev,
                                                     attn_implementation="sdpa")
    lm.config.use_cache = False

    cfg_lora = LoraConfig(r=32, lora_alpha=64, lora_dropout=0.0,
                          target_modules=LORA_TARGETS, bias="none")
    lm = get_peft_model(lm, cfg_lora)

    stats = patch_bidirectional(lm, True)
    gates = [n for n, _ in lm.named_parameters() if n.endswith("rev_gate")]
    print("[patch]", stats, "| rev_gates:", len(gates), flush=True)
    for n, p in lm.named_parameters():
        if n.endswith("rev_gate"):
            p.requires_grad_(True)

    mask_id = tok.convert_tokens_to_ids("<|fim_pad|>")
    dlm = DiffusionLM(lm, mask_id=mask_id)
    dlm.mask_embed.requires_grad_(True)
    n_tr = sum(p.numel() for p in dlm.parameters() if p.requires_grad)
    print(f"[model] trainable {n_tr/1e6:.1f}M | mask id {mask_id}", flush=True)

    n_rows = 2000 if smoke else 60000
    rows = iter_rows("tatsu-lab/alpaca", n=n_rows, seed=0)
    ds = build_dataset(tok, rows, max_len=max_len, max_prompt=max_prompt,
                       eos_token="<|im_end|>", chat_kwargs={"enable_thinking": False})

    cfg = Stage1Cfg(
        max_steps=steps, batch_size=batch, grad_accum=accum, lr=2e-4,
        gate_lr=2e-3, mask_lr=1e-3, warmup=50, t_min=0.05, weighting="elbo",
        log_every=20, sample_every=500, save_every=1000, save_dir=f"{OUT}/ckpt",
        seed=0,
    )
    prompts = ["Explain why the sky is blue.", "Write a short poem about the sea.",
               "What is the capital of France?"]
    train_stage1(dlm, tok, ds, cfg, device=dev, sample_prompts=prompts)

    dlm.eval()
    for p_ in prompts[:1]:
        print("[sample]", repr(generate(dlm, tok, p_)), flush=True)

    print("[merge] merging LoRA", flush=True)
    merged = dlm.lm.merge_and_unload()
    gates = rev_gate_state(merged)
    for name, mod in list(merged.named_modules()):
        if hasattr(mod, "rev_gate"):
            delattr(mod, "rev_gate")
    os.makedirs(f"{OUT}/final", exist_ok=True)
    merged.save_pretrained(f"{OUT}/final", safe_serialization=True)
    tok.save_pretrained(f"{OUT}/final")
    torch.save({"mask_embed": dlm.mask_embed.detach().cpu(),
                "rev_gates": gates,
                "mask_id": mask_id,
                "base": BASE}, f"{OUT}/final/dllm_extras.pt")
    print(f"[saved] {OUT}/final", flush=True)

    return {"steps": steps, "tokens": steps * batch * accum * max_len}


def generate(dlm, tok, prompt, gen_len=96, steps=32, block=32):
    import torch
    from dqwen.core import diffusion_generate
    eos_id = tok.convert_tokens_to_ids("<|im_end|>")
    text = tok.apply_chat_template([{"role": "user", "content": prompt}], tokenize=False,
                                   add_generation_prompt=True, enable_thinking=False)
    ids = torch.tensor([tok(text, add_special_tokens=False)["input_ids"]], device="cuda")
    with torch.no_grad():
        out = diffusion_generate(dlm, ids, gen_len=gen_len, steps=steps, block_len=block, eos_id=eos_id)
    seq = out[0].tolist()
    if eos_id in seq:
        seq = seq[:seq.index(eos_id)]
    return tok.decode(seq, skip_special_tokens=True)


@app.local_entrypoint()
def main(smoke: bool = False, steps: int = 4000, batch: int = 8, accum: int = 4):
    print(train.remote(smoke=smoke, steps=steps, batch=batch, accum=accum))