"""Paired example images for the public dataset (non-sexual concepts only).

Predeclared selection, not cherry-picked: the first 12 target and 12 labelled-neighbour test prompts
by prompt_id, and the first 12 indirect probes, each with its first frozen seed, rendered by every
model listed below under the frozen inference config. Output is HF imagefolder format:
<out>/<model>/<prompt_id>.png plus metadata.jsonl with prompt, seed, checkpoint and detector scores.

    python scripts/make_examples.py --variant tiny --concept vangogh --out runs/examples/tiny-vangogh
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import pandas as pd
from huggingface_hub import snapshot_download

from uvw import ROOT
from uvw.generate import GenItem, generate, load_inference_config
from uvw.models import apply_edited_tensors, load_pipeline, load_registry
from uvw.score import load_scoring_config, scorer_for

MODELS = {
    "original": None,
    "C2-UCE": "tiny/vangogh/C2/c2-uce-tiny-vangogh/uce-preserve-x2",
    "C3-ESD": "tiny/vangogh/C3/c3-esd-tiny-vangogh/esd-esd-x",
    "C5-r2": "tiny/vangogh/C5/c5-tiny-vangogh-r2/vg_value_remap_bounded",
    "C6-r1": "tiny/vangogh/C6/c6-tiny-vangogh-r1/vg_value_remap_strong",
    "C6-r2": "tiny/vangogh/C6/c6-tiny-vangogh-r2/vg_preserve_regularized",
    "C6-r3": "tiny/vangogh/C6/c6-tiny-vangogh-r3/vg_value_rewrite_r10",
}
NEIGHBOUR_ARTIST_CATEGORIES = ("subject-", "style-")


def select_prompts(heldout: Path, concept: str, n: int) -> pd.DataFrame:
    test = pd.read_json(heldout / "test" / f"{concept}.jsonl", lines=True).sort_values("prompt_id")
    probes = pd.read_json(heldout / "probes" / f"{concept}_indirect.jsonl", lines=True).sort_values("prompt_id")
    targets = test[test.role == "target"].head(n).assign(set="test-target")
    nb = test[(test.role == "neighbour") & test.category.str.startswith(NEIGHBOUR_ARTIST_CATEGORIES)].head(n)
    return pd.concat([targets, nb.assign(set="test-neighbour"), probes.head(n).assign(set="indirect-probe")])


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--variant", required=True)
    p.add_argument("--concept", required=True)
    p.add_argument("--n", type=int, default=12)
    p.add_argument("--out", required=True)
    a = p.parse_args()
    if a.concept == "nudity":
        raise SystemExit("no example images are produced for nudity")

    heldout = Path(snapshot_download("javadtaghia/unlearning-via-weight-heldout", repo_type="dataset",
                                     local_dir=ROOT / "heldout"))
    ckroot = Path(snapshot_download("javadtaghia/unlearning-via-weight",
                                    allow_patterns=[f"checkpoints/{c}/*" for c in MODELS.values() if c],
                                    local_dir=ROOT / "hub-checkpoints")) / "checkpoints"
    prompts = select_prompts(heldout, a.concept, a.n)
    cfg, scorer = load_inference_config(), scorer_for(a.concept, load_scoring_config())
    out = Path(a.out)
    rows = []
    for model, ckpt in MODELS.items():
        pipe = load_pipeline(load_registry()[a.variant])
        if ckpt:
            apply_edited_tensors(pipe, ckroot / ckpt / "edited_tensors.safetensors")
        items = [GenItem(r.prompt_id, r.text, int(r.seed)) for r in prompts.itertuples()]
        images = [g.image for g in generate(pipe, items, cfg)]
        scores = scorer.score(images)
        (out / model).mkdir(parents=True, exist_ok=True)
        for r, img, s in zip(prompts.itertuples(), images, scores):
            name = f"{model}/{r.prompt_id}.png"
            img.save(out / name)
            rows.append({"file_name": name, "model": model, "checkpoint": ckpt or "original", "set": r.set,
                         "prompt_id": r.prompt_id, "category": r.category, "prompt": r.text, "seed": int(r.seed),
                         "style_top1_clip_l14": s.get("clip_l14_top1"), "style_top1_clip_h14": s.get("clip_h14_top1"),
                         "van_gogh_detected": bool(s.get("target_hit"))})
        print("rendered", model, flush=True)
    with open(out / "metadata.jsonl", "w") as f:
        for row in rows:
            f.write(json.dumps(row) + "\n")


if __name__ == "__main__":
    main()