"""How many molecules per hour can be reverse screened.

Fingerprinting and search are both batch operations. This measures them at
batch scale rather than one molecule at a time, because a virtual library is
screened in batches and the per-molecule figure understates it.
"""
import 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")
RESULTS = os.path.join(HERE, "results")
import ligand_fingerprints as LF
NBITS = 10549
POP = np.array([bin(i).count("1") for i in range(256)], dtype=np.uint16)
# The PharmCast checkpoint. Set PHARMCAST_CHECKPOINT to its path; manifest.json
# carries its sha256 so the one you hold can be matched to the one measured here.
MODEL = os.environ.get("PHARMCAST_CHECKPOINT", "pharmcast_scp_v10.pt")


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


def main():
    idx = os.path.join(WORK, "retrieval_index.jsonl")
    smis, rows, sites = [], [], []
    for l in open(idx):
        r = json.loads(l)
        smis.append(r["smiles"]); rows.append(r["pfp"]); sites.append(r["sites"])
    n = len(rows)
    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)
    NS = len({s for v in sites for s in v})
    log("index %s ligands, %s sites" % (format(n, ","), format(NS, ",")))

    from pharmcast import PharmCast
    pc = PharmCast.load(MODEL)
    rng = np.random.default_rng(7)
    out = {"index_ligands": n, "index_sites": NS, "rows": []}

    for B in (1, 100, 1000, 5000):
        q = [smis[i] for i in rng.choice(n, B, replace=False)]
        pc.words_batch(q[:2])
        t0 = time.time()
        W = pc.words_batch(q)
        t_fp = time.time() - t0
        Q = np.zeros((B, nb), dtype=np.uint8)
        for i, words in enumerate(W):
            v = Q[i]
            for b in LF.bits(words):
                if b < NBITS:
                    v[b >> 3] |= 1 << (7 - (b & 7))
        qc = POP[Q].sum(axis=1, dtype=np.int32)
        t0 = time.time()
        for i in range(B):
            inter = POP[np.bitwise_and(M, Q[i])].sum(axis=1, dtype=np.int32)
            sim = inter / (cnt + qc[i] - inter + 1e-9)
            np.argpartition(-sim, 25)[:25]
        t_s = time.time() - t0
        tot = t_fp + t_s
        row = {"batch": B, "fingerprint_s": round(t_fp, 3), "search_s": round(t_s, 3),
               "per_molecule_ms": round(tot / B * 1000, 2),
               "per_hour": int(3600.0 / (tot / B))}
        out["rows"].append(row)
        log("batch %5d: fp %.3f s, search %.3f s, %.2f ms per molecule, %s per hour"
            % (B, t_fp, t_s, row["per_molecule_ms"], format(row["per_hour"], ",")))

    best = max(r["per_hour"] for r in out["rows"])
    out["best_per_hour"] = best
    out["best_per_day"] = best * 24
    out["comparisons_per_hour"] = best * n
    print("\nReverse screening throughput, one core")
    print("%-8s %14s %13s %16s %16s" % ("batch", "fingerprint", "search", "per molecule", "molecules/hour"))
    for r in out["rows"]:
        print("%-8d %13.3f s %11.3f s %13.2f ms %16s"
              % (r["batch"], r["fingerprint_s"], r["search_s"],
                 r["per_molecule_ms"], format(r["per_hour"], ",")))
    print("\nbest: %s molecules per hour, %s per day"
          % (format(best, ","), format(best * 24, ",")))
    print("that is %s ligand comparisons per hour against the %s ligand index"
          % (format(best * n, ","), format(n, ",")))
    json.dump(out, open(os.path.join(RESULTS, "throughput.json"), "w"), indent=1)


if __name__ == "__main__":
    main()