"""Replace the co-occurrence pairing with a contact-verified pairing.

Section 2.2 pairs a ligand with every UniProt accession carried by its PDB
entry. This rebuilds that map from deposited coordinates instead: a ligand is
paired with an accession only when at least one detected interaction of that
ligand touches a chain belonging to that accession.

Inputs, all already on disk and read only:
  contactmaps/work/entries.jsonl    polymer entity uniprot_ids and auth_asym_ids
  contactmaps/images/<shelf>/<PDB>_<LIG>.json   per contact chain
  work/retrieval_index.jsonl        the published index and its ground truth

Outputs:
  pairing_audit.json   coverage and shrinkage, the numbers the editor asked for
  retrieval_index_contact.jsonl   the index rebuilt under the contact rule
"""
import argparse, glob, json, os, sys, time
from collections import defaultdict

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


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("--root", required=True)
    ap.add_argument("--index", required=True)
    ap.add_argument("--outdir", required=True)
    ap.add_argument("--strong-only", action="store_true",
                    help="count only typed interactions, not bare close contacts")
    a = ap.parse_args()
    os.makedirs(a.outdir, exist_ok=True)

    # chain to accession, per entry
    chain_acc = {}
    ent_acc = {}
    with open(os.path.join(a.root, "contactmaps/work/entries.jsonl")) as f:
        for l in f:
            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 {}
                us = ci.get("uniprot_ids") or []
                ch = ci.get("auth_asym_ids") or []
                for u in us:
                    ea.add("UNP:" + u)
                    for c in ch:
                        cm[c].add("UNP:" + u)
            chain_acc[pid] = cm
            ent_acc[pid] = ea
    log("entries with a chain map %s" % format(len(chain_acc), ","))

    types = [t for t in TYPES if not (a.strong_only and t in ("contact", "hydrophobic"))]

    verified = defaultdict(set)          # comp_id -> accessions touched
    seen_comp_entry = set()
    no_chain_map = 0
    no_contacts = 0
    unmapped_chain = 0
    files = glob.glob(os.path.join(a.root, "contactmaps/images/*/*.json"))
    log("sidecars %s" % format(len(files), ","))
    t0 = time.time()
    for i, p in enumerate(files):
        try:
            d = json.load(open(p))
        except Exception:
            continue
        pid = d.get("pdb_id")
        comp = ((d.get("ligand") or {}).get("comp_id"))
        if not pid or not comp:
            continue
        seen_comp_entry.add((comp, pid))
        cm = chain_acc.get(pid)
        if not cm:
            no_chain_map += 1
            continue
        I = d.get("interactions") or {}
        chains = set()
        for t in types:
            for r in (I.get(t) or []):
                c = r.get("chain")
                if c:
                    chains.add(c)
        if not chains:
            no_contacts += 1
            continue
        for c in chains:
            accs = cm.get(c)
            if not accs:
                unmapped_chain += 1
                continue
            verified[comp] |= accs
        if i % 5000 == 0:
            log("  %s of %s, %.0f s" % (format(i, ","), format(len(files), ","),
                                        time.time() - t0))
    log("sidecars read, %.0f s" % (time.time() - t0))

    # apply to the published index
    rows = [json.loads(l) for l in open(a.index)]
    pub_pairs = sum(len(r["sites"]) for r in rows)
    kept_pairs = 0
    kept_rows = []
    dropped_rows = 0
    covered = 0
    out = open(os.path.join(a.outdir, "retrieval_index_contact.jsonl"), "w")
    for r in rows:
        v = verified.get(r["comp_id"])
        if v is None:
            # no coordinates ever examined for this component
            continue
        covered += 1
        keep = sorted(set(r["sites"]) & v)
        if not keep:
            dropped_rows += 1
            continue
        kept_pairs += len(keep)
        r2 = dict(r); r2["sites"] = keep
        out.write(json.dumps(r2) + "\n")
        kept_rows.append(r2)
    out.close()

    pub_sites = len({s for r in rows for s in r["sites"]})
    new_sites = len({s for r in kept_rows for s in r["sites"]})

    # how much of the published map came from multi accession entries
    multi = sum(1 for p in chain_acc if len(ent_acc.get(p, ())) > 1)

    audit = {
        "interaction_types_counted": types,
        "published_index_ligands": len(rows),
        "published_ligand_accession_pairs": pub_pairs,
        "published_target_sites": pub_sites,
        "entries_with_chain_map": len(chain_acc),
        "entries_multi_accession": multi,
        "entries_multi_accession_fraction": round(multi / max(len(chain_acc), 1), 3),
        "sidecars_read": len(files),
        "ligand_entry_pairs_with_coordinates": len(seen_comp_entry),
        "sidecars_without_chain_map": no_chain_map,
        "sidecars_without_any_contact": no_contacts,
        "contacts_on_unmappable_chain": unmapped_chain,
        "index_ligands_with_coordinates": covered,
        "index_ligands_coverage": round(covered / max(len(rows), 1), 3),
        "contact_index_ligands": len(kept_rows),
        "contact_ligand_accession_pairs": kept_pairs,
        "contact_target_sites": new_sites,
        "pairs_retained_fraction": round(kept_pairs / max(pub_pairs, 1), 3),
        "ligands_losing_every_accession": dropped_rows,
    }
    json.dump(audit, open(os.path.join(a.outdir, "pairing_audit.json"), "w"), indent=1)
    for k, v in audit.items():
        print("%-42s %s" % (k, v))


if __name__ == "__main__":
    main()