"""One command from a SMILES to ranked targets and, optionally, docked poses.

Stage 1 retrieves with code/reverse_screen.py. Stage 2 resolves each retrieved
UniProt accession to a deposited entry that actually contains the component
that retrieved it, then docks the query into that site with
code/dock_into_site.py.

    python code/reverse_screen_run.py --name Orforglipron \
        --smiles "..." --pools 25 --dock 3 --out results/rs_orforglipron_run.json

--dock 0 stops after retrieval, which is the 44 ms half of the pipeline.
"""
import argparse, json, os, subprocess, sys, time, urllib.request

HERE = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
PY = sys.executable
CODE = os.path.join(HERE, "code")
WORK = os.path.join(HERE, "work")
SEARCH = "https://search.rcsb.org/rcsbsearch/v2/query"


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


def entry_for(comp_id, accession):
    """A deposited entry holding this component and this UniProt accession.

    Sorted by resolution so the best available coordinates define the box.
    """
    acc = accession.replace("UNP:", "")
    q = {"query": {"type": "group", "logical_operator": "and", "nodes": [
        {"type": "terminal", "service": "text",
         "parameters": {"attribute": "rcsb_nonpolymer_entity_container_identifiers."
                                     "nonpolymer_comp_id",
                        "operator": "exact_match", "value": comp_id}},
        {"type": "terminal", "service": "text",
         "parameters": {"attribute": "rcsb_polymer_entity_container_identifiers."
                                     "reference_sequence_identifiers.database_accession",
                        "operator": "exact_match", "value": acc}}]},
        "return_type": "entry",
        "request_options": {"paginate": {"start": 0, "rows": 5},
                            "sort": [{"sort_by": "rcsb_entry_info.resolution_combined",
                                      "direction": "asc"}]}}
    try:
        r = urllib.request.urlopen(urllib.request.Request(
            SEARCH, data=json.dumps(q).encode(),
            headers={"Content-Type": "application/json"}), timeout=120)
        d = json.load(r)
    except Exception as exc:
        log("  entry lookup failed for %s / %s: %s" % (comp_id, acc, exc))
        return None
    ids = [x["identifier"] for x in d.get("result_set", [])]
    return ids[0] if ids else None


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--smiles", required=True)
    ap.add_argument("--name", default="query")
    ap.add_argument("--pools", type=int, default=25)
    ap.add_argument("--top", type=int, default=15)
    ap.add_argument("--truth", default="")
    ap.add_argument("--dock", type=int, default=0,
                    help="how many of the highest ranked retrieved sites to dock")
    ap.add_argument("--exhaustiveness", type=int, default=32)
    ap.add_argument("--cpu", type=int, default=8)
    ap.add_argument("--out", default="")
    a = ap.parse_args()

    t0 = time.time()
    retr = os.path.join(HERE, "work", "_rs_%s.json" % a.name)
    cmd = [PY, os.path.join(CODE, "reverse_screen.py"), "--smiles", a.smiles,
           "--name", a.name, "--pools", str(a.pools), "--top", str(a.top),
           "--out", retr]
    if a.truth:
        cmd += ["--truth", a.truth]
    r = subprocess.run(cmd, capture_output=True, text=True)
    if r.returncode != 0:
        sys.stderr.write(r.stdout + r.stderr)
        sys.exit("retrieval failed")
    res = json.load(open(retr))
    os.unlink(retr)
    res["retrieval_seconds"] = round(time.time() - t0, 2)
    log("retrieved %d sites in %.2f s" % (res["pool_size"], res["retrieval_seconds"]))

    docked = []
    for site in res["pool"][:a.dock]:
        pdb = entry_for(site["via"], site["site"])
        if not pdb:
            docked.append({"site": site["site"], "via": site["via"],
                           "pdb": None, "status": "no entry"})
            continue
        log("docking into %s (%s, via %s)" % (pdb, site["site"], site["via"]))
        t1 = time.time()
        dr = subprocess.run(
            [PY, os.path.join(CODE, "dock_into_site.py"), "--pdb", pdb,
             "--ref-comp", site["via"], "--smiles", a.smiles,
             "--name", "%s_%s" % (a.name, site["site"].replace("UNP:", "")),
             "--exhaustiveness", str(a.exhaustiveness), "--cpu", str(a.cpu)],
            capture_output=True, text=True)
        run_dir = os.path.join(HERE, "docking", "runs", "%s_%s_in_%s" % (
            a.name, site["site"].replace("UNP:", ""), pdb.upper()))
        rj = os.path.join(run_dir, "result.json")
        row = {"site": site["site"], "via": site["via"], "pdb": pdb,
               "tanimoto": site["tanimoto"], "family": site.get("family", ""),
               "seconds": round(time.time() - t1, 1), "dir": run_dir}
        if os.path.isfile(rj):
            d = json.load(open(rj))
            row["scores"] = d.get("scores", [])[:5]
            row["top_score"] = d.get("scores", [None])[0]
            row["box"] = d.get("box")
            row["status"] = "ok"
        else:
            row["status"] = "docking failed"
            sys.stderr.write(dr.stdout[-2000:] + dr.stderr[-2000:])
        docked.append(row)

    res["docked"] = docked
    res["total_seconds"] = round(time.time() - t0, 2)
    if a.out:
        json.dump(res, open(a.out, "w"), indent=1)
        log("wrote %s" % a.out)
    else:
        print(json.dumps(res, indent=1))


if __name__ == "__main__":
    main()