"""Recompute PR curves from an image_similarity .store file (messagepack).

Replicates DescriptorStore::get_stats semantics:
- micro-averaged over base-image queries
- TP_m(t): mutant m of query found within Hamming distance t
- FP: other base images count for every mutator; other images' mutants
  count only for their own mutator
"""
import sys
import msgpack
import numpy as np

TMAX = 33  # thresholds 0..32

def load(path):
    with open(path, "rb") as f:
        m = msgpack.unpack(f, strict_map_key=False)
    if isinstance(m, (list, tuple)):
        # versioned format: [format, descriptor, descriptor_version, map]
        print(f"{path}: {m[1]} v{m[2]} (store format {m[0]})")
        return m[3]
    return m  # legacy format: bare map

def detect_mutators(buckets):
    tags = set()
    for bucket in buckets:
        for name in bucket:
            if name.startswith("mut."):
                tags.add(name.split(".")[1])
    return sorted(tags)

def main(path, label):
    m = load(path)
    global MUTATORS
    MUTATORS = detect_mutators(m.values())
    print(f"mutators: {MUTATORS}")
    keys = np.array(list(m.keys()), dtype=np.uint64)
    buckets = list(m.values())
    K = len(keys)

    name2hash = {}
    n_base = np.zeros(K, dtype=np.float64)
    n_mut = {mut: np.zeros(K, dtype=np.float64) for mut in MUTATORS}
    for ki, bucket in enumerate(buckets):
        for name in bucket:
            name2hash[name] = keys[ki]
            if name.startswith("mut."):
                for mut in MUTATORS:
                    if name.startswith(f"mut.{mut}."):
                        n_mut[mut][ki] += 1
                        break
            else:
                n_base[ki] += 1

    bases = [n for n in name2hash if not n.startswith("mut.")]
    N = len(bases)
    q = np.array([name2hash[b] for b in bases], dtype=np.uint64)

    # TP_m(t): distance from each base to its own mutant, cumulative over t
    tp = {}
    for mut in MUTATORS:
        d = np.array(
            [bin(int(name2hash[b]) ^ int(name2hash[f"mut.{mut}.{b}"])).count("1")
             for b in bases])
        tp[mut] = np.cumsum(np.bincount(d, minlength=TMAX)[:TMAX])

    # Histogram of (query, key) distances weighted by bucket composition
    hist_base = np.zeros(TMAX)
    hist_mut = {mut: np.zeros(TMAX) for mut in MUTATORS}
    CHUNK = 512
    for i in range(0, N, CHUNK):
        d = np.bitwise_count(q[i:i + CHUNK, None] ^ keys[None, :]).astype(np.uint8)
        flat = d.ravel()
        sel = flat < TMAX
        flat = flat[sel]
        rows = d.shape[0]
        hist_base += np.bincount(flat, weights=np.broadcast_to(n_base, (rows, K)).ravel()[sel], minlength=TMAX)[:TMAX]
        for mut in MUTATORS:
            hist_mut[mut] += np.bincount(flat, weights=np.broadcast_to(n_mut[mut], (rows, K)).ravel()[sel], minlength=TMAX)[:TMAX]

    cum_base = np.cumsum(hist_base) - N          # exclude self (d=0 always)
    print(f"\n=== {label} ===")
    print(f"{'t':>2} | " + " | ".join(f"{mut:>22}" for mut in MUTATORS))
    print(f"{'':>2} | " + " | ".join(f"{'recall':>10} {'precis':>11}" for _ in MUTATORS))
    curves = {}
    for mut in MUTATORS:
        fp = (np.cumsum(hist_mut[mut]) - tp[mut]) + cum_base
        rec = tp[mut] / N
        prec = np.divide(tp[mut], tp[mut] + fp,
                         out=np.zeros(TMAX), where=(tp[mut] + fp) > 0)
        curves[mut] = (rec, prec)
    for t in range(TMAX):
        row = " | ".join(f"{curves[mut][0][t]:>10.4f} {curves[mut][1][t]:>11.6f}" for mut in MUTATORS)
        print(f"{t:>2} | {row}")
    return curves

if __name__ == "__main__":
    curves_by_store = {}
    for path, label in [("dct.store", "DCT"), ("median.store", "Median")]:
        full = f"/home/mark/workspace/repos/image-similarity/{path}"
        curves_by_store[label] = main(full, label)

    try:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
        fig, axes = plt.subplots(1, 2, figsize=(13, 5.5), sharey=True)
        for ax, (label, curves) in zip(axes, curves_by_store.items()):
            for mut, (rec, prec) in curves.items():
                ax.plot(rec, prec, marker=".", label=mut)
            ax.set_title(f"{label} hash — 24,988 Flickr images, thresholds 0–32")
            ax.set_xlabel("Recall")
            ax.set_ylabel("Precision")
            ax.grid(alpha=.3)
            ax.legend()
        fig.tight_layout()
        fig.savefig("/tmp/imgsim/pr-full.png", dpi=110)
        print("\nplot: /tmp/imgsim/pr-full.png")
    except ImportError:
        print("\nmatplotlib not available; table output only")