"""Reverse screening: one ligand in, the targets it might interact with out.

PharmCast end to end. The query is predicted from 2D structure and every
indexed ligand was predicted the same way, so both sides of the comparison are
the same calculation and no conformer ensemble is generated at query time.

The index is every ligand solved against a target in a structure good enough to
place it, which is X-ray at 2.5 A or better and cryo-EM at 4.0 A or better. The
site fingerprint work keeps its own stricter rule; retrieval does not need the
coordinates, only the association.
"""
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")
import ligand_fingerprints as LF
from rdkit import Chem, RDLogger
from rdkit.Chem import Descriptors
RDLogger.DisableLog("rdApp.*")

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")
_IDX = {}


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


def pack(bits, nb):
    v = np.zeros(nb, dtype=np.uint8)
    for b in bits:
        if b < NBITS:
            v[b >> 3] |= 1 << (7 - (b & 7))
    return v


def load_index(path):
    if path in _IDX:
        return _IDX[path]
    comps, sites, rows, smis = [], [], [], []
    for l in open(path):
        try:
            r = json.loads(l)
        except Exception:
            continue
        comps.append(r["comp_id"]); sites.append(r["sites"]); rows.append(r["pfp"])
        smis.append(r.get("ikey", "") or r.get("smiles", ""))
    nb = (NBITS + 7) // 8
    M = np.zeros((len(rows), nb), dtype=np.uint8)
    for i, bits in enumerate(rows):
        M[i] = pack(bits, nb)
    cnt = POP[M].sum(axis=1, dtype=np.int32)
    _IDX[path] = (comps, sites, M, cnt, smis)
    return _IDX[path]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--smiles", required=True)
    ap.add_argument("--name", default="query")
    ap.add_argument("--index", default=os.path.join(WORK, "retrieval_index.jsonl"))
    ap.add_argument("--model", default=MODEL)
    ap.add_argument("--pools", type=int, default=25)
    ap.add_argument("--top", type=int, default=15)
    ap.add_argument("--exclude", default="",
                    help="component ids to remove from the index for this run")
    ap.add_argument("--truth", default="", help="comma separated UniProt accessions to score against")
    ap.add_argument("--out", default="")
    a = ap.parse_args()

    mol = Chem.MolFromSmiles(a.smiles)
    if mol is None:
        print("SMILES would not parse"); return
    from pharmcast import PharmCast
    pc = PharmCast.load(a.model)
    t0 = time.time()
    words = pc.words_batch([Chem.MolToSmiles(mol)])[0]
    qbits = LF.bits(words)
    t_fp = time.time() - t0

    comps, sites, M, cnt, smis = load_index(a.index)
    nb = M.shape[1]
    qv = pack(qbits, nb); qn = int(POP[qv].sum())
    t0 = time.time()
    inter = POP[np.bitwise_and(M, qv)].sum(axis=1, dtype=np.int32)
    sim = inter / (cnt + qn - inter + 1e-9)
    t_search = time.time() - t0

    # HOLD THE QUERY OUT. A molecule already deposited against its target would
    # otherwise retrieve itself at Tanimoto 1.000, which measures nothing.
    try:
        qcanon = Chem.MolToInchiKey(mol).split("-")[0]
    except Exception:
        qcanon = Chem.MolToSmiles(mol)
    held = []
    for i, sm in enumerate(smis):
        if sm and sm == qcanon:
            held.append(comps[i]); sim[i] = -1.0
    for x in [y.strip() for y in a.exclude.split(",") if y.strip()]:
        for i, c in enumerate(comps):
            if c == x:
                held.append(c); sim[i] = -1.0

    order = np.argsort(-sim)
    nsites = len({s for v in sites for s in v})
    fam = json.load(open(os.path.join(WORK, "chembl_families.json")))

    def famname(t):
        return (fam.get(t.replace("UNP:", "")) or {}).get("name", "")

    truth = {x if x.startswith("UNP:") else "UNP:" + x
             for x in a.truth.split(",") if x.strip()}
    pool, prov = {}, {}
    for i in order[:a.pools]:
        for s in sites[i]:
            if s not in pool:
                pool[s] = round(float(sim[i]), 3); prov[s] = comps[i]
    rank_true, tan_true, via_true = None, None, None
    if truth:
        for r, i in enumerate(order, 1):
            if set(sites[i]) & truth:
                rank_true, tan_true, via_true = r, round(float(sim[i]), 3), comps[i]
                break

    res = {"name": a.name, "smiles": a.smiles,
           "mw": round(Descriptors.MolWt(mol), 1),
           "heavy_atoms": mol.GetNumHeavyAtoms(),
           "query_bits": qn, "indexed_ligands": len(comps), "indexed_sites": nsites,
           "held_out": sorted(set(held)),
           "fingerprint_ms": round(t_fp * 1000, 1),
           "search_ms": round(t_search * 1000, 1),
           "pools": a.pools, "pool_size": len(pool),
           "best_tanimoto": round(float(sim[order[0]]), 3),
           "truth": sorted(truth), "true_rank": rank_true,
           "true_tanimoto": tan_true, "true_via": via_true,
           "true_in_pool": bool(truth & set(pool)),
           "top": [{"comp": comps[i], "tanimoto": round(float(sim[i]), 3),
                    "sites": sites[i], "family": famname(sites[i][0]) if sites[i] else ""}
                   for i in order[:a.top]],
           "pool": [{"site": s, "tanimoto": v, "via": prov[s], "family": famname(s),
                     "is_true": s in truth}
                    for s, v in sorted(pool.items(), key=lambda kv: -kv[1])]}
    print("\n%s  %.1f Da, %d heavy atoms" % (a.name, res["mw"], res["heavy_atoms"]))
    print("index %s ligands, %s target sites" % (format(len(comps), ","), format(nsites, ",")))
    print("fingerprint %.1f ms, search %.1f ms" % (res["fingerprint_ms"], res["search_ms"]))
    if held:
        print("held out of the index (same molecule): %s" % ", ".join(sorted(set(held))))
    print("\nnearest indexed ligands")
    for t in res["top"][:8]:
        mark = "  <== TRUE" if set(t["sites"]) & truth else ""
        print("  %-8s %.3f  %s%s" % (t["comp"], t["tanimoto"], ",".join(t["sites"][:3]), mark))
    if truth:
        print("\ntrue target %s" % ",".join(sorted(truth)))
        print("  first ligand binding it: rank %s of %s, tanimoto %s, via %s"
              % (rank_true, format(len(comps), ","), tan_true, via_true))
        print("  in the pool at %d neighbors: %s (%d sites)"
              % (a.pools, res["true_in_pool"], len(pool)))
    if a.out:
        json.dump(res, open(a.out, "w"), indent=1)
        log("wrote %s" % a.out)


if __name__ == "__main__":
    main()