#!/usr/bin/env python
"""Compare polar-atom selection conventions for 3D PSA on a finished run.

The element-based selection (N, O, polar H) and the charge-threshold selection
used in the reference work give different absolute PSA values on the same
trajectory. This prints both, sweeping the charge threshold, so a gap against
published numbers can be attributed to convention rather than to conformation.
Charges come from the run's AM1-BCC cache, so nothing is recomputed.

Takes a run name rather than a path and resolves <name>/<solvent>/ against the
current directory, so run it from the md folder. Solvent defaults to chloroform.

Usage:
    python psa_method_probe.py egfrvhl_rank4 [solvent]
"""
import json, glob, numpy as np, mdtraj as md
import xml.etree.ElementTree as ET

def cached_charges(cache_path):
    d = json.load(open(cache_path))
    entry = list(list(d.values())[0].values())[0]      # {'openff-2.2.0': {'1': {'ffxml': ...}}}
    ffxml = entry["ffxml"]
    root = ET.fromstring(ffxml)
    # type -> charge from NonbondedForce
    q_by_type = {}
    for a in root.iter("Atom"):
        if a.get("charge") is not None and a.get("type") is not None:
            q_by_type[a.get("type")] = float(a.get("charge"))
    # atom order from the Residue template
    res = root.find(".//Residues/Residue")
    order = [a.get("type") for a in res.findall("Atom")]
    return np.array([q_by_type[t] for t in order])

def main():
    import sys
    NAME = sys.argv[1] if len(sys.argv) > 1 else "egfrvhl_rank4"
    SOL = sys.argv[2] if len(sys.argv) > 2 else "chloroform"
    q = cached_charges(f"{NAME}/{SOL}/ff_cache.json")
    print(f"charges from cache: n={len(q)}  range [{q.min():.2f},{q.max():.2f}]  "
          f"#|q|>1.0={(np.abs(q) > 1.0).sum()}")

    top = glob.glob(f"{NAME}/{SOL}/*_solvated.pdb")[0]
    reps = sorted(glob.glob(f"{NAME}/{SOL}/rep*/{NAME}_prod.dcd"))
    tr = md.join([md.load(r, top=top) for r in reps])
    lig = [a.index for r in tr.topology.residues if not r.is_water and r.n_atoms > 5
           for a in r.atoms]
    t = tr.atom_slice(lig); t.center_coordinates(); L = t.topology
    syms = [a.element.symbol for a in L.atoms]
    if len(syms) != len(q):
        print(f"atom-count mismatch: cache {len(q)} vs MD {len(syms)}"); return
    sasa = md.shrake_rupley(t, mode="atom", probe_radius=0.14)   # nm^2
    def area(mask): return (sasa[:, mask].sum(1) * 100.0).mean()

    nbr = {}
    for b in L.bonds:
        nbr.setdefault(b[0].index, []).append(b[1].index)
        nbr.setdefault(b[1].index, []).append(b[0].index)
    elem = np.array([i for i, a in enumerate(L.atoms)
        if a.element.symbol in ("N", "O") or (a.element.symbol == "H" and
           any(L.atom(j).element.symbol in ("N", "O") for j in nbr.get(a.index, [])))])

    print(f"\nelement selection (N, O, polar H): {area(elem):.0f} A2  ({len(elem)} atoms)")
    print("charge-threshold selection:")
    print(f"  {'|q|>thr':>8s}{'#atoms':>8s}{'PSA':>7s}")
    for thr in [0.10, 0.15, 0.20, 0.25, 0.30, 0.40, 0.50]:
        m = np.abs(q) > thr
        print(f"  {thr:>8.2f}{int(m.sum()):>8d}{area(m):>7.0f}")
    print("\nreference permeable band = 190-265 A2")

if __name__ == "__main__":
    main()