#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import math
import re
import sys
from pathlib import Path

import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer


def find_repo_root(start: Path) -> Path:
    for parent in [start, *start.parents]:
        if (parent / "flowtext_lab").is_dir():
            return parent
    raise RuntimeError(f"Could not find repo root from {start}")


REPO_ROOT = find_repo_root(Path(__file__).resolve())
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from flowtext_lab.genppl import filter_generated_texts, summarize_token_diversity
from flowtext_lab.model import EndpointPredictor
from flowtext_lab.tokenization import BpeTextTokenizer


SPECIAL_RE = re.compile(r"\s+")


def strip_common_special(text: str) -> str:
    text = text.replace("[CLS]", " ").replace("[SEP]", " ").replace("[PAD]", " ")
    text = text.replace("<|endoftext|>", " ")
    return SPECIAL_RE.sub(" ", text).strip()


def write_raw_samples(path: Path, title: str, summary: dict, texts: list[str]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as f:
        f.write(title + "\n")
        f.write(json.dumps(summary, ensure_ascii=False, indent=2) + "\n\n")
        for i, text in enumerate(texts):
            f.write(f"===== sample {i:03d} =====\n")
            f.write(text + "\n\n")


def build_endpoint_model(ckpt: dict, tokenizer: BpeTextTokenizer, device: torch.device) -> EndpointPredictor:
    train_args = ckpt.get("args", {})
    token_vocab = tokenizer.vocab_size
    train_vocab = int(ckpt.get("train_vocab_size") or token_vocab + 1)
    train_vocab = max(train_vocab, token_vocab + 1)
    model = EndpointPredictor(
        vocab_size=train_vocab,
        max_len=int(train_args.get("max_len", 128)),
        d_model=int(train_args.get("d_model", 768)),
        n_heads=int(train_args.get("n_heads", 12)),
        n_layers=int(train_args.get("n_layers", 12)),
        dim_ff=int(train_args.get("dim_ff", 3072)),
        dropout=0.0,
        model_type=str(train_args.get("model_type", "ddit")),
        cond_dim=int(train_args.get("cond_dim", 128)),
        input_format=str(train_args.get("state_format", train_args.get("input_format", "prob"))),
    ).to(device)
    model.load_state_dict(ckpt["model"], strict=True)
    model.eval()
    return model


def log_linear_alpha(t: torch.Tensor, eps: float) -> torch.Tensor:
    return 1.0 - (1.0 - eps) * t


def sample_categorical(weights: torch.Tensor) -> torch.Tensor:
    flat = weights.to(torch.float64).clamp_min(0).reshape(-1, weights.size(-1))
    row_sum = flat.sum(dim=-1, keepdim=True)
    probs = flat / row_sum.clamp_min(1e-300)
    return torch.multinomial(probs, num_samples=1).view(*weights.shape[:-1])


def mdlm_log_probs(
    model: EndpointPredictor,
    x: torch.Tensor,
    sigma: torch.Tensor,
    attn: torch.Tensor,
    *,
    token_vocab_size: int,
    mask_id: int,
) -> torch.Tensor:
    logits = model(x, sigma, attn).float()
    logits[:, :, mask_id] = -1_000_000.0
    log_probs = logits - torch.logsumexp(logits, dim=-1, keepdim=True)

    # Official MDLM's subs parameterization makes already-unmasked tokens absorbing.
    unmasked = x != mask_id
    log_probs = log_probs.clone()
    log_probs[unmasked] = -1_000_000.0
    log_probs[unmasked, x[unmasked].clamp_max(token_vocab_size - 1)] = 0.0
    return log_probs


@torch.inference_mode()
def decode_mdlm_ddpm_cache(
    model: EndpointPredictor,
    tokenizer: BpeTextTokenizer,
    *,
    n_samples: int,
    batch_size: int,
    max_len: int,
    steps: int,
    eps: float,
    final: str,
    time_conditioned: bool,
    seed: int,
    device: torch.device,
) -> tuple[list[list[int]], list[str], dict]:
    torch.manual_seed(seed)
    token_vocab_size = tokenizer.vocab_size
    mask_id = token_vocab_size
    all_ids: list[list[int]] = []
    all_texts: list[str] = []
    remaining = n_samples
    while remaining > 0:
        bs = min(batch_size, remaining)
        x = torch.full((bs, max_len), mask_id, dtype=torch.long, device=device)
        attn = torch.ones((bs, max_len), dtype=torch.bool, device=device)
        timesteps = torch.linspace(1.0, eps, steps + 1, device=device)
        dt = (1.0 - eps) / max(steps, 1)
        p_x0_cache: torch.Tensor | None = None

        for step in range(steps):
            t = torch.full((bs, 1), float(timesteps[step].item()), dtype=torch.float32, device=device)
            alpha_t = log_linear_alpha(t, eps)
            alpha_s = log_linear_alpha(t - dt, eps)
            if p_x0_cache is None:
                sigma = -torch.log(alpha_t.clamp_min(1e-8)).squeeze(-1)
                p_x0_cache = mdlm_log_probs(
                    model,
                    x,
                    sigma,
                    attn,
                    token_vocab_size=token_vocab_size,
                    mask_id=mask_id,
                ).exp()

            q_xs = p_x0_cache * (alpha_s - alpha_t).clamp_min(0)[:, :, None]
            q_xs[:, :, mask_id] = (1.0 - alpha_s).clamp_min(0).squeeze(-1)[:, None]
            proposal = sample_categorical(q_xs)
            x_next = torch.where(x != mask_id, x, proposal)

            # Official ddpm_cache reuses p(x0|x_t) until x changes; for our
            # time-conditioned checkpoint, sigma changes every step, so cache must reset.
            if time_conditioned or not torch.equal(x_next, x):
                p_x0_cache = None
            x = x_next

        if final != "none" and (x == mask_id).any():
            t0 = torch.full((bs, 1), float(timesteps[-1].item()), dtype=torch.float32, device=device)
            alpha_t0 = log_linear_alpha(t0, eps)
            sigma0 = -torch.log(alpha_t0.clamp_min(1e-8)).squeeze(-1)
            log_probs = mdlm_log_probs(
                model,
                x,
                sigma0,
                attn,
                token_vocab_size=token_vocab_size,
                mask_id=mask_id,
            )
            if final == "greedy":
                proposal = log_probs[:, :, :token_vocab_size].argmax(dim=-1)
            elif final == "ancestral":
                proposal = sample_categorical(log_probs.exp())
            else:
                raise ValueError(f"unknown final={final!r}")
            x = torch.where(x != mask_id, x, proposal)

        x = x.clamp_max(token_vocab_size - 1)
        rows = x.detach().cpu().tolist()
        all_ids.extend(rows)
        all_texts.extend(tokenizer.decode(row, stop_at_eos=False, skip_special_tokens=False) for row in rows)
        remaining -= bs
        print(f"[mdlm ddpm_cache final={final}] generated {n_samples - remaining}/{n_samples}", flush=True)

    decode = {
        "kind": "mdlm",
        "steps": steps,
        "decode_rule": "ddpm_cache_absorbing_state",
        "final": final,
        "eps": eps,
        "time_conditioned": time_conditioned,
        "n_samples": n_samples,
        "seed": seed,
    }
    return all_ids, all_texts, decode


@torch.no_grad()
def score_with_loaded(
    texts: list[str],
    scorer,
    scorer_tok,
    *,
    batch_size: int,
    max_length: int,
    device: torch.device,
) -> dict:
    total_nll = 0.0
    total_tokens = 0
    skipped = 0
    for start in range(0, len(texts), batch_size):
        batch_texts = texts[start : start + batch_size]
        enc = scorer_tok(
            batch_texts,
            return_tensors="pt",
            return_token_type_ids=False,
            return_attention_mask=True,
            padding=True,
            truncation=True,
            max_length=max_length,
        ).to(device)
        input_ids = enc["input_ids"]
        attention_mask = enc["attention_mask"]
        if input_ids.size(1) < 2:
            skipped += len(batch_texts)
            continue
        logits = scorer(input_ids=input_ids, attention_mask=attention_mask).logits.transpose(-1, -2)
        token_nll = F.cross_entropy(logits[..., :-1].float(), input_ids[..., 1:], reduction="none")
        if scorer_tok.eos_token_id is not None:
            first_eos = (input_ids == scorer_tok.eos_token_id).cumsum(-1) == 1
            token_mask = input_ids != scorer_tok.eos_token_id
            shift_mask = (first_eos[..., 1:] | token_mask[..., 1:]).contiguous()
        else:
            shift_mask = attention_mask[..., 1:].contiguous().bool()
        total_nll += float(token_nll[shift_mask].sum().detach().cpu())
        total_tokens += int(shift_mask.sum().detach().cpu())
    nll = total_nll / max(total_tokens, 1)
    return {
        "ppl": float(math.exp(min(nll, 50.0))),
        "nll_per_token": float(nll),
        "tokens": total_tokens,
        "kept_samples": len(texts),
        "total_samples": len(texts),
        "empty_rate": 0.0,
        "skipped_samples": skipped,
    }


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--checkpoint", required=True)
    p.add_argument("--tokenizer_path", required=True)
    p.add_argument("--scorer", required=True)
    p.add_argument("--out_dir", required=True)
    p.add_argument("--name", default="mdlm_ddpm_cache")
    p.add_argument("--n_samples", type=int, default=256)
    p.add_argument("--max_len", type=int, default=128)
    p.add_argument("--steps", type=int, default=1024)
    p.add_argument("--decode_batch", type=int, default=32)
    p.add_argument("--score_batch", type=int, default=8)
    p.add_argument("--score_max_length", type=int, default=256)
    p.add_argument("--seed", type=int, default=20260506)
    p.add_argument("--eps", type=float, default=1e-5)
    p.add_argument("--final", choices=["ancestral", "greedy", "none"], default="ancestral")
    p.add_argument("--no_time_conditioned", action="store_true")
    args = p.parse_args()

    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    tokenizer = BpeTextTokenizer.from_file(args.tokenizer_path)

    ckpt = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
    model = build_endpoint_model(ckpt, tokenizer, device)
    ids, raw_texts, decode = decode_mdlm_ddpm_cache(
        model,
        tokenizer,
        n_samples=args.n_samples,
        batch_size=args.decode_batch,
        max_len=args.max_len,
        steps=args.steps,
        eps=args.eps,
        final=args.final,
        time_conditioned=not args.no_time_conditioned,
        seed=args.seed,
        device=device,
    )
    del model
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    scorer_tok = AutoTokenizer.from_pretrained(args.scorer)
    if scorer_tok.pad_token_id is None:
        scorer_tok.pad_token = scorer_tok.eos_token
        scorer_tok.pad_token_id = scorer_tok.eos_token_id
    scorer = AutoModelForCausalLM.from_pretrained(args.scorer).to(device).eval()
    if getattr(scorer.config, "pad_token_id", None) is None:
        scorer.config.pad_token_id = scorer_tok.pad_token_id

    stripped = [strip_common_special(t) for t in raw_texts]
    kept_raw, _ = filter_generated_texts(raw_texts, min_chars=1, normalize_whitespace=False, drop_empty=True)
    kept_stripped, _ = filter_generated_texts(stripped, min_chars=1, normalize_whitespace=True, drop_empty=True)
    raw_ppl = score_with_loaded(
        kept_raw,
        scorer,
        scorer_tok,
        batch_size=args.score_batch,
        max_length=args.score_max_length,
        device=device,
    )
    stripped_ppl = score_with_loaded(
        kept_stripped,
        scorer,
        scorer_tok,
        batch_size=args.score_batch,
        max_length=args.score_max_length,
        device=device,
    )
    diversity = summarize_token_diversity(ids).__dict__
    summary = {
        "type": "summary",
        "name": args.name,
        "kind": "mdlm",
        "checkpoint": args.checkpoint,
        "step": ckpt.get("step"),
        "decode": decode,
        "raw_genppl": raw_ppl,
        "stripped_genppl": stripped_ppl,
        "diversity": diversity,
    }

    with (out_dir / f"{args.name}_scored.jsonl").open("w", encoding="utf-8") as f:
        f.write(json.dumps(summary, ensure_ascii=False) + "\n")
        for i, (raw, clean) in enumerate(zip(raw_texts, stripped)):
            f.write(
                json.dumps(
                    {
                        "type": "sample",
                        "name": args.name,
                        "index": i,
                        "raw_text": raw,
                        "stripped_text": clean,
                    },
                    ensure_ascii=False,
                )
                + "\n"
            )
    write_raw_samples(out_dir / f"{args.name}_raw_samples.txt", args.name, summary, raw_texts)
    with (out_dir / "summary.jsonl").open("w", encoding="utf-8") as f:
        f.write(json.dumps(summary, ensure_ascii=False) + "\n")
    with (out_dir / "summary.tsv").open("w", encoding="utf-8") as f:
        f.write("name\tstep\tfinal\traw_genppl\tstripped_genppl\tentropy\tdistinct_2\ttop_token_mass\n")
        f.write(
            "\t".join(
                [
                    args.name,
                    str(ckpt.get("step")),
                    args.final,
                    f"{raw_ppl['ppl']:.6f}",
                    f"{stripped_ppl['ppl']:.6f}",
                    f"{diversity['sample_entropy']:.6f}",
                    f"{diversity['distinct_2']:.6f}",
                    f"{diversity['top_token_mass']:.6f}",
                ]
            )
            + "\n"
        )
    print("[summary]", json.dumps(summary, ensure_ascii=False), flush=True)
    print(f"[done] {out_dir}", flush=True)


if __name__ == "__main__":
    main()