"""Widen the retrieval index to structures the site work excludes.

The site fingerprints need 1.8 A X-ray with a low R free, because a
pseudo-ligand point is only as good as the coordinates under it. RETRIEVAL
needs neither. It needs a ligand fingerprint and the target that ligand was
solved against, and a 3.2 A cryo-EM structure supplies both.

That distinction matters because entire target classes are cryo-EM only. The
GLP-1 receptor has 31 entries carrying a drug-like ligand and not one of them
is X-ray, so under the site rule the receptor does not exist.

This fetches the cryo-EM entries, maps their ligands to UniProt accessions, and
writes the pairs. Fingerprints for components new to the index are built
separately from the chemical component dictionary.
"""
import argparse, json, os, sys, time, urllib.request
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")
SEARCH = "https://search.rcsb.org/rcsbsearch/v2/query"
GRAPHQL = "https://data.rcsb.org/graphql"


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


def post(url, payload, tries=4):
    for a in range(tries):
        try:
            r = urllib.request.urlopen(urllib.request.Request(
                url, data=json.dumps(payload).encode(),
                headers={"Content-Type": "application/json"}), timeout=300)
            return json.load(r)
        except Exception as exc:
            if a == tries - 1:
                raise
            log("  retry %d after %s" % (a + 1, exc))
            time.sleep(3 * (a + 1))


def t(attr, op, val):
    return {"type": "terminal", "service": "text",
            "parameters": {"attribute": attr, "operator": op, "value": val}}


Q = """{ entries(entry_ids: %s) {
 rcsb_id
 rcsb_entry_info { resolution_combined experimental_method }
 polymer_entities {
   rcsb_polymer_entity_container_identifiers { uniprot_ids }
 }
 nonpolymer_entities {
   rcsb_nonpolymer_entity_container_identifiers { nonpolymer_comp_id }
   nonpolymer_comp { chem_comp { id formula_weight } }
 }
} }"""


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--method", default="ELECTRON MICROSCOPY")
    ap.add_argument("--max-res", type=float, default=4.0)
    ap.add_argument("--max-rfree", type=float, default=0.0)
    ap.add_argument("--mw-lo", type=float, default=250.0)
    ap.add_argument("--mw-hi", type=float, default=900.0)
    ap.add_argument("--out", default=os.path.join(WORK, "extra_sites.json"))
    a = ap.parse_args()

    nodes = [t("exptl.method", "exact_match", a.method),
             t("rcsb_entry_info.resolution_combined", "less_or_equal", a.max_res),
             t("chem_comp.formula_weight", "range",
               {"from": a.mw_lo, "to": a.mw_hi})]
    if a.max_rfree:
        nodes.append(t("refine.ls_R_factor_R_free", "less_or_equal", a.max_rfree))
    ids, start, rows = [], 0, 10000
    while True:
        d = post(SEARCH, {"query": {"type": "group", "logical_operator": "and",
                                    "nodes": nodes},
                          "return_type": "entry",
                          "request_options": {"paginate": {"start": start, "rows": rows},
                                              "results_content_type": ["experimental"]}})
        got = [r["identifier"] for r in d.get("result_set", [])]
        ids += got
        total = d.get("total_count", 0)
        log("  entries %s of %s" % (format(len(ids), ","), format(total, ",")))
        start += rows
        if len(ids) >= total or not got:
            break

    pairs = {}     # comp_id -> set of UNP accessions
    where = {}     # (comp_id, unp) -> a pdb id that holds it
    res = {}       # pdb id -> {res, method}
    B = 250
    for i in range(0, len(ids), B):
        chunk = ids[i:i + B]
        d = post(GRAPHQL, {"query": Q % json.dumps(chunk)})
        for e in (d.get("data", {}).get("entries") or []):
            if not e:
                continue
            pid = e["rcsb_id"]
            # The resolution the query already returns, kept per entry: it is
            # what ranks a promiscuous ligand's sites in the service.
            info = e.get("rcsb_entry_info") or {}
            rr = info.get("resolution_combined") or []
            if rr:
                res[pid] = {"res": round(min(rr), 2),
                            "method": info.get("experimental_method") or ""}
            unps = set()
            for pe in (e.get("polymer_entities") or []):
                for u in ((pe.get("rcsb_polymer_entity_container_identifiers") or {})
                          .get("uniprot_ids") or []):
                    unps.add("UNP:" + u)
            if not unps:
                continue
            for ne in (e.get("nonpolymer_entities") or []):
                cc = ((ne.get("nonpolymer_comp") or {}).get("chem_comp") or {})
                cid, mw = cc.get("id"), cc.get("formula_weight")
                if not cid or mw is None:
                    continue
                if not (a.mw_lo <= float(mw) <= a.mw_hi):
                    continue
                for u in unps:
                    pairs.setdefault(cid, set()).add(u)
                    where.setdefault(cid + "|" + u, pid)
        if (i // B) % 8 == 0:
            log("  mapped %s of %s entries, %s components"
                % (format(min(i + B, len(ids)), ","), format(len(ids), ","),
                   format(len(pairs), ",")))
    out = {"method": a.method, "max_res": a.max_res, "entries": len(ids),
           "components": len(pairs),
           "sites": len({u for v in pairs.values() for u in v}),
           "map": {c: sorted(v) for c, v in pairs.items()},
           "where": where,
           "res": res}
    json.dump(out, open(a.out, "w"))
    log("components %s, sites %s, wrote %s"
        % (format(len(pairs), ","), format(out["sites"], ","), a.out))


if __name__ == "__main__":
    main()