"""Reproduce the MolProbity worked example in the ProteinIQ MolProbity score guide.

Downloads each structure listed in STRUCTURES from RCSB PDB or the AlphaFold
Protein Structure Database, runs CCTBX's aggregate MolProbity validation
(mmtbx.validation.molprobity) with Reduce and Probe, and writes the headline
numbers to molprobity_results.csv. It then recomputes the MolProbity score
from clashscore, rotamer outliers and Ramachandran favored with the published
formula, so the composite can be checked by hand.

Environment used for the published numbers (same pins as ProteinIQ's
MolProbity tool):
  cctbx-base 2026.4 (conda-forge), Python 3.11
  probe   rlabduke/probe          835c2c63dd14b8c30599900b07c2296925cc933c
  reduce  rlabduke/reduce         d723303c0ae7a9991a4f723b74f9bfa268d33e87
  Top8000 rlabduke/reference_data cbe300663a2bb55015cf9e9978c720cb12dbf52f
  CCP4 monomer library            15e7134789cc8f835c725eb13483ec61338a30ca
Copy the Top8000 rotamer, Ramachandran, CaBLAM and Rama-Z grids into
cctbx's chem_data, run mmtbx.rebuild_rotarama_cache and
mmtbx.rebuild_cablam_cache, put probe and reduce on PATH (CCTBX calls Reduce
as molprobity.reduce, so add a symlink with that name), set
MMTBX_CCP4_MONOMER_LIB to the monomer library checkout and run
prepare_monomer_library.py once.

Usage: python run_molprobity.py [output_dir]
"""

import csv
import io
import math
import shutil
import sys
import urllib.request
from pathlib import Path

STRUCTURES = [
    # (label, description, download URL); all three are sperm whale myoglobin
    ("1MBN", "X-ray 2.0 Å (1973)", "https://files.rcsb.org/download/1MBN.pdb"),
    ("1A6M", "X-ray 1.0 Å (1998)", "https://files.rcsb.org/download/1A6M.pdb"),
    ("AF-P02185-F1", "AlphaFold DB model v6", "https://alphafold.ebi.ac.uk/files/AF-P02185-F1-model_v6.pdb"),
]


def molprobity_score(clashscore, rota_out, rama_fav):
    """MolProbity score as implemented in mmtbx.validation.utils."""
    return (
        0.426 * math.log(1 + clashscore)
        + 0.33 * math.log(1 + max(0.0, rota_out - 1))
        + 0.25 * math.log(1 + max(0.0, (100.0 - rama_fav) - 2))
        + 0.5
    )


def patch_probe_lookup():
    """Point CCTBX at a probe binary on PATH (as ProteinIQ's worker does)."""
    import libtbx
    import libtbx.load_env  # noqa: F401  (populates libtbx.env)

    probe_path = shutil.which("probe")
    if not probe_path:
        raise RuntimeError("probe is not on PATH")
    if not shutil.which("molprobity.reduce"):
        raise RuntimeError("molprobity.reduce is not on PATH (symlink it to reduce)")
    has_module = libtbx.env.has_module
    under_build = libtbx.env.under_build
    libtbx.env.has_module = lambda name: True if name == "probe" else has_module(name)
    libtbx.env.under_build = lambda path, return_relocatable_path=False: (
        probe_path
        if path.replace("\\", "/") == "probe/exe/probe"
        else under_build(path, return_relocatable_path=return_relocatable_path)
    )


def validate(path):
    import iotbx.pdb
    import mmtbx.model
    from libtbx.utils import null_out
    from mmtbx.validation import molprobity

    model = mmtbx.model.manager(model_input=iotbx.pdb.input(file_name=str(path)), log=null_out())
    model.process(make_restraints=True)
    result = molprobity.molprobity(
        model=model,
        keep_hydrogens=True,
        outliers_only=True,
        file_name=str(path),
        flags=molprobity.molprobity_flags(),
    )
    report = io.StringIO()
    result.show(out=report, outliers_only=True, show_percentiles=True)
    return result, report.getvalue()


def main():
    out_dir = Path(sys.argv[1] if len(sys.argv) > 1 else ".")
    out_dir.mkdir(parents=True, exist_ok=True)
    patch_probe_lookup()
    rows = []
    for label, kind, url in STRUCTURES:
        path = out_dir / f"{label}.pdb"
        if not path.exists():
            request = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
            with urllib.request.urlopen(request) as response:
                path.write_bytes(response.read())
        result, report = validate(path)
        (out_dir / f"{label}_molprobity.txt").write_text(report, encoding="utf-8")
        clash = result.clashscore()
        rota = result.rota_outliers()
        fav = result.rama_favored()
        geometry = result.model_statistics_geometry_result
        cablam = getattr(geometry, "cablam", None)
        rows.append({
            "structure": label,
            "type": kind,
            "clashscore": round(clash, 2),
            "rama_favored_pct": round(fav, 2),
            "rama_allowed_pct": round(result.rama_allowed(), 2),
            "rama_outliers_pct": round(result.rama_outliers(), 2),
            "rotamer_outliers_pct": round(rota, 2),
            "cbeta_deviations": result.cbeta_outliers(),
            "cablam_outliers_pct": round(cablam.outliers, 2) if cablam else "",
            "cablam_disfavored_pct": round(cablam.disfavored, 2) if cablam else "",
            "cablam_ca_outliers_pct": round(cablam.ca_outliers, 2) if cablam else "",
            "molprobity_score": round(result.molprobity_score(), 2),
            "molprobity_score_recomputed": round(molprobity_score(clash, rota, fav), 2),
        })
        print(rows[-1])
    with open(out_dir / "molprobity_results.csv", "w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(rows[0]))
        writer.writeheader()
        writer.writerows(rows)


if __name__ == "__main__":
    main()