#!/usr/bin/env python
"""3D polar surface area over an MD ensemble.

Element-based PSA: N, O and polar H, Bondi radii, 1.4 A probe, mdtraj
shrake_rupley. validate_psa.py checks this code path against the published
grazoprevir numbers. Lower PSA in the membrane mimic tracks higher passive
permeability.

Polar atoms are indexed from each run's SDF, so the SDF and MD atom orders must
agree. A run whose ligand atom count does not match its SDF is reported and
dropped from the table rather than summed over the wrong atoms.

Usage:
    python psa_ensemble.py egfrvhl_rank4 ms39_ref --solvent chloroform --out psa.npy
"""
import argparse, glob, os
import numpy as np
import mdtraj as md
from rdkit import Chem


def polar_idx_from_sdf(sdf):
    m = Chem.MolFromMolFile(sdf, removeHs=False)
    idx = []
    for a in m.GetAtoms():
        s = a.GetSymbol()
        if s in ("N", "O"):
            idx.append(a.GetIdx())
        elif s == "H" and any(nb.GetSymbol() in ("N", "O") for nb in a.GetNeighbors()):
            idx.append(a.GetIdx())
    return idx, m.GetNumAtoms()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("runs", nargs="+", help="run directories, one per compound")
    ap.add_argument("--solvent", default="chloroform")
    ap.add_argument("--stride", type=int, default=20)
    ap.add_argument("--out", default=None, help="optional .npy dump of the per-frame values")
    a = ap.parse_args()

    print(f"{'compound':16s}{'n_pol':6s}{'3D PSA (A^2)  mean(std)':24s}{'frames':8s}")
    print("-" * 56)
    out = {}
    for run in a.runs:
        name = os.path.basename(run.rstrip("/"))
        pol, nat = polar_idx_from_sdf(os.path.join(run, f"{name}.sdf"))
        tops = glob.glob(os.path.join(run, a.solvent, "*_solvated.pdb"))
        if not tops:
            print(f"{name:16s}  no {a.solvent} topology"); continue
        vals = []
        for dcd in sorted(glob.glob(os.path.join(run, a.solvent, "rep*", f"{name}_prod.dcd"))):
            tr = md.load(dcd, top=tops[0], stride=a.stride)
            lig = [at.index for r in tr.topology.residues
                   if not r.is_water and r.n_atoms > 5 for at in r.atoms]
            sub = tr.atom_slice(lig)
            if sub.n_atoms != nat:
                print(f"{name:16s}  atom count mismatch md={sub.n_atoms} sdf={nat}")
                vals = None
                break
            sasa = md.shrake_rupley(sub, mode="atom", probe_radius=0.14)   # nm^2
            vals.extend((sasa[:, pol].sum(1) * 100.0).tolist())            # A^2 per frame
        if vals:
            v = np.array(vals)
            print(f"{name:16s}{len(pol):<6d}{f'{v.mean():.0f} ({v.std():.0f})':24s}{len(v):<8d}")
            out[name] = dict(mean=float(v.mean()), std=float(v.std()), vals=v.tolist())
    if a.out and out:
        np.save(a.out, out, allow_pickle=True)
        print(f"\nsaved {a.out}")


if __name__ == "__main__":
    main()