#!/usr/bin/env python
"""Sanity check for the 3D PSA code path used by psa_ensemble.py.

Runs the same element-based PSA (N, O, polar H; Bondi radii; 1.4 A probe; mdtraj
shrake_rupley) on grazoprevir, whose solution-accessible 3D PSA is reported as
144-205 A^2 (174 A^2 for the crystal conformer) using a charge-threshold
selection in PyMOL. Agreement means the two conventions land on the same scale,
so PROTAC numbers from psa_ensemble.py can be read against published values.

Needs no finished run: grazoprevir and its conformers are generated here, so it
takes seconds.

Usage:
    python validate_psa.py
"""
import numpy as np, tempfile, os
from rdkit import Chem
from rdkit.Chem import AllChem
import mdtraj as md

# grazoprevir, PubChem CID 44603531 (SMILES field, verbatim)
SMI = "CC(C)(C)[C@H]1C(=O)N2C[C@@H](C[C@H]2C(=O)N[C@@]3(C[C@H]3C=C)C(=O)NS(=O)(=O)C4CC4)OC5=NC6=C(C=CC(=C6)OC)N=C5CCCCC[C@@H]7C[C@H]7OC(=O)N1"

m = Chem.AddHs(Chem.MolFromSmiles(SMI))
p = AllChem.ETKDGv3(); p.randomSeed = 1; p.useRandomCoords = True
cids = AllChem.EmbedMultipleConfs(m, numConfs=25, params=p)
AllChem.MMFFOptimizeMoleculeConfs(m)
print(f"grazoprevir: {m.GetNumAtoms()} atoms, {len(cids)} conformers embedded")

def is_polar(atom):
    s = atom.GetSymbol()
    if s in ("N", "O"): return True
    if s == "H":
        return any(nb.GetSymbol() in ("N", "O") for nb in atom.GetNeighbors())
    return False
polar = [a.GetIdx() for a in m.GetAtoms() if is_polar(a)]
print(f"polar atoms (N,O,polar-H): {len(polar)}")

vals = []
for cid in cids:
    pdb = Chem.MolToPDBBlock(m, confId=cid)
    with tempfile.NamedTemporaryFile("w", suffix=".pdb", delete=False) as fh:
        fh.write(pdb); path = fh.name
    t = md.load(path); os.unlink(path)
    sasa = md.shrake_rupley(t, mode="atom", probe_radius=0.14)[0]  # nm^2
    vals.append(sasa[polar].sum() * 100.0)  # A^2
vals = np.array(vals)
rg = []
for cid in cids:
    conf = m.GetConformer(cid)
    xyz = np.array([list(conf.GetAtomPosition(i)) for i in range(m.GetNumAtoms())])
    xyz -= xyz.mean(0); rg.append(np.sqrt((xyz**2).sum(1).mean()))
rg = np.array(rg)

print(f"\nelement-based 3D PSA over {len(vals)} conformers:")
print(f"  min {vals.min():.0f}  median {np.median(vals):.0f}  max {vals.max():.0f} A2")
print(f"  (most folded conformer, Rg={rg.min():.2f} A): PSA = {vals[np.argmin(rg)]:.0f} A2")
print(f"\nreported (PyMOL, |q|>1 selection): grazoprevir SA 3D PSA = 144-205 A2, crystal 174")
print(f"  grazoprevir 2D TPSA (PubChem) = 204 A2")
inrange = ((vals >= 140) & (vals <= 210)).mean() * 100
print(f"\n-> {inrange:.0f}% of the conformers fall in the reported 144-205 range")