"""Fusion Perception 1 v0.1 — two-stage evaluation: our global descriptor + AMES rerank.

The headline rows in `README_hf.md` are two-stage. This repository's head produces the
first-stage ranking; a top-1600 shortlist of that ranking is then rescored by AMES
(Suma et al., ECCV 2024) over local descriptors, and the two scores are combined. This
script runs the second stage and scores the result, so the published two-stage cells can
be reproduced from what ships here.

    # one-time: pull the AMES authors' checkpoint and local descriptors (not redistributed)
    python rerank.py --dataset roxford5k --ames-dir ./ames_assets --fetch-only

    # rerank, computing our descriptors from the benchmark images with inference.py
    python rerank.py --dataset roxford5k --ames-dir ./ames_assets \
        --images-root /data/roxford5k/jpg --head standard --out roxford_rerank.json

    # rerank from descriptors you already have (a .pt holding {"q": [Q, 512], "db": [N, 512]})
    python rerank.py --dataset rparis6k --ames-dir ./ames_assets \
        --descriptors rparis_desc.pt --out rparis_rerank.json

Scoring, for one shortlist position, is

    score = lambda * global_cosine + (1 - lambda) * sigmoid(temp * ames_logit)

with `lambda = 0.55`, `temp = 0.3` and shortlist `k = 1600`. That is the AMES paper
default and it is the cell behind every two-stage number in `results.json`. Other lambda
and temperature values can be passed as comma-separated lists; they are diagnostics and
are not reported as results.

The shortlist is junk-aware, matching the AMES and DELG evaluation code: for each query,
ground-truth junk ids (junk, plus easy under the Hard setting) are moved behind the
shortlist before the top-k is cut. Medium and Hard are scored separately because their
junk sets differ, so the model forward runs twice. mAP comes from the official revisitop
`compute_map`, vendored below, which a stage-0 audit scored against the upstream file at
a maximum difference of 0.0.

Third-party assets, none of which are redistributed here
--------------------------------------------------------
AMES code: https://github.com/pavelsuma/ames, Apache-2.0. Fetched with `torch.hub` on
first use, or point `--ames-repo` at your own clone. Only the authors' code is used; the
reranker is not reimplemented.

AMES checkpoint `dinov2_ames.pt`: downloaded by their model class from
http://ptak.felk.cvut.cz/personal/sumapave/public/ames/networks/ .

AMES local descriptors, per dataset, from
http://ptak.felk.cvut.cz/personal/sumapave/public/ames/data/<dataset>/ :
`dinov2_query_local.hdf5`, `dinov2_gallery_local.hdf5`. `--fetch` downloads them. These
are the authors' published DINOv2-B local features, used unchanged, and they are the
reason a reader should apply the DINOv2 pretraining caveat in `README_hf.md` to the
two-stage rows.

Ground truth `gnd_<dataset>.pkl` from the revisited Oxford / Paris release
(https://github.com/filipradenovic/revisitop). Also mirrored on the AMES data host, which
is where `--fetch` takes it from.

Please cite AMES (arXiv 2408.03282) for any use of the two-stage numbers.

Requires: torch, numpy, h5py, plus what `inference.py` needs if descriptors are computed
here rather than supplied. A GPU is expected; the reranker is a transformer over 600
local descriptors per image and 70 x 1600 pairs per difficulty setting per dataset.
Reported numbers were produced on an A10G.
"""
from __future__ import annotations

import argparse
import itertools
import json
import os
import pickle
import sys
import urllib.request
from copy import deepcopy
from typing import List, Optional, Sequence, Tuple

import numpy as np
import torch

AMES_REPO = "pavelsuma/ames"
AMES_DATA_HOST = "http://ptak.felk.cvut.cz/personal/sumapave/public/ames/data"
DATASETS = ("roxford5k", "rparis6k")
LOCAL_FILES = ("dinov2_query_local.hdf5", "dinov2_gallery_local.hdf5")
DESC_NUM = 600
DEFAULT_K = 1600
DEFAULT_LAMBDA = 0.55
DEFAULT_TEMP = 0.3

# Set by load_ames_model so the authors' loaders can be imported afterwards. torch.hub
# takes its checkout off sys.path once the model is built, so the path is kept here.
_AMES_SRC: Optional[str] = None


# ------------------------------------------------------------------ official revisitop
def compute_ap(ranks: np.ndarray, nres: int) -> float:
    """Trapezoidal average precision for one query, from revisitop `python/evaluate.py`."""
    ap = 0.0
    recall_step = 1.0 / nres
    for j in range(len(ranks)):
        rank = ranks[j]
        precision_0 = 1.0 if rank == 0 else float(j) / rank
        precision_1 = float(j + 1) / (rank + 1)
        ap += (precision_0 + precision_1) * recall_step / 2.0
    return ap


def compute_map(ranks: np.ndarray, gnd: Sequence[dict]) -> float:
    """mAP over a [n_gallery, n_query] rank matrix, from revisitop `python/evaluate.py`.

    Each `gnd` entry carries `ok` (positives) and `junk` (ignored). Junk images are
    removed from the ranking by shifting positive ranks down, not by deleting rows.
    """
    map_ = 0.0
    nq = len(gnd)
    nempty = 0
    for i in np.arange(nq):
        qgnd = np.array(gnd[i]["ok"])
        if qgnd.shape[0] == 0:
            nempty += 1
            continue
        try:
            qgndj = np.array(gnd[i]["junk"])
        except Exception:
            qgndj = np.empty(0)
        pos = np.arange(ranks.shape[0])[np.isin(ranks[:, i], qgnd)]
        junk = np.arange(ranks.shape[0])[np.isin(ranks[:, i], qgndj)]
        k = 0
        ij = 0
        if len(junk):
            ip = 0
            while ip < len(pos):
                while ij < len(junk) and pos[ip] > junk[ij]:
                    k += 1
                    ij += 1
                pos[ip] = pos[ip] - k
                ip += 1
        map_ += compute_ap(pos, len(qgnd))
    return map_ / max(nq - nempty, 1)


def maps_from_ranks(ranks: np.ndarray, gnd: Sequence[dict]) -> Tuple[float, float]:
    """(Medium, Hard) mAP as fractions. Medium: ok = easy + hard, junk = junk.
    Hard: ok = hard, junk = junk + easy."""
    gnd_m = [{"ok": np.concatenate([np.array(g["easy"]), np.array(g["hard"])]).astype(int),
              "junk": np.array(g["junk"]).astype(int)} for g in gnd]
    gnd_h = [{"ok": np.array(g["hard"]).astype(int),
              "junk": np.concatenate([np.array(g["junk"]),
                                      np.array(g["easy"])]).astype(int)} for g in gnd]
    return (round(compute_map(ranks, gnd_m), 4), round(compute_map(ranks, gnd_h), 4))


# ------------------------------------------------------------------------ AMES assets
def asset_paths(ames_dir: str, dataset: str) -> dict:
    d = os.path.join(ames_dir, dataset)
    return {"dir": d,
            "query": os.path.join(d, "dinov2_query_local.hdf5"),
            "gallery": os.path.join(d, "dinov2_gallery_local.hdf5"),
            "gnd": os.path.join(d, f"gnd_{dataset}.pkl")}


def fetch_assets(ames_dir: str, dataset: str) -> dict:
    """Download the authors' local descriptors and the revisitop ground truth.

    Written to a `.tmp` name and renamed, so an interrupted download is never mistaken
    for a complete one. The two HDF5 files are several gigabytes each.
    """
    paths = asset_paths(ames_dir, dataset)
    os.makedirs(paths["dir"], exist_ok=True)
    got = {}
    for fn in LOCAL_FILES + (f"gnd_{dataset}.pkl",):
        dst = os.path.join(paths["dir"], fn)
        if os.path.exists(dst):
            got[fn] = f"present ({os.path.getsize(dst)} bytes)"
            continue
        url = f"{AMES_DATA_HOST}/{dataset}/{fn}"
        print(f"fetching {url}", flush=True)
        tmp = dst + ".tmp"
        try:
            urllib.request.urlretrieve(url, tmp)
        except Exception as exc:
            if os.path.exists(tmp):
                os.remove(tmp)
            raise RuntimeError(f"download failed for {url}: {exc}") from exc
        os.replace(tmp, dst)
        got[fn] = f"downloaded ({os.path.getsize(dst)} bytes)"
    return got


def require_assets(ames_dir: str, dataset: str, need_gnd: bool = True) -> dict:
    """Fail with an actionable message rather than a stack trace deep in a loader."""
    paths = asset_paths(ames_dir, dataset)
    wanted = ("query", "gallery", "gnd") if need_gnd else ("query", "gallery")
    missing = [k for k in wanted if not os.path.exists(paths[k])]
    if missing:
        names = ", ".join(os.path.basename(paths[k]) for k in missing)
        raise SystemExit(
            f"AMES assets missing under {paths['dir']}: {names}\n"
            f"They are the authors' files and are not redistributed with this model.\n"
            f"Fetch them with:\n"
            f"    python rerank.py --dataset {dataset} --ames-dir {ames_dir} --fetch-only\n"
            f"The gallery file is several gigabytes, so for a resumable download prefer:\n"
            f"    wget -c -P {os.path.join(ames_dir, dataset)} "
            f"{AMES_DATA_HOST}/{dataset}/{{dinov2_query_local.hdf5,"
            f"dinov2_gallery_local.hdf5,gnd_{dataset}.pkl}}")
    return paths


def load_ames_model(variant: str = "dinov2_ames.pt",
                    ames_repo: Optional[str] = None) -> torch.nn.Module:
    """The authors' AMES model with their pretrained weights.

    With `ames_repo` set, imports from that clone of github.com/pavelsuma/ames.
    Otherwise `torch.hub` fetches the repository. Either way the checkpoint itself is
    downloaded by their own model class from their host on first use and cached under
    `TORCH_HOME`.
    """
    global _AMES_SRC
    binarized = "dist" in variant
    if ames_repo:
        if not os.path.isdir(os.path.join(ames_repo, "src", "models")):
            raise SystemExit(f"--ames-repo {ames_repo} is not a clone of {AMES_REPO} "
                             f"(no src/models). Clone it with:\n"
                             f"    git clone https://github.com/pavelsuma/ames")
        _AMES_SRC = os.path.abspath(ames_repo)
        _add_ames_path()
        from src.models.ames import AMES
        model = AMES(desc_name="dinov2", local_dim=768, pretrained=variant,
                     binarized=binarized)
    else:
        entry = "dinov2_ames_dist" if binarized else "dinov2_ames"
        model = torch.hub.load(AMES_REPO, entry, pretrained=True, trust_repo=True)
        hub = torch.hub.get_dir()
        cached = sorted(d for d in os.listdir(hub)
                        if d.startswith(AMES_REPO.replace("/", "_")))
        if not cached:
            raise SystemExit(f"torch.hub loaded {AMES_REPO} but no checkout is under "
                             f"{hub}; pass --ames-repo instead")
        _AMES_SRC = os.path.join(hub, cached[-1])
        _add_ames_path()
    return model.eval()


def _add_ames_path() -> None:
    for p in (os.path.join(_AMES_SRC, "src"), _AMES_SRC):
        if p not in sys.path:
            sys.path.insert(0, p)


def ames_query_dataset(dataset: str, desc_dir: str, gnd: Sequence[dict]):
    """The authors' TestDataset over their query local descriptors."""
    if _AMES_SRC is None:
        raise SystemExit("load_ames_model must run before the AMES loaders are used")
    _add_ames_path()
    from src.utils.tensor_dataset import TestDataset
    return TestDataset(dataset, desc_dir, "dinov2_query_local.hdf5",
                       desc_num=DESC_NUM, gnd_data=gnd)


class GalleryLocals(torch.utils.data.Dataset):
    """Gallery-side local descriptors: the benchmark database, then optional distractors.

    The benchmark part is the authors' `dinov2_gallery_local.hdf5`. The distractor part
    is a list of HDF5 shards in the same layout, appended in order, so gallery index
    `i >= n_database` addresses the (i - n_database)-th distractor. Shards are opened
    lazily, which is what makes the +1M setting fit in memory.
    """

    def __init__(self, db_path: str, shard_paths: Sequence[str],
                 shard_sizes: Sequence[int], desc_num: int = DESC_NUM):
        import h5py
        self.db = h5py.File(db_path, "r")
        self.ndb = len(self.db["features"])
        self.desc_num = desc_num
        self.paths = list(shard_paths)
        self.bounds = np.cumsum([0] + list(shard_sizes))
        self.handles: dict = {}

    def row(self, gi: int):
        import h5py
        if gi < self.ndb:
            a = self.db["features"][gi]
        else:
            r = gi - self.ndb
            si = int(np.searchsorted(self.bounds, r, side="right") - 1)
            if si not in self.handles:
                self.handles[si] = h5py.File(self.paths[si], "r")
            a = self.handles[si]["features"][r - self.bounds[si]]
        if a.dtype.names:
            return (a["descriptor"][:self.desc_num].astype(np.float32),
                    a["metadata"][:self.desc_num, 3])
        return (a[:self.desc_num, 5:].astype(np.float32), a[:self.desc_num, 3])

    def __len__(self) -> int:
        return int(self.bounds[-1]) + self.ndb

    def __getitem__(self, batch_index):
        feats, masks = [], []
        for gi in batch_index:
            f, m = self.row(int(gi))
            feats.append(torch.from_numpy(f))
            masks.append(torch.from_numpy(m.astype(bool)))
        return (torch.stack(feats).float(), torch.stack(masks)), batch_index


# ------------------------------------------------------------------ global descriptors
def load_descriptors(path: str) -> Tuple[torch.Tensor, torch.Tensor]:
    """Read a descriptor file: `{"q": [Q, D], "db": [N, D]}`, L2-normalized, any dtype."""
    d = torch.load(path, map_location="cpu", weights_only=False)
    for k in ("q", "db"):
        if k not in d:
            raise SystemExit(f"{path} has no '{k}' tensor; expected keys 'q' and 'db'")
    return d["q"].float(), d["db"].float()


def embed_benchmark(images_root: str, gnd_pkl: dict, head: str,
                    model_root: Optional[str] = None,
                    batch: int = 64) -> Tuple[torch.Tensor, torch.Tensor]:
    """Compute our first-stage descriptors for one benchmark with `inference.py`.

    Queries are cropped to the ground-truth bounding box; gallery images are not. This
    is the revisitop query protocol and it is what every number in `results.json` uses.
    """
    from PIL import Image
    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
    from inference import FusionPerceptionRetrieval

    root = model_root or os.path.dirname(os.path.abspath(__file__))
    fp = FusionPerceptionRetrieval.from_pretrained(root, protocol=head)

    def run(names: List[str], boxes: Optional[List]) -> torch.Tensor:
        out = []
        for s in range(0, len(names), batch):
            chunk = names[s:s + batch]
            pils = [Image.open(os.path.join(images_root, f"{n}.jpg")) for n in chunk]
            bb = None if boxes is None else boxes[s:s + batch]
            out.append(fp.embed(pils, bbox=bb))
            for im in pils:
                im.close()
            print(f"  embedded {min(s + batch, len(names))}/{len(names)}", flush=True)
        return torch.cat(out)

    db = run(list(gnd_pkl["imlist"]), None)
    qbox = [list(g["bbx"]) for g in gnd_pkl["gnd"]]
    q = run(list(gnd_pkl["qimlist"]), qbox)
    return q.float(), db.float()


def load_distractors(locals_dir: str, desc_path: str,
                     protocol: str) -> Tuple[torch.Tensor, List[str], List[int]]:
    """Distractor global descriptors plus the ordered list of local-descriptor shards.

    `locals_dir` holds `r1m_order.json` (the canonical shard order and per-shard counts)
    and the `locals_XXXX.hdf5` shards it names. `desc_path` holds `{"desc": [M, D]}` in
    exactly that order. Both are produced by our distractor pipeline; see REPRODUCE.md.
    """
    order_path = os.path.join(locals_dir, "r1m_order.json")
    if not os.path.exists(order_path):
        raise SystemExit(f"no r1m_order.json under {locals_dir}; see REPRODUCE.md for "
                         f"how the distractor shards are produced")
    order = json.load(open(order_path))
    paths = [os.path.join(locals_dir, f"locals_{s['arch']:04d}.hdf5")
             for s in order["shards"]]
    sizes = [int(s["n"]) for s in order["shards"]]
    absent = [p for p in paths if not os.path.exists(p)]
    if absent:
        raise SystemExit(f"{len(absent)} distractor local shards missing under "
                         f"{locals_dir}, first is {os.path.basename(absent[0])}")
    d = torch.load(desc_path, map_location="cpu", weights_only=False)
    desc = (d["desc"] if isinstance(d, dict) else d).float()
    if len(desc) != sum(sizes):
        raise SystemExit(f"{desc_path} has {len(desc)} descriptors but the shards hold "
                         f"{sum(sizes)}; the two must be in the same order")
    print(f"distractors: {len(desc)} images over {len(paths)} shards "
          f"({protocol} head)", flush=True)
    return desc, paths, sizes


# ----------------------------------------------------------------------------- rerank
def _first(batch):
    """Collate for the authors' loaders, which already batch inside the dataset."""
    return batch[0]


def junk_aware_shortlist(cache_nn: torch.Tensor, gnd: Sequence[dict],
                         hard: bool) -> Tuple[torch.Tensor, torch.Tensor]:
    """Move ground-truth junk behind the ranking, per query, before the top-k is cut.

    This is the AMES and DELG evaluation convention: junk images are not reranked and
    do not consume shortlist slots. Under Hard the easy positives count as junk, so the
    two difficulty settings get different shortlists and are reranked separately.
    """
    nn = cache_nn.clone()
    for i in range(cache_nn.shape[1]):
        junk_ids = list(gnd[i]["junk"]) + (list(gnd[i]["easy"]) if hard else [])
        is_junk = torch.from_numpy(np.isin(cache_nn[1, i].numpy(), junk_ids))
        nn[:, i] = torch.cat((cache_nn[:, i, ~is_junk], cache_nn[:, i, is_junk]), dim=1)
    return nn[0], nn[1].long()


def ames_scores(model, q_loader, g_loader, nn_inds: torch.Tensor,
                max_k: int, device: str) -> torch.Tensor:
    """Raw AMES logits for every (query, shortlist candidate) pair: [n_query, max_k].

    One gallery-loader step handles one shortlist position across all queries, which is
    how the authors' `rerank()` drives it.
    """
    scores = []
    for q_f, i in q_loader:
        q_score = []
        g_loader.batch_sampler.sampler = nn_inds[i, :max_k].T.tolist()
        for db_f, _ in g_loader:
            cur = model(*[x.to(device, non_blocking=True) for x in q_f],
                        *[x.to(device, non_blocking=True) for x in db_f])
            q_score.append(cur.cpu().data)
        scores.append(torch.stack(q_score).T)
    return torch.cat(scores)


def fuse_and_score(raw_sim: torch.Tensor, nn_sims: torch.Tensor, nn_inds: torch.Tensor,
                   gnd: Sequence[dict], hard: bool, topk: Sequence[int],
                   lambs: Sequence[float], temps: Sequence[float]) -> dict:
    """Fuse global and reranked scores over the (k, lambda, temp) grid and score mAP.

    `score = lambda * global_cosine + (1 - lambda) * sigmoid(temp * ames_logit)`. Only
    the first k entries are reordered; the tail of the ranking is left as the global
    descriptor left it, which is what makes the shortlist size a real parameter.
    """
    cells = {}
    for k, lam, temp in itertools.product(topk, lambs, temps):
        s = 1.0 / (1.0 + torch.exp(-temp * raw_sim))
        s = lam * nn_sims[:, :k] + (1.0 - lam) * s[:, :k]
        _, indices = torch.sort(s, dim=-1, descending=True)
        closest = torch.gather(nn_inds[:, :k], -1, indices)
        ranks = deepcopy(nn_inds)
        ranks[:, :k] = closest
        M, H = maps_from_ranks(ranks.numpy().T, gnd)
        cells[f"k{k}_l{lam}_t{temp}"] = H if hard else M
    return cells


def rerank_dataset(dataset: str, ames_dir: str, q: torch.Tensor, db: torch.Tensor,
                   gnd: Sequence[dict], topk: Sequence[int], lambs: Sequence[float],
                   temps: Sequence[float], model, device: str,
                   distractors: Optional[Tuple[torch.Tensor, List[str], List[int]]] = None,
                   workers: int = 4) -> dict:
    """Global-only and two-stage mAP for one dataset."""
    from torch.utils.data import BatchSampler, DataLoader, SequentialSampler

    paths = asset_paths(ames_dir, dataset)
    gallery_desc = db if distractors is None else torch.cat([db, distractors[0]])
    shard_paths = [] if distractors is None else distractors[1]
    shard_sizes = [] if distractors is None else distractors[2]

    sims = q @ gallery_desc.T
    sims_sorted, inds_sorted = torch.sort(sims, dim=1, descending=True)
    gM, gH = maps_from_ranks(inds_sorted.T.numpy(), gnd)
    out = {"n_query": int(q.shape[0]), "n_gallery": int(gallery_desc.shape[0]),
           "global_only": {"M": gM, "H": gH}}
    print(f"[{dataset}] global only: M {gM * 100:.2f}  H {gH * 100:.2f}", flush=True)

    cache_nn = torch.stack((sims_sorted, inds_sorted.float()))
    qset = ames_query_dataset(dataset, paths["dir"], gnd)
    gset = GalleryLocals(paths["gallery"], shard_paths, shard_sizes)
    if len(gset) != gallery_desc.shape[0]:
        raise SystemExit(f"{dataset}: {gallery_desc.shape[0]} global descriptors but "
                         f"{len(gset)} rows of local descriptors; they must line up")

    def loader(ds, nw):
        return DataLoader(ds, sampler=BatchSampler(SequentialSampler(ds), batch_size=300,
                                                   drop_last=False),
                          batch_size=1, num_workers=nw, collate_fn=_first)

    q_loader, g_loader = loader(qset, 2), loader(gset, workers)
    max_k = max(topk)
    cells: dict = {}
    with torch.no_grad():
        for hard in (False, True):
            nn_sims, nn_inds = junk_aware_shortlist(cache_nn, gnd, hard)
            raw_sim = ames_scores(model, q_loader, g_loader, nn_inds, max_k, device)
            got = fuse_and_score(raw_sim, nn_sims, nn_inds, gnd, hard, topk, lambs, temps)
            for key, val in got.items():
                cells.setdefault(key, {})["H" if hard else "M"] = val
            print(f"[{dataset}] reranked {'Hard' if hard else 'Medium'}: " +
                  "  ".join(f"{k} {v * 100:.2f}" for k, v in got.items()), flush=True)
    out["reranked"] = cells
    return out


# -------------------------------------------------------------------------------- CLI
def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="AMES reranking of the Fusion Perception first-stage ranking.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="AMES (Suma et al., ECCV 2024) is third-party work under Apache-2.0. Its "
               "checkpoint and local descriptors are downloaded from the authors' host "
               "and are not redistributed with this model.")
    p.add_argument("--dataset", default="roxford5k", choices=DATASETS)
    p.add_argument("--ames-dir", default="./ames_assets",
                   help="directory holding <dataset>/dinov2_*_local.hdf5 and the gnd pickle")
    p.add_argument("--ames-repo", default=None,
                   help="path to a clone of github.com/pavelsuma/ames; "
                        "without it, torch.hub fetches the repository")
    p.add_argument("--ames-variant", default="dinov2_ames.pt",
                   help="AMES checkpoint name on the authors' host")
    p.add_argument("--fetch", action="store_true",
                   help="download any missing AMES assets before running")
    p.add_argument("--fetch-only", action="store_true",
                   help="download the AMES assets and exit")

    p.add_argument("--descriptors", default=None,
                   help="a .pt holding our first-stage descriptors as "
                        "{'q': [Q, 512], 'db': [N, 512]}, L2-normalized")
    p.add_argument("--images-root", default=None,
                   help="benchmark jpg directory; descriptors are computed with "
                        "inference.py when --descriptors is not given")
    p.add_argument("--head", default="standard", choices=("standard", "decon"))
    p.add_argument("--model-root", default=None,
                   help="this model's directory; defaults to the directory of this file")
    p.add_argument("--gnd", default=None,
                   help="path to gnd_<dataset>.pkl; defaults to the copy in --ames-dir")
    p.add_argument("--save-descriptors", default=None,
                   help="write the computed descriptors here so a rerun can skip them")

    p.add_argument("--distractor-locals", default=None,
                   help="directory with r1m_order.json and locals_XXXX.hdf5; "
                        "enables the +1M setting")
    p.add_argument("--distractor-desc", default=None,
                   help="a .pt holding {'desc': [M, 512]} for the same distractors, "
                        "in r1m_order.json order")

    p.add_argument("--topk", default=str(DEFAULT_K),
                   help="shortlist sizes, comma-separated")
    p.add_argument("--lambdas", default=str(DEFAULT_LAMBDA),
                   help="global-score weights, comma-separated")
    p.add_argument("--temps", default=str(DEFAULT_TEMP),
                   help="sigmoid temperatures, comma-separated")
    p.add_argument("--device", default=None)
    p.add_argument("--workers", type=int, default=4,
                   help="gallery loader workers; use 0 on platforms without fork, "
                        "since the HDF5 handles are not picklable")
    p.add_argument("--out", default=None, help="write the result JSON here")
    return p.parse_args(argv)


def main(argv: Optional[Sequence[str]] = None) -> dict:
    a = parse_args(argv)
    if a.fetch or a.fetch_only:
        print(json.dumps(fetch_assets(a.ames_dir, a.dataset), indent=1))
        if a.fetch_only:
            return {}
    paths = require_assets(a.ames_dir, a.dataset, need_gnd=not a.gnd)

    gnd_path = a.gnd or paths["gnd"]
    with open(gnd_path, "rb") as fh:
        gnd_pkl = pickle.load(fh)
    gnd = gnd_pkl["gnd"]

    if a.descriptors:
        q, db = load_descriptors(a.descriptors)
    elif a.images_root:
        q, db = embed_benchmark(a.images_root, gnd_pkl, a.head, a.model_root)
        if a.save_descriptors:
            torch.save({"q": q, "db": db}, a.save_descriptors)
            print(f"wrote {a.save_descriptors}", flush=True)
    else:
        raise SystemExit("give either --descriptors or --images-root")
    if len(gnd) != q.shape[0]:
        raise SystemExit(f"{len(gnd)} ground-truth queries but {q.shape[0]} query "
                         f"descriptors")

    distractors = None
    if a.distractor_locals or a.distractor_desc:
        if not (a.distractor_locals and a.distractor_desc):
            raise SystemExit("the +1M setting needs both --distractor-locals and "
                             "--distractor-desc")
        distractors = load_distractors(a.distractor_locals, a.distractor_desc, a.head)

    device = a.device or ("cuda" if torch.cuda.is_available() else "cpu")
    model = load_ames_model(a.ames_variant, a.ames_repo).to(device)
    res = rerank_dataset(a.dataset, a.ames_dir, q, db, gnd,
                         [int(x) for x in a.topk.split(",")],
                         [float(x) for x in a.lambdas.split(",")],
                         [float(x) for x in a.temps.split(",")],
                         model, device, distractors, a.workers)
    payload = {"model": "fusion-perception-1 v0.1-preview", "head": a.head,
               "dataset": a.dataset, "reranker": a.ames_variant,
               "setting": "plus_1m" if distractors else "no_distractors",
               "fusion": "lambda * global_cosine + (1 - lambda) * sigmoid(temp * ames_logit)",
               "metric": "mAP as a fraction, revisitop Medium and Hard", **res}
    if a.out:
        with open(a.out, "w") as fh:
            json.dump(payload, fh, indent=1)
        print(f"wrote {a.out}", flush=True)
    print("RERANK:", json.dumps(payload, indent=1), flush=True)
    return payload


if __name__ == "__main__":
    main()