#!/usr/bin/env python3
"""Count how many approved small-molecule drugs trigger PAINS alerts.

Data for the ProteinIQ guide /guides/pains-filters.

Population
  The 2,104 approved drugs of the Lipinski analysis
  (scripts/lipinski-ro5-chembl.py), which this script imports unchanged: ChEMBL
  parent compounds with max_phase 4, selected with the same rules as
  ProteinIQ's approved small-molecule drug library. The import reuses its
  ChEMBL download and cache, so both guides count the same drugs.

Alerts (RDKit FilterCatalog, run on the ChEMBL parent structure as parsed from
its canonical SMILES, without explicit hydrogens, as ProteinIQ's PAINS filter
does)
  * PAINS_A, PAINS_B, PAINS_C: the 480 substructures of Baell and Holloway
    (2010) in RDKit's SMARTS translation (16 + 55 + 409 entries). A drug is
    flagged when it matches at least one.
  * BRENK: the 105 unwanted-group rules of Brenk et al. (2008), as a
    comparison.
  * Family: the alert name without its parenthesized class size and its final
    letter suffix, so anil_di_alk_A(478) and anil_di_alk_C(246) both count as
    anil_di_alk. A drug counts once per family.

Outputs (public/data/guides/pains-filters/)
  approved-drugs-pains.csv  one row per drug
  summary.json              every aggregate quoted in the guide
  analyze.py                a copy of this script

Run with a Python that has RDKit (tested with RDKit 2025.03.5):
  python scripts/pains-alerts-approved-drugs.py [--cache-dir DIR]
"""

from __future__ import annotations

import argparse
import collections
import csv
import importlib.util
import json
import math
import re
import shutil
import tempfile
from pathlib import Path

import rdkit
from rdkit import Chem, RDLogger
from rdkit.Chem.FilterCatalog import FilterCatalog, FilterCatalogParams

RDLogger.DisableLog("rdApp.*")

REPO_ROOT = Path(__file__).resolve().parents[1]
OUT_DIR = REPO_ROOT / "public/data/guides/pains-filters"
LIPINSKI_SCRIPT = REPO_ROOT / "scripts/lipinski-ro5-chembl.py"

PAINS_CATALOGS = ("PAINS_A", "PAINS_B", "PAINS_C")
MIN_CLASS_SIZE = 30

# Plain-language names for the families that appear in approved drugs.
FAMILY_LABELS = {
    "catechol": "Catechols",
    "quinone": "Quinones",
    "anil_di_alk": "Dialkylanilines",
    "indol_3yl_alk": "3-Alkylindoles",
    "azo": "Azo compounds",
    "het_thio_666": "Phenothiazines",
    "mannich": "Phenolic Mannich bases",
    "keto_keto_beta": "Cyclic 1,3-diketones",
    "anthranil_one": "Aminoaryl ketones",
    "dyes5A": "Fused pyrrole dyes",
    "keto_naphthol": "Keto-naphthols",
    "hzone_phenol": "Phenolic hydrazones",
    "hzone_pipzn": "Piperazinyl hydrazones",
    "amino_acridine": "Aminoacridines",
    "styrene": "Styrenes",
    "ene_rhod": "Ene-rhodanines",
    "ene_five_het": "Ene five-membered heterocycles",
    "imine_one": "Imine ketones",
    "imine_one_fives": "Five-membered imine ketones",
    "pyrrole": "Pyrroles",
    "het_pyridiniums": "Pyridinium heterocycles",
    "anil_no_alk": "Anilines (no alkyl)",
    "steroid": "Steroid enones",
    "imidazole": "Imidazoles",
}


def load_lipinski_module():
    spec = importlib.util.spec_from_file_location("lipinski_ro5_chembl", LIPINSKI_SCRIPT)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def catalog(name: str) -> FilterCatalog:
    params = FilterCatalogParams()
    params.AddCatalog(getattr(FilterCatalogParams.FilterCatalogs, name))
    return FilterCatalog(params)


def alert_name(description: str) -> str:
    return re.sub(r"\(\d+\)$", "", description)


def family_of(description: str) -> str:
    return re.sub(r"_[A-Z]$", "", alert_name(description))


def pct(part: int, whole: int) -> float:
    return round(100 * part / whole, 1) if whole else 0.0


def wilson(part: int, whole: int, z: float = 1.96) -> list[float]:
    """95% Wilson score interval, in percent."""
    if not whole:
        return [0.0, 0.0]
    p = part / whole
    centre = (p + z * z / (2 * whole)) / (1 + z * z / whole)
    half = z * math.sqrt(p * (1 - p) / whole + z * z / (4 * whole * whole)) / (1 + z * z / whole)
    return [round(100 * (centre - half), 1), round(100 * (centre + half), 1)]


def share(rows: list[dict], key: str = "pains") -> dict:
    n = len(rows)
    flagged = sum(1 for r in rows if r[key])
    return {"n": n, "flagged": flagged, "pct": pct(flagged, n), "ci95": wilson(flagged, n)}


def build_rows(lipinski, chembl: dict) -> list[dict]:
    ro5_rows, _ = lipinski.build_rows(chembl)
    catalogs = {name: catalog(name) for name in (*PAINS_CATALOGS, "BRENK")}
    rows = []
    for base in ro5_rows:
        smiles = chembl["by_id"][base["chembl_id"]]["molecule_structures"]["canonical_smiles"]
        mol = Chem.MolFromSmiles(smiles)
        matches = {name: sorted({e.GetDescription() for e in cat.GetMatches(mol)}) for name, cat in catalogs.items()}
        pains = [(name, d) for name in PAINS_CATALOGS for d in matches[name]]
        families = sorted({family_of(d) for _, d in pains})
        rows.append(
            {
                "chembl_id": base["chembl_id"],
                "name": base["name"],
                "first_approval": base["first_approval"],
                "oral": base["oral"],
                "withdrawn": base["withdrawn"],
                "atc_level2": base["atc_level2"],
                "therapeutic_class": base["therapeutic_class"],
                "pains": bool(pains),
                "pains_a": bool(matches["PAINS_A"]),
                "pains_b": bool(matches["PAINS_B"]),
                "pains_c": bool(matches["PAINS_C"]),
                "pains_alerts": ";".join(alert_name(d) for _, d in pains),
                "pains_families": ";".join(families),
                "brenk": bool(matches["BRENK"]),
                "brenk_alerts": ";".join(matches["BRENK"]),
            }
        )
    return rows


def summarize(rows: list[dict], chembl: dict, lipinski) -> dict:
    flagged = [r for r in rows if r["pains"]]

    family_drugs: dict[str, list[dict]] = collections.defaultdict(list)
    for row in flagged:
        for family in row["pains_families"].split(";"):
            family_drugs[family].append(row)
    families = sorted(family_drugs.items(), key=lambda item: (-len(item[1]), item[0]))

    alert_counts = collections.Counter(a for r in flagged for a in r["pains_alerts"].split(";"))
    alert_catalog = {}
    for name in PAINS_CATALOGS:
        cat = catalog(name)
        for index in range(cat.GetNumEntries()):
            alert_catalog[alert_name(cat.GetEntryWithIdx(index).GetDescription())] = name

    decades = collections.defaultdict(list)
    for row in rows:
        if row["first_approval"]:
            decades[lipinski.decade_label(int(row["first_approval"]))].append(row)
    by_decade = [
        {"decade": label, **share(decades[label]), "brenk": share(decades[label], "brenk")}
        for label in sorted(decades, key=lambda d: (d != "Before 1950", d))
    ]

    by_class_rows = collections.defaultdict(list)
    for row in rows:
        if row["atc_level2"]:
            by_class_rows[(row["atc_level2"], row["therapeutic_class"])].append(row)
    by_class = sorted(
        (
            {"atc": code, "class": label, **share(members)}
            for (code, label), members in by_class_rows.items()
            if len(members) >= MIN_CLASS_SIZE
        ),
        key=lambda item: (-item["pct"], item["atc"]),
    )

    brenk_counts = collections.Counter(a for r in rows if r["brenk"] for a in r["brenk_alerts"].split(";"))
    alert_counts_per_drug = collections.Counter(len(r["pains_alerts"].split(";")) for r in flagged)
    dated = [r for r in rows if r["first_approval"]]

    return {
        "chemblVersion": chembl["status"]["chembl_db_version"],
        "chemblReleaseDate": chembl["status"]["chembl_release_date"],
        "rdkitVersion": rdkit.__version__,
        "catalogEntries": {name: catalog(name).GetNumEntries() for name in (*PAINS_CATALOGS, "BRENK")},
        "drugs": len(rows),
        "pains": share(rows),
        "painsByCatalog": {name: share(rows, name.lower()) for name in PAINS_CATALOGS},
        "painsAOrB": share([{**r, "ab": r["pains_a"] or r["pains_b"]} for r in rows], "ab"),
        "alertsPerFlaggedDrug": {str(k): alert_counts_per_drug[k] for k in sorted(alert_counts_per_drug)},
        "distinctAlertsMatched": len(alert_counts),
        "distinctFamiliesMatched": len(families),
        "oral": share([r for r in rows if r["oral"]]),
        "nonOral": share([r for r in rows if not r["oral"]]),
        "withdrawn": share([r for r in rows if r["withdrawn"]]),
        "withFirstApproval": len(dated),
        "datedSince2000": share([r for r in dated if int(r["first_approval"]) >= 2000]),
        "datedBefore2000": share([r for r in dated if int(r["first_approval"]) < 2000]),
        "families": [
            {
                "family": family,
                "label": FAMILY_LABELS.get(family, family),
                "drugs": len(members),
                "pctOfAll": pct(len(members), len(rows)),
                "pctOfFlagged": pct(len(members), len(flagged)),
                "examples": sorted(r["name"] for r in members),
            }
            for family, members in families
        ],
        "alerts": [
            {"alert": alert, "catalog": alert_catalog.get(alert, ""), "drugs": count}
            for alert, count in sorted(alert_counts.items(), key=lambda item: (-item[1], item[0]))
        ],
        "byDecade": by_decade,
        "byTherapeuticClass": by_class,
        "brenk": share(rows, "brenk"),
        "brenkTopAlerts": [{"alert": a, "drugs": c, "pct": pct(c, len(rows))} for a, c in brenk_counts.most_common(15)],
        "painsAndBrenk": sum(1 for r in rows if r["pains"] and r["brenk"]),
        "painsOnly": sum(1 for r in rows if r["pains"] and not r["brenk"]),
        "brenkOnly": sum(1 for r in rows if r["brenk"] and not r["pains"]),
        "neither": sum(1 for r in rows if not (r["pains"] or r["brenk"])),
    }


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--cache-dir", default=str(Path(tempfile.gettempdir()) / "proteiniq-lipinski-ro5"))
    parser.add_argument("--out-dir", default=str(OUT_DIR))
    args = parser.parse_args()
    cache = Path(args.cache_dir)
    cache.mkdir(parents=True, exist_ok=True)
    out = Path(args.out_dir)
    out.mkdir(parents=True, exist_ok=True)

    lipinski = load_lipinski_module()
    chembl = lipinski.load_chembl(cache)
    rows = build_rows(lipinski, chembl)
    rows.sort(key=lambda row: row["name"].lower())
    summary = summarize(rows, chembl, lipinski)
    summary["libraryCheck"] = lipinski.check_against_library(rows)

    with open(out / "approved-drugs-pains.csv", "w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(rows[0]), lineterminator="\n")
        writer.writeheader()
        for row in rows:
            writer.writerow({k: (str(v).lower() if isinstance(v, bool) else v) for k, v in row.items()})
    (out / "summary.json").write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8")
    shutil.copyfile(Path(__file__), out / "analyze.py")

    keys = ("chemblVersion", "rdkitVersion", "drugs", "pains", "painsByCatalog", "brenk", "datedBefore2000", "datedSince2000", "libraryCheck")
    print(json.dumps({k: summary[k] for k in keys}, indent=2))


if __name__ == "__main__":
    main()