#!/usr/bin/env python3
"""Rolling Dirichlet-20k decode search with final distribution sampling.

The previous sweep used argmax at the final state. This script keeps the same
Dirichlet bridge update, but tries a small set of final sampling/top-p settings
and stops as soon as entropy and gen-PPL hit the requested target.
"""

from __future__ import annotations

import argparse
import csv
import importlib.util
import json
import math
import sys
from dataclasses import dataclass
from pathlib import Path

import torch
import torch.nn.functional as F


BASE = Path(__file__).with_name("eval_c1024_decode_sweep_20260507.py")
spec = importlib.util.spec_from_file_location("eval_c1024_decode_sweep_20260507", BASE)
if spec is None or spec.loader is None:
    raise RuntimeError(f"cannot import {BASE}")
base = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = base
spec.loader.exec_module(base)


@dataclass(frozen=True)
class DecodeConfig:
    name: str
    model_t_mode: str = "post"
    support_power: float = 1.0
    semantic_power: float = 1.0
    final_from: str = "blend"
    endpoint_temp: float = 1.3
    concentration_min: float = 1.0
    concentration_max: float = 16.0
    update: str = "resample"
    final_decode: str = "sample"  # "argmax" or "sample"
    final_temp: float = 1.0
    top_p: float = 1.0
    top_k: int = 0
    sample_frac: float = 1.0
    select_by: str = "all"  # "all" or "low_conf"


def fmt(x: float) -> str:
    return ("%g" % x).replace(".", "p")


def build_configs() -> list[DecodeConfig]:
    configs: list[DecodeConfig] = []

    # Start from the previous best area, then release diversity in a controlled
    # way: only sample low-confidence positions and/or only from top-k.
    anchors = [
        (16, 1.30, 1.0, 1.0),
        (12, 1.30, 1.0, 1.0),
        (24, 1.30, 1.0, 1.0),
        (16, 1.35, 1.0, 1.0),
        (16, 1.25, 1.0, 1.0),
        (8, 1.30, 1.0, 1.0),
        (32, 1.30, 1.0, 1.0),
        (16, 1.30, 0.8, 1.0),
        (16, 1.30, 1.2, 1.0),
        (16, 1.30, 1.0, 0.8),
        (16, 1.30, 1.0, 1.3),
    ]
    for cmax, endpoint_temp, support_power, semantic_power in anchors:
        for top_k in [2, 3, 5, 10]:
            for sample_frac in [0.10, 0.20, 0.35, 0.50, 1.00]:
                for final_temp in [1.0, 1.3, 1.6, 2.0]:
                    # Small top-k with all-position sampling can be OK; larger
                    # top-k is safer when restricted to low-confidence tokens.
                    if sample_frac == 1.0 and top_k > 3:
                        continue
                    configs.append(
                        DecodeConfig(
                            name=(
                                f"sel_c{cmax}_et{fmt(endpoint_temp)}_sp{fmt(support_power)}"
                                f"_sem{fmt(semantic_power)}_k{top_k}_frac{fmt(sample_frac)}"
                                f"_ft{fmt(final_temp)}"
                            ),
                            endpoint_temp=endpoint_temp,
                            support_power=support_power,
                            semantic_power=semantic_power,
                            concentration_max=float(cmax),
                            final_temp=final_temp,
                            top_k=top_k,
                            sample_frac=sample_frac,
                            select_by="all" if sample_frac >= 0.999 else "low_conf",
                        )
                    )
    return configs


def top_p_filter(probs: torch.Tensor, top_p: float) -> torch.Tensor:
    if top_p >= 0.999:
        return probs
    sorted_probs, sorted_idx = torch.sort(probs, dim=-1, descending=True)
    keep = torch.cumsum(sorted_probs, dim=-1) <= top_p
    keep[..., 0] = True
    filtered = torch.zeros_like(probs)
    filtered.scatter_(-1, sorted_idx, sorted_probs * keep.to(sorted_probs.dtype))
    return filtered / filtered.sum(dim=-1, keepdim=True).clamp_min(1e-8)


def top_k_filter(probs: torch.Tensor, top_k: int) -> torch.Tensor:
    if top_k <= 0 or top_k >= probs.size(-1):
        return probs
    values, indices = torch.topk(probs, k=top_k, dim=-1)
    filtered = torch.zeros_like(probs)
    filtered.scatter_(-1, indices, values)
    return filtered / filtered.sum(dim=-1, keepdim=True).clamp_min(1e-8)


@torch.no_grad()
def decode_config(
    model,
    tokenizer,
    cfg: DecodeConfig,
    *,
    n_samples: int,
    batch_size: int,
    max_len: int,
    steps: int,
    seed: int,
    device: torch.device,
) -> tuple[list[list[int]], list[str], dict]:
    torch.manual_seed(seed)
    eps = 1e-8
    all_ids: list[list[int]] = []
    all_texts: list[str] = []
    remaining = n_samples
    while remaining > 0:
        bs = min(batch_size, remaining)
        probs = base.sample_noise_simplex(
            (bs, max_len),
            tokenizer.vocab_size,
            device,
            eps,
            noise_mode="dirichlet",
            target_prob=1.0,
            noise_sigma=-1.0,
            dirichlet_concentration=1.0,
        )
        attn = torch.ones((bs, max_len), dtype=torch.bool, device=device)
        last_endpoint = probs
        for step in range(steps):
            support_t = ((step + 1) / max(steps, 1)) ** cfg.support_power
            semantic_t = ((step + 1) / max(steps, 1)) ** cfg.semantic_power
            t = base.model_time(cfg.model_t_mode, step, steps, bs, device)
            logits = model(base.state_for_model(model, probs, eps), t, attn).float() / cfg.endpoint_temp
            endpoint = F.softmax(logits, dim=-1)
            last_endpoint = endpoint
            anchor = probs.clamp_min(eps)
            anchor = anchor / anchor.sum(dim=-1, keepdim=True).clamp_min(eps)
            forward_endpoint = (1.0 - semantic_t) * anchor + semantic_t * endpoint
            forward_endpoint = forward_endpoint.clamp_min(eps)
            forward_endpoint = forward_endpoint / forward_endpoint.sum(dim=-1, keepdim=True).clamp_min(eps)
            mean = (1.0 - support_t) / float(tokenizer.vocab_size) + support_t * forward_endpoint
            mean = mean.clamp_min(eps)
            mean = mean / mean.sum(dim=-1, keepdim=True).clamp_min(eps)
            log_min = math.log(max(cfg.concentration_min, eps))
            log_max = math.log(max(cfg.concentration_max, cfg.concentration_min))
            conc = math.exp(log_min + support_t * (log_max - log_min))
            alpha = (mean * conc).clamp_min(eps)
            probs = torch._standard_gamma(alpha).clamp_min(eps)
            probs = probs / probs.sum(dim=-1, keepdim=True).clamp_min(eps)

        if cfg.final_from == "endpoint":
            final = last_endpoint
        elif cfg.final_from == "state":
            final = probs
        elif cfg.final_from == "blend":
            final = 0.5 * probs + 0.5 * last_endpoint
        else:
            raise ValueError(cfg.final_from)
        final = final / final.sum(dim=-1, keepdim=True).clamp_min(eps)
        if cfg.final_decode == "sample":
            sample_probs = (final.clamp_min(eps).log() / cfg.final_temp).softmax(dim=-1)
            sample_probs = top_p_filter(sample_probs, cfg.top_p)
            sample_probs = top_k_filter(sample_probs, cfg.top_k)
            sampled = torch.multinomial(sample_probs.reshape(-1, sample_probs.size(-1)), 1).reshape(bs, max_len)
            argmax = final.argmax(dim=-1)
            if cfg.select_by == "all" or cfg.sample_frac >= 0.999:
                ids_t = sampled
            elif cfg.select_by == "low_conf":
                k_pos = max(1, int(round(max_len * cfg.sample_frac)))
                conf = final.max(dim=-1).values
                low_idx = torch.topk(-conf, k=k_pos, dim=-1).indices
                mask = torch.zeros((bs, max_len), dtype=torch.bool, device=device)
                mask.scatter_(1, low_idx, True)
                ids_t = torch.where(mask, sampled, argmax)
            else:
                raise ValueError(cfg.select_by)
        elif cfg.final_decode == "argmax":
            ids_t = final.argmax(dim=-1)
        else:
            raise ValueError(cfg.final_decode)
        ids = ids_t.detach().cpu().tolist()
        all_ids.extend(ids)
        all_texts.extend(tokenizer.decode(row, stop_at_eos=False, skip_special_tokens=False) for row in ids)
        remaining -= bs
    decode = {
        "kind": "dirichlet20k_final_sampling",
        "config": cfg.name,
        "steps": steps,
        "model_t_mode": cfg.model_t_mode,
        "support_power": cfg.support_power,
        "semantic_power": cfg.semantic_power,
        "endpoint_temp": cfg.endpoint_temp,
        "concentration_min": cfg.concentration_min,
        "concentration_max": cfg.concentration_max,
        "final_from": cfg.final_from,
        "update": cfg.update,
        "final_decode": cfg.final_decode,
        "final_temp": cfg.final_temp,
        "top_p": cfg.top_p,
        "top_k": cfg.top_k,
        "sample_frac": cfg.sample_frac,
        "select_by": cfg.select_by,
        "n_samples": n_samples,
        "seed": seed,
    }
    return all_ids, all_texts, decode


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("--n_samples", type=int, default=64)
    p.add_argument("--max_len", type=int, default=128)
    p.add_argument("--steps", type=int, default=512)
    p.add_argument("--decode_batch", type=int, default=16)
    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=20260508)
    p.add_argument("--target_entropy", type=float, default=4.2)
    p.add_argument("--target_ppl", type=float, default=35.0)
    p.add_argument("--keep_running_after_hit", 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 = base.BpeTextTokenizer.from_file(args.tokenizer_path)
    ckpt = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
    model = base.build_model(ckpt, tokenizer, device)
    scorer_tok = base.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 = base.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

    rows: list[dict] = []
    summary_jsonl = out_dir / "summary.jsonl"
    summary_jsonl.write_text("", encoding="utf-8")
    for cfg in build_configs():
        print(f"[decode] {cfg.name}", flush=True)
        ids, raw_texts, decode = decode_config(
            model,
            tokenizer,
            cfg,
            n_samples=args.n_samples,
            batch_size=args.decode_batch,
            max_len=args.max_len,
            steps=args.steps,
            seed=args.seed,
            device=device,
        )
        detok_texts = [base.lm1b_detokenizer(text) for text in raw_texts]
        detok_ppl = base.score_texts(
            [text for text in detok_texts if text.strip()],
            scorer,
            scorer_tok,
            batch_size=args.score_batch,
            max_length=args.score_max_length,
            device=device,
        )
        diversity = base.summarize_token_diversity(ids).__dict__
        summary = {
            "type": "summary",
            "name": cfg.name,
            "checkpoint": args.checkpoint,
            "step": ckpt.get("step"),
            "decode": decode,
            "detok_genppl": detok_ppl,
            "diversity": diversity,
        }
        base.write_samples(out_dir / f"{cfg.name}_raw_samples.txt", cfg.name, summary, raw_texts)
        base.write_samples(out_dir / f"{cfg.name}_detok_raw_samples.txt", cfg.name + "_detok", summary, detok_texts)
        with (out_dir / f"{cfg.name}_scored.jsonl").open("w", encoding="utf-8") as f:
            f.write(json.dumps(summary, ensure_ascii=False) + "\n")
            for i, (raw, detok) in enumerate(zip(raw_texts, detok_texts)):
                f.write(json.dumps({"type": "sample", "index": i, "raw_text": raw, "detok_text": detok}, ensure_ascii=False) + "\n")
        with summary_jsonl.open("a", encoding="utf-8") as f:
            f.write(json.dumps(summary, ensure_ascii=False) + "\n")
        row = {
            "name": cfg.name,
            "step": ckpt.get("step"),
            "detok_genppl": detok_ppl["ppl"],
            "sample_entropy": diversity["sample_entropy"],
            "distinct_2": diversity["distinct_2"],
            "top_token_mass": diversity["top_token_mass"],
            "model_t_mode": cfg.model_t_mode,
            "support_power": cfg.support_power,
            "semantic_power": cfg.semantic_power,
            "final_from": cfg.final_from,
            "endpoint_temp": cfg.endpoint_temp,
            "concentration_max": cfg.concentration_max,
            "final_decode": cfg.final_decode,
            "final_temp": cfg.final_temp,
            "top_p": cfg.top_p,
            "top_k": cfg.top_k,
            "sample_frac": cfg.sample_frac,
            "select_by": cfg.select_by,
        }
        rows.append(row)
        keys = list(rows[0].keys())
        with (out_dir / "summary.tsv").open("w", encoding="utf-8") as f:
            writer = csv.DictWriter(f, fieldnames=keys, delimiter="\t")
            writer.writeheader()
            writer.writerows(rows)
        print("[summary]", json.dumps(row, ensure_ascii=False), flush=True)
        if row["sample_entropy"] > args.target_entropy and row["detok_genppl"] < args.target_ppl:
            print("[hit]", json.dumps(row, ensure_ascii=False), flush=True)
            (out_dir / "HIT.json").write_text(json.dumps(row, ensure_ascii=False, indent=2), encoding="utf-8")
            if not args.keep_running_after_hit:
                break


if __name__ == "__main__":
    main()