"""The retrieval index, built with PharmCast.

Retrieval needs two things: a pharmacophore fingerprint for every ligand ever
solved against a target, and the targets it was solved against. It does not
need the coordinates. That distinction is what lets the retrieval index admit
structures the site fingerprints cannot use, and cryo-EM only receptor classes
are the reason it matters.

PharmCast predicts the fingerprint from 2D structure, so a component costs
microseconds instead of the seconds a 100 conformer ensemble takes. Every
ligand in the index and every query are predicted the same way, so both sides
of a comparison are the same calculation.
"""
import argparse, json, os, sys, time
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 ccd as CCD
import ligand_fingerprints as LF
# The Morgan fingerprint is defined in exactly one place. See morgan_fp.py for
# why, and selftest.py for the pinned hash that keeps it that way.
import morgan_fp as MFP
from rdkit import Chem, RDLogger
from rdkit.Chem import Descriptors
RDLogger.DisableLog("rdApp.*")

# 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():
    ap = argparse.ArgumentParser()
    ap.add_argument("--map", default=os.path.join(WORK, "retrieval_map.json"))
    ap.add_argument("--blocks", default=os.path.join(WORK, "ccd_blocks_all.jsonl"))
    ap.add_argument("--model", default=MODEL)
    ap.add_argument("--batch", type=int, default=512)
    ap.add_argument("--out", default=os.path.join(WORK, "retrieval_index.jsonl"))
    a = ap.parse_args()

    m = json.load(open(a.map))["map"]
    want = sorted(m)
    log("components wanted %s" % format(len(want), ","))

    if not os.path.exists(a.blocks):
        log("extracting component definitions, one pass over the dictionary")
        CCD.extract(want, a.blocks)
    blocks = CCD.load(a.blocks)
    log("component definitions available %s" % format(len(blocks), ","))

    smis, ids = [], []
    bad = 0
    # One Morgan generator for the whole build. Constructing one per molecule
    # would be correct and wasteful.
    _mgen = MFP.generator()
    for c in want:
        blk = blocks.get(c)
        if not blk:
            bad += 1; continue
        try:
            mol, _ = LF.mol_from_block(blk)
            if mol is None:
                bad += 1; continue
            mw = Descriptors.MolWt(mol)
            if not (250.0 <= mw <= 900.0):
                bad += 1; continue
            s = Chem.MolToSmiles(mol)
            if not s:
                bad += 1; continue
            # the connectivity layer of the InChIKey, so a query can be held
            # out of its own index even when the deposited component differs
            # in stereo perception or in SMILES canonicalization
            try:
                key = Chem.MolToInchiKey(mol).split("-")[0]
            except Exception:
                key = ""
            smis.append(s)
            # MORGAN IS COMPUTED HERE, FROM THIS SAME PARSED MOLECULE, rather
            # than from the SMILES again at write time. One parse, so the
            # Morgan bits, the InChIKey and the heavy-atom count can never
            # describe three slightly different readings of the same string.
            ids.append((c, round(mw, 2), mol.GetNumHeavyAtoms(), key,
                        MFP.bits(mol, _mgen)))
        except Exception:
            bad += 1
    log("buildable %s, skipped %s" % (format(len(smis), ","), format(bad, ",")))

    from pharmcast import PharmCast
    pc = PharmCast.load(a.model)
    fh = open(a.out, "w")
    t0 = time.time()
    n = 0
    for i in range(0, len(smis), a.batch):
        chunk = smis[i:i + a.batch]
        try:
            W = pc.words_batch(chunk)
        except Exception as exc:
            log("  batch %d failed: %s" % (i, exc))
            continue
        for k, words in enumerate(W):
            c, mw, ha, key, mbits = ids[i + k]
            fh.write(json.dumps({"comp_id": c, "mw": mw, "heavy_atoms": ha,
                                 "smiles": chunk[k], "ikey": key,
                                 "pfp": LF.bits(words),
                                 "morgan": mbits,
                                 "sites": m[c]}) + "\n")
            n += 1
        if (i // a.batch) % 10 == 0:
            log("  %s of %s, %.1f s" % (format(i + len(chunk), ","),
                                        format(len(smis), ","), time.time() - t0))
    fh.close()
    log("index written: %s components in %.1f s -> %s"
        % (format(n, ","), time.time() - t0, a.out))


if __name__ == "__main__":
    main()