#!/usr/bin/env python
"""danbooru-tagger-v1 + wd-eva02-large-tagger-v3, blended.

This is the configuration that actually beats WD on ranking quality. On 11,639
held-out Danbooru posts with id > 8,600,750 (after every incumbent's training
cutoff), against wd-eva02-large-tagger-v3, paired bootstrap over images:

    macro-AP  0.5090 vs 0.4668   +0.0421  95% CI [+0.039, +0.044]
    fine AP   0.6113 vs 0.5979   +0.0133  95% CI [+0.002, +0.022]

Either model alone loses one of those. The two disagree in useful ways
(probability correlation r = 0.876): v1 is stronger on tags WD handles badly,
WD is stronger on tags it handles well, and averaging keeps both.

Cost: two forward passes instead of one. If you only want a single model, use
predict.py -- it wins macro-F1 and coverage but ties macro-AP and loses fine AP.

    pip install torch timm pillow numpy huggingface_hub
    python ensemble.py image.png --thr 0.38
"""
import argparse, csv, os
import numpy as np
import torch
from PIL import Image

SIDE = 448
WD_REPO = "SmilingWolf/wd-eva02-large-tagger-v3"


def to_rgb(im):
    if im.mode in ("RGBA", "LA", "P"):
        im = im.convert("RGBA")
        bg = Image.new("RGBA", im.size, (255, 255, 255))
        bg.alpha_composite(im)
        return bg.convert("RGB")
    return im.convert("RGB")


def pad_square(im):
    w, h = im.size
    if w == h:
        return im
    s = max(w, h)
    bg = Image.new("RGB", (s, s), (255, 255, 255))
    bg.paste(im, ((s - w) // 2, (s - h) // 2))
    return bg


def preprocess(path):
    """Both models share this exactly: white pad, bicubic 448, [-1,1], RGB->BGR."""
    im = pad_square(to_rgb(Image.open(path))).resize((SIDE, SIDE), Image.BICUBIC)
    x = torch.from_numpy(np.asarray(im, dtype=np.float32) / 255.0).permute(2, 0, 1)
    x = (x - 0.5) / 0.5
    return x[[2, 1, 0]][None]


def snapshot_download_wd():
    from huggingface_hub import snapshot_download
    return snapshot_download(WD_REPO, allow_patterns=["config.json", "model.safetensors",
                                                      "selected_tags.csv"])


def load_wd(device, path):
    import json, timm
    from safetensors.torch import load_file
    cfg = json.load(open(f"{path}/config.json"))
    m = timm.create_model(cfg["architecture"], pretrained=False,
                          num_classes=cfg["num_classes"], **cfg.get("model_args", {}))
    missing, unexpected = m.load_state_dict(load_file(f"{path}/model.safetensors"), strict=False)
    if missing or unexpected:
        raise SystemExit(f"WD weights did not load cleanly: {missing[:3]} {unexpected[:3]}")
    rows = list(csv.DictReader(open(f"{path}/selected_tags.csv")))
    # general tags only (category 0). The published comparison was computed over
    # each model's general vocabulary, so blending WD's character/copyright rows
    # here would not match the numbers in the README -- 'chocobo' is general in
    # our vocabulary and category 4 in WD's, and is the one tag this excludes.
    return m.eval().to(device), [r["name"] if int(r["category"]) == 0 else None
                                 for r in rows]


def load_ours(device, ckpt, base_arch_from):
    import timm, json
    ck = torch.load(ckpt, map_location="cpu", weights_only=False)
    cfg = json.load(open(f"{base_arch_from}/config.json"))
    m = timm.create_model(cfg["architecture"], pretrained=False,
                          num_classes=len(ck["tags"]), **cfg.get("model_args", {}))
    m.load_state_dict({k: v.float() for k, v in ck["model"].items()})
    return m.eval().to(device), list(ck["tags"])


def main():
    here = os.path.dirname(os.path.abspath(__file__))
    ap = argparse.ArgumentParser()
    ap.add_argument("images", nargs="+")
    ap.add_argument("--ckpt", default=os.path.join(here, "weights.fp16.pt"))
    ap.add_argument("--thr", type=float, default=0.38)
    ap.add_argument("--w", type=float, default=0.5, help="weight on our model; 0.5 is what was measured")
    ap.add_argument("--wd-path", default="", help="local dir holding WD's config.json/model.safetensors/"
                    "selected_tags.csv, instead of downloading from the Hub")
    a = ap.parse_args()

    dev = "cuda" if torch.cuda.is_available() else "cpu"
    wd_path = a.wd_path or snapshot_download_wd()
    wd, wd_tags = load_wd(dev, wd_path)
    ours, our_tags = load_ours(dev, a.ckpt, wd_path)

    our_idx = {t: i for i, t in enumerate(our_tags)}
    # blend only where both models have the tag; elsewhere ours stands alone
    shared = [(our_idx[t], j) for j, t in enumerate(wd_tags) if t is not None and t in our_idx]
    oi = np.array([s[0] for s in shared]); wi = np.array([s[1] for s in shared])
    print(f"{len(our_tags):,} tags, {len(shared):,} shared with WD -> blended; "
          f"{len(our_tags)-len(shared):,} ours only")

    for path in a.images:
        x = preprocess(path).to(dev)
        with torch.inference_mode():
            p = torch.sigmoid(ours(x).float())[0].cpu().numpy()
            q = torch.sigmoid(wd(x).float())[0].cpu().numpy()
        p[oi] = a.w * p[oi] + (1.0 - a.w) * q[wi]
        order = np.argsort(-p)
        hit = [(our_tags[i], float(p[i])) for i in order if p[i] >= a.thr]
        print(f"\n=== {path}  ({len(hit)} tags >= {a.thr})")
        print(", ".join(f"{t}:{v:.2f}" for t, v in hit) or "(nothing above threshold)")


if __name__ == "__main__":
    main()