"""Recall and cost of the reverse screen, on the PharmCast index.

Every query is one indexed ligand. It is held out of the index by the
connectivity layer of its InChIKey, so a query never retrieves a site through
itself or through a stereoisomer of itself. A hit is any of the query's own
target sites appearing in the retrieved pool.
"""
import argparse, json, os, sys, time
import numpy as np
HERE = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, os.path.join(HERE, "code"))
WORK = os.path.join(HERE, "work")
NBITS = 10549
POP = np.array([bin(i).count("1") for i in range(256)], dtype=np.uint16)
SEC_PER_SITE = 5.1


def log(m):
    sys.stderr.write("[%s] %s\n" % (time.strftime("%H:%M:%S"), m)); sys.stderr.flush()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--index", default=os.path.join(WORK, "retrieval_index.jsonl"))
    ap.add_argument("--sample", type=int, default=3000)
    ap.add_argument("--pools", default="1,5,10,25,50")
    ap.add_argument("--out", default=os.path.join(HERE, "results", "bench_retrieval.json"))
    a = ap.parse_args()

    comps, sites, keys, rows = [], [], [], []
    for l in open(a.index):
        try:
            r = json.loads(l)
        except Exception:
            continue
        comps.append(r["comp_id"]); sites.append(set(r["sites"]))
        keys.append(r.get("ikey", "")); rows.append(r["pfp"])
    n = len(comps)
    nb = (NBITS + 7) // 8
    M = np.zeros((n, nb), dtype=np.uint8)
    for i, bits in enumerate(rows):
        v = M[i]
        for b in bits:
            if b < NBITS:
                v[b >> 3] |= 1 << (7 - (b & 7))
    cnt = POP[M].sum(axis=1, dtype=np.int32)
    allsites = sorted({s for v in sites for s in v})
    NS = len(allsites)
    log("index %s ligands, %s target sites" % (format(n, ","), format(NS, ",")))

    # ligands sharing a connectivity key are held out together
    kidx = {}
    for i, k in enumerate(keys):
        kidx.setdefault(k or comps[i], []).append(i)

    rng = np.random.default_rng(20260906)
    samp = rng.choice(n, min(a.sample, n), replace=False)
    pools = [int(x) for x in a.pools.split(",")]
    res = {P: [0, 0, []] for P in pools}
    t0 = time.time()
    for c, qi in enumerate(samp):
        q = int(qi)
        truth = sites[q]
        if not truth:
            continue
        inter = POP[np.bitwise_and(M, M[q])].sum(axis=1, dtype=np.int32)
        sim = inter / (cnt + cnt[q] - inter + 1e-9)
        for j in kidx.get(keys[q] or comps[q], [q]):
            sim[j] = -1.0
        order = np.argsort(-sim)
        for P in pools:
            pool = set()
            for i in order[:P]:
                pool |= sites[i]
            r = res[P]
            r[0] += len(pool & truth) > 0
            r[1] += 1
            r[2].append(len(pool))
        if c % 500 == 0:
            log("  %s of %s, %.0f s" % (format(c, ","), format(len(samp), ","),
                                        time.time() - t0))

    rows_out = []
    for P in pools:
        hit, tot, ps = res[P]
        mps = float(np.mean(ps)) if ps else 0.0
        rows_out.append({"neighbors": P, "recall": round(hit / max(tot, 1), 3),
                         "mean_sites": round(mps, 1),
                         "dock_seconds": round(mps * SEC_PER_SITE, 1),
                         "reduction": round(NS / max(mps, 1e-9), 1)})
    out = {"index_ligands": n, "index_sites": NS, "queries": int(len(samp)),
           "sec_per_site": SEC_PER_SITE,
           "exhaustive_hours": round(NS * SEC_PER_SITE / 3600, 1), "rows": rows_out}
    os.makedirs(os.path.dirname(a.out), exist_ok=True)
    json.dump(out, open(a.out, "w"), indent=1)
    print("\nReverse screen, %s ligands, %s target sites, %s held-out queries\n"
          % (format(n, ","), format(NS, ","), format(len(samp), ",")))
    print("%-11s %8s %12s %13s %11s" % ("neighbors", "recall", "sites", "dock time", "reduction"))
    for r in rows_out:
        t = r["dock_seconds"]
        ts = "%.0f s" % t if t < 120 else ("%.1f min" % (t / 60) if t < 3600 else "%.1f h" % (t / 3600))
        print("%-11d %8.3f %12.1f %13s %10.0fx"
              % (r["neighbors"], r["recall"], r["mean_sites"], ts, r["reduction"]))
    print("\nexhaustive docking of all %s sites: %.1f hours" % (format(NS, ","), out["exhaustive_hours"]))


if __name__ == "__main__":
    main()