"""Cross product inflation, measured only where coordinates were examined.

Conditioning on a contact map existing removes the bias in a whole index
rebuild, where an accession can be dropped merely because no sidecar covers the
entry that would have verified it.

Unit is the (ligand, entry, accession) triple. For every sidecar, the entry's
accessions come from contactmaps/work/entries.jsonl and the contacted
accessions from the chains in the sidecar's own interactions.
"""
import glob, json, os, sys, time
from collections import defaultdict
import numpy as np

TYPES = ("hbond", "halogen", "salt_bridge", "pi_stacking", "cation_pi",
         "metal", "hydrophobic", "contact", "water_bridge")

chain_acc, ent_acc = {}, {}
for l in open("contactmaps/work/entries.jsonl"):
    try:
        d = json.loads(l)
    except Exception:
        continue
    pid = d.get("rcsb_id")
    if not pid:
        continue
    cm, ea = defaultdict(set), set()
    for pe in (d.get("polymer_entities") or []):
        ci = pe.get("rcsb_polymer_entity_container_identifiers") or {}
        for u in (ci.get("uniprot_ids") or []):
            ea.add("UNP:" + u)
            for c in (ci.get("auth_asym_ids") or []):
                cm[c].add("UNP:" + u)
    chain_acc[pid] = cm; ent_acc[pid] = ea

trip_tot = trip_hit = 0
multi_tot = multi_hit = 0
single_tot = single_hit = 0
ent_single = ent_multi = 0
nocontact = 0
files = glob.glob("contactmaps/images/*/*.json")
t0 = time.time()
for i, p in enumerate(files):
    try:
        d = json.load(open(p))
    except Exception:
        continue
    pid = d.get("pdb_id")
    ea = ent_acc.get(pid)
    cm = chain_acc.get(pid)
    if not ea or not cm:
        continue
    I = d.get("interactions") or {}
    chains = {r.get("chain") for t in TYPES for r in (I.get(t) or []) if r.get("chain")}
    if not chains:
        nocontact += 1
        continue
    touched = set()
    for c in chains:
        touched |= cm.get(c, set())
    k = len(ea)
    hit = len(ea & touched)
    trip_tot += k; trip_hit += hit
    if k > 1:
        ent_multi += 1; multi_tot += k; multi_hit += hit
    else:
        ent_single += 1; single_tot += k; single_hit += hit
    if i % 10000 == 0:
        sys.stderr.write("  %d of %d, %.0f s\n" % (i, len(files), time.time() - t0))

out = {
 "sidecars_used": ent_single + ent_multi,
 "sidecars_without_contacts": nocontact,
 "entries_single_accession": ent_single,
 "entries_multi_accession": ent_multi,
 "triples_all": trip_tot, "triples_contact_supported": trip_hit,
 "supported_fraction_all": round(trip_hit / max(trip_tot, 1), 3),
 "triples_multi_accession_entries": multi_tot,
 "supported_fraction_multi_accession_entries": round(multi_hit / max(multi_tot, 1), 3),
 "triples_single_accession_entries": single_tot,
 "supported_fraction_single_accession_entries": round(single_hit / max(single_tot, 1), 3),
}
json.dump(out, open(sys.argv[1], "w"), indent=1)
for k, v in out.items():
    print("%-48s %s" % (k, v))