"""How many molecules per hour can be reverse screened, BOTH ARMS.

WHY THIS FILE EXISTS. `throughput.py` measures one arm. It loads `r["pfp"]`,
packs 10,549 pharmacophore bits, calls PharmCast once and scans that single
matrix. The published figure of 37.6 ms per molecule and about 96,000 per hour
comes from it, and the method it is quoted for has two arms. A two-arm method
priced with a one-arm measurement is the first thing a referee asking for
computational cost will check.

WHAT THE SECOND ARM ACTUALLY COSTS. The Morgan fingerprint is not an extra
fingerprinting step: PharmCast consumes those same 2,048 bits as its input, so
they are computed either way and the arm reuses them. The added cost is one
more Tanimoto scan, over a matrix a fifth as wide. Both are measured and
reported separately so the claim can be checked rather than taken.

MEASURED IN BATCHES, as `throughput.py` is, because a virtual library is
screened in batches and the per-molecule figure understates it. The service's
own per-query timing on a shared Cloud Run vCPU is a different measurement of a
different machine and the two are not interchangeable.
"""
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"))
# The Morgan definition lives in ONE place, the served tree, and is asserted
# byte-identical to its mirror by the service selftest. Importing it here
# rather than restating RADIUS and NBITS is the whole point of that rule.
sys.path.insert(0, os.path.join(HERE, "webapp", "service", "buildcode"))
WORK = os.path.join(HERE, "work")
RESULTS = os.path.join(HERE, "results")
import ligand_fingerprints as LF
import morgan_fp as MFP

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 pack(bitsets, nbits, n):
    nb = (nbits + 7) // 8
    M = np.zeros((n, nb), dtype=np.uint8)
    for i, bits in enumerate(bitsets):
        v = M[i]
        for b in bits:
            if b < nbits:
                v[b >> 3] |= 1 << (7 - (b & 7))
    return M


def main():
    idx = os.path.join(WORK, "retrieval_index.jsonl")
    smis, pfps = [], []
    for l in open(idx):
        r = json.loads(l)
        smis.append(r["smiles"])
        pfps.append(r["pfp"])
    n = len(smis)
    log("index %s ligands" % format(n, ","))

    M = pack(pfps, NBITS, n)
    cnt = POP[M].sum(axis=1, dtype=np.int32)
    log("pharmacophore matrix packed: %s x %d bits" % (format(n, ","), NBITS))

    # The Morgan side of the index. The Sep 6 index predates the Morgan arm, so
    # the bits are recomputed here from the same SMILES the service uses,
    # through the shared definition.
    from rdkit import Chem, RDLogger
    RDLogger.DisableLog("rdApp.*")
    GEN = MFP.generator()      # held and reused, as the index build does
    mbits, nfail = [], 0
    t0 = time.time()
    for s in smis:
        b = MFP.bits_from_smiles(s, GEN)
        if b is None:
            mbits.append([])
            nfail += 1
            continue
        mbits.append(b)
    log("morgan bits for the index in %.1f s, %d components would not parse"
        % (time.time() - t0, nfail))
    Mm = pack(mbits, MFP.NBITS, n)
    cntm = POP[Mm].sum(axis=1, dtype=np.int32)

    from pharmcast import PharmCast
    pc = PharmCast.load(MODEL)
    rng = np.random.default_rng(7)
    out = {"index_ligands": n, "morgan_nbits": MFP.NBITS,
           "pharmacophore_nbits": NBITS, "morgan_unparseable": nfail,
           "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])                      # warm

        # ONE FINGERPRINTING STEP, SHARED. PharmCast's input IS the Morgan
        # fingerprint, so this is measured once and attributed once.
        t0 = time.time()
        QM = pack([MFP.bits_from_smiles(s, GEN) or [] for s in q],
                  MFP.NBITS, B)
        t_mfp = time.time() - t0
        qmc = POP[QM].sum(axis=1, dtype=np.int32)

        t0 = time.time()
        W = pc.words_batch(q)
        t_fp = time.time() - t0
        Q = pack([LF.bits(w) for w in W], NBITS, B)
        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_sp = time.time() - t0

        t0 = time.time()
        for i in range(B):
            interm = POP[np.bitwise_and(Mm, QM[i])].sum(axis=1, dtype=np.int32)
            simm = interm / (cntm + qmc[i] - interm + 1e-9)
            np.argpartition(-simm, 25)[:25]
        t_sm = time.time() - t0

        one = t_mfp + t_fp + t_sp
        two = one + t_sm
        row = {"batch": B,
               "morgan_fingerprint_s": round(t_mfp, 3),
               "pharmcast_prediction_s": round(t_fp, 3),
               "pharmacophore_scan_s": round(t_sp, 3),
               "morgan_scan_s": round(t_sm, 3),
               "one_arm_ms": round(one / B * 1000, 2),
               "two_arm_ms": round(two / B * 1000, 2),
               "one_arm_per_hour": int(3600.0 / (one / B)),
               "two_arm_per_hour": int(3600.0 / (two / B))}
        out["rows"].append(row)
        log("batch %5d: one arm %.2f ms (%s/h), two arms %.2f ms (%s/h)"
            % (B, row["one_arm_ms"], format(row["one_arm_per_hour"], ","),
               row["two_arm_ms"], format(row["two_arm_per_hour"], ",")))

    best1 = max(r["one_arm_per_hour"] for r in out["rows"])
    best2 = max(r["two_arm_per_hour"] for r in out["rows"])
    fastest = min(r["two_arm_ms"] for r in out["rows"])
    out["best_one_arm_per_hour"] = best1
    out["best_two_arm_per_hour"] = best2
    out["best_two_arm_ms"] = fastest

    print("\nReverse screening throughput, one core, both arms")
    print("%-7s %12s %12s %12s %11s %11s %11s %16s"
          % ("batch", "morgan fp", "pharmcast", "pharm scan", "morgan scan",
             "one arm", "two arms", "two arms/hour"))
    for r in out["rows"]:
        print("%-7d %10.3f s %10.3f s %10.3f s %9.3f s %8.2f ms %8.2f ms %16s"
              % (r["batch"], r["morgan_fingerprint_s"],
                 r["pharmcast_prediction_s"], r["pharmacophore_scan_s"],
                 r["morgan_scan_s"], r["one_arm_ms"], r["two_arm_ms"],
                 format(r["two_arm_per_hour"], ",")))
    print("\nBOTH ARMS: %.2f ms per molecule, %s molecules per hour"
          % (fastest, format(best2, ",")))
    print("one arm for comparison: %s per hour" % format(best1, ","))
    json.dump(out, open(os.path.join(RESULTS, "throughput_two_arm.json"), "w"),
              indent=1)
    print("\nwrote results/throughput_two_arm.json")


if __name__ == "__main__":
    main()