#!/usr/bin/env python3
"""Count how many approved small-molecule drugs break Lipinski's rule of five.

Data for the ProteinIQ guide /guides/lipinski-rule-of-five.

Population
  ChEMBL parent compounds with max_phase 4 (approved in at least one
  jurisdiction ChEMBL tracks), selected with the same rules as ProteinIQ's
  approved small-molecule drug library
  (scripts/curated-datasets/generate-approved-drug-library.py): salts and
  hydrates collapse into their parent; small molecules are kept, plus
  structurally defined peptides and oligosaccharides of at most 900 Da that
  ChEMBL types as protein, oligosaccharide or unknown; structures without
  carbon, with metals or other elements outside H, B, C, N, O, F, Si, P, S,
  Cl, Se, Br and I, or with several disconnected components are excluded.
  When the generated library file is present, the script checks that both
  select the same ChEMBL IDs.

Properties (all recalculated with RDKit from the ChEMBL parent structure)
  * MW: average molecular weight of the neutral parent (Descriptors.MolWt).
  * cLogP: Wildman-Crippen logP (Crippen.MolLogP), the ALogP ChEMBL reports.
    Lipinski's paper used CLogP > 5 or MLogP > 4.15; neither is in RDKit.
  * Donors and acceptors, two definitions:
      "lipinski": Lipinski's own counts, OH plus NH (Lipinski.NHOHCount)
                  and N plus O atoms (Lipinski.NOCount). Headline definition.
      "rdkit":    RDKit pharmacophore donors and acceptors
                  (CalcNumHBD / CalcNumHBA), the counts ChEMBL uses.
  * Violations: number of MW > 500, cLogP > 5, donors > 5, acceptors > 10.
    Lipinski's alert ("poor absorption or permeation is more likely") fires
    at two or more violations.
  * Oral: ChEMBL's oral route flag on the parent or any of its salts.
  * First approval: earliest ChEMBL first_approval year across the family.
  * Kinase inhibitor: any ATC code in L01E (protein kinase inhibitors).
  * Macrocycle: any ring of 12 or more atoms.

Outputs (public/data/guides/lipinski-rule-of-five/)
  approved-drugs-ro5.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/lipinski-ro5-chembl.py [--cache-dir DIR]

ChEMBL responses are cached in --cache-dir. Delete the cache to pick up a new
ChEMBL release.
"""

from __future__ import annotations

import argparse
import collections
import csv
import json
import re
import shutil
import statistics
import tempfile
import time
import urllib.error
import urllib.parse
import urllib.request
from pathlib import Path

import rdkit
from rdkit import Chem, RDLogger
from rdkit.Chem import Crippen, Descriptors, Lipinski, rdMolDescriptors

RDLogger.DisableLog("rdApp.*")

REPO_ROOT = Path(__file__).resolve().parents[1]
OUT_DIR = REPO_ROOT / "public/data/guides/lipinski-rule-of-five"
LIBRARY_TS = REPO_ROOT / "src/lib/curated-datasets/approved-drug-library.generated.ts"

CHEMBL_API = "https://www.ebi.ac.uk/chembl/api/data"
USER_AGENT = "ProteinIQ Lipinski analysis (https://proteiniq.io/guides/lipinski-rule-of-five)"

ALLOWED_ELEMENTS = {"H", "B", "C", "N", "O", "F", "Si", "P", "S", "Cl", "Se", "Br", "I"}
NON_SMALL_MOLECULE_TYPES_KEPT = {"Protein", "Oligosaccharide", "Unknown"}
NON_SMALL_MOLECULE_MAX_MW = 900.0
MACROCYCLE_MIN_RING = 12
MIN_CLASS_SIZE = 30


# --------------------------------------------------------------------------
# ChEMBL access with an on-disk cache


def http_json(url: str, retries: int = 5) -> dict:
    for attempt in range(retries):
        try:
            request = urllib.request.Request(url, headers={"User-Agent": USER_AGENT})
            with urllib.request.urlopen(request, timeout=120) as response:
                return json.loads(response.read())
        except urllib.error.HTTPError as error:
            if error.code == 404 or attempt == retries - 1:
                raise
        except Exception:
            if attempt == retries - 1:
                raise
        time.sleep(2 * (attempt + 1))
    raise RuntimeError("unreachable")


def paginate(resource: str, params: dict[str, str]) -> list[dict]:
    query = urllib.parse.urlencode({**params, "limit": "1000", "format": "json"})
    url: str | None = f"{CHEMBL_API}/{resource}.json?{query}"
    items: list[dict] = []
    key = None
    while url:
        payload = http_json(url)
        key = key or next(k for k in payload if k != "page_meta")
        items.extend(payload[key])
        next_url = payload["page_meta"].get("next")
        url = f"https://www.ebi.ac.uk{next_url}" if next_url else None
    return items


def cached(cache: Path, key: str, fetch):
    path = cache / f"{key}.json"
    if path.exists():
        return json.loads(path.read_text(encoding="utf-8"))
    value = fetch()
    path.write_text(json.dumps(value), encoding="utf-8")
    return value


MOLECULE_FIELDS = ",".join(
    [
        "molecule_chembl_id", "pref_name", "molecule_type", "molecule_hierarchy",
        "molecule_structures", "molecule_properties", "inorganic_flag", "first_approval",
        "oral", "parenteral", "topical", "atc_classifications", "withdrawn_flag", "max_phase",
    ]
)


def load_chembl(cache: Path) -> dict:
    status = cached(cache, "status", lambda: http_json(f"{CHEMBL_API}/status.json"))
    approved = cached(cache, "max_phase_4", lambda: paginate("molecule", {"max_phase": "4", "only": MOLECULE_FIELDS}))
    by_id = {m["molecule_chembl_id"]: m for m in approved}
    parent_ids = sorted({parent_of(m) for m in approved})
    missing = [pid for pid in parent_ids if pid not in by_id]

    def fetch_missing() -> list[dict]:
        found: list[dict] = []
        for start in range(0, len(missing), 40):
            batch = ",".join(missing[start : start + 40])
            found += paginate("molecule", {"molecule_chembl_id__in": batch, "only": MOLECULE_FIELDS})
        return found

    for molecule in cached(cache, "missing_parents", fetch_missing):
        by_id[molecule["molecule_chembl_id"]] = molecule
    atc = cached(cache, "atc_classes", lambda: paginate("atc_class", {}))
    return {"status": status, "approved": approved, "by_id": by_id, "parent_ids": parent_ids, "atc": atc}


def parent_of(molecule: dict) -> str:
    return (molecule.get("molecule_hierarchy") or {}).get("parent_chembl_id") or molecule["molecule_chembl_id"]


# --------------------------------------------------------------------------
# Drug table


def largest_ring(mol: Chem.Mol) -> int:
    rings = mol.GetRingInfo().AtomRings()
    return max((len(ring) for ring in rings), default=0)


def build_rows(chembl: dict) -> tuple[list[dict], collections.Counter]:
    by_id = chembl["by_id"]
    families: dict[str, list[dict]] = collections.defaultdict(list)
    for molecule in by_id.values():
        families[parent_of(molecule)].append(molecule)
    atc_level2 = {e["level5"]: (e["level2"], e["level2_description"].strip().capitalize()) for e in chembl["atc"]}

    excluded: collections.Counter = collections.Counter()
    rows: list[dict] = []
    for parent_id in chembl["parent_ids"]:
        parent = by_id.get(parent_id)
        if parent is None:
            excluded["parent record missing"] += 1
            continue
        molecule_type = parent.get("molecule_type")
        smiles = (parent.get("molecule_structures") or {}).get("canonical_smiles") or ""
        if molecule_type != "Small molecule" and molecule_type not in NON_SMALL_MOLECULE_TYPES_KEPT:
            excluded["biologic"] += 1
            continue
        if not smiles:
            excluded["no structure" if molecule_type == "Small molecule" else "biologic"] += 1
            continue
        mol = Chem.MolFromSmiles(smiles)
        if mol is None:
            excluded["structure does not parse"] += 1
            continue
        if molecule_type != "Small molecule" and Descriptors.MolWt(mol) > NON_SMALL_MOLECULE_MAX_MW:
            excluded["biologic"] += 1
            continue
        elements = {atom.GetSymbol() for atom in mol.GetAtoms()}
        if "C" not in elements or parent.get("inorganic_flag") == 1:
            excluded["inorganic"] += 1
            continue
        if not elements <= ALLOWED_ELEMENTS:
            excluded["metals or other non-organic elements"] += 1
            continue
        if len(Chem.GetMolFrags(mol)) > 1:
            excluded["several disconnected components"] += 1
            continue
        family = families[parent_id]
        name = parent.get("pref_name") or next((m.get("pref_name") for m in family if m.get("pref_name")), None)
        if not name:
            excluded["no preferred name"] += 1
            continue

        mw = Descriptors.MolWt(mol)
        clogp = Crippen.MolLogP(mol)
        hbd_lipinski = Lipinski.NHOHCount(mol)
        hba_lipinski = Lipinski.NOCount(mol)
        hbd_rdkit = rdMolDescriptors.CalcNumHBD(mol)
        hba_rdkit = rdMolDescriptors.CalcNumHBA(mol)
        violations = sum([mw > 500, clogp > 5, hbd_lipinski > 5, hba_lipinski > 10])
        violations_rdkit = sum([mw > 500, clogp > 5, hbd_rdkit > 5, hba_rdkit > 10])
        tpsa = rdMolDescriptors.CalcTPSA(mol)
        rotatable = rdMolDescriptors.CalcNumRotatableBonds(mol)

        atc_codes = sorted({code for m in family for code in (m.get("atc_classifications") or [])})
        level2 = collections.Counter(atc_level2[c] for c in atc_codes if c in atc_level2)
        therapeutic = sorted(level2.items(), key=lambda item: (-item[1], item[0][0]))[0][0] if level2 else ("", "")
        first_years = [m["first_approval"] for m in family if m.get("first_approval")]
        chembl_violations = (parent.get("molecule_properties") or {}).get("num_ro5_violations")

        rows.append(
            {
                "chembl_id": parent_id,
                "name": name.capitalize() if name.isupper() else name,
                "molecule_type": molecule_type,
                "first_approval": min(first_years) if first_years else "",
                "oral": any(m.get("oral") for m in family),
                "parenteral": any(m.get("parenteral") for m in family),
                "topical": any(m.get("topical") for m in family),
                "withdrawn": any(m.get("withdrawn_flag") for m in family),
                "atc_level2": therapeutic[0],
                "therapeutic_class": therapeutic[1],
                "kinase_inhibitor": any(code.startswith("L01E") for code in atc_codes),
                "largest_ring": largest_ring(mol),
                "macrocycle": largest_ring(mol) >= MACROCYCLE_MIN_RING,
                "mw": round(mw, 2),
                "clogp": round(clogp, 2),
                "hbd": hbd_lipinski,
                "hba": hba_lipinski,
                "ro5_violations": violations,
                "ro5_alert": violations >= 2,
                "hbd_rdkit": hbd_rdkit,
                "hba_rdkit": hba_rdkit,
                "ro5_violations_rdkit": violations_rdkit,
                "tpsa": round(tpsa, 1),
                "rotatable_bonds": rotatable,
                "veber_pass": tpsa <= 140 and rotatable <= 10,
                "chembl_num_ro5_violations": "" if chembl_violations is None else chembl_violations,
            }
        )
    return rows, excluded


# --------------------------------------------------------------------------
# Aggregates


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


def violation_summary(rows: list[dict], key: str = "ro5_violations") -> dict:
    n = len(rows)
    counts = collections.Counter(min(row[key], 2) for row in rows)
    return {
        "n": n,
        "zero": counts[0],
        "one": counts[1],
        "twoOrMore": counts[2],
        "zeroPct": pct(counts[0], n),
        "onePct": pct(counts[1], n),
        "anyPct": pct(counts[1] + counts[2], n),
        "twoOrMorePct": pct(counts[2], n),
    }


def median(values: list[float]) -> float:
    return round(statistics.median(values), 2) if values else 0.0


def decade_label(year: int) -> str:
    return "Before 1950" if year < 1950 else f"{year // 10 * 10}s"


def summarize(rows: list[dict], chembl: dict, excluded: collections.Counter) -> dict:
    oral = [r for r in rows if r["oral"]]
    non_oral = [r for r in rows if not r["oral"]]
    # ChEMBL sets no route flag for many drugs approved only outside the US.
    other_route = [r for r in non_oral if r["parenteral"] or r["topical"]]
    route_unknown = [r for r in non_oral if not (r["parenteral"] or r["topical"])]
    criteria = {
        "mwOver500": sum(r["mw"] > 500 for r in rows),
        "clogpOver5": sum(r["clogp"] > 5 for r in rows),
        "donorsOver5": sum(r["hbd"] > 5 for r in rows),
        "acceptorsOver10": sum(r["hba"] > 10 for r in rows),
    }
    exact = collections.Counter(r["ro5_violations"] for r in rows)

    decades = collections.defaultdict(list)
    for row in rows:
        if row["first_approval"]:
            decades[decade_label(int(row["first_approval"]))].append(row)
    by_decade = []
    for label in sorted(decades, key=lambda d: (d != "Before 1950", d)):
        group = decades[label]
        group_oral = [r for r in group if r["oral"]]
        by_decade.append(
            {
                "decade": label,
                **violation_summary(group),
                "medianMw": median([r["mw"] for r in group]),
                "medianClogp": median([r["clogp"] for r in group]),
                "meanMw": round(statistics.mean(r["mw"] for r in group), 1),
                "mwOver500Pct": pct(sum(r["mw"] > 500 for r in group), len(group)),
                "oral": violation_summary(group_oral),
                "oralMedianMw": median([r["mw"] for r in group_oral]),
                "oralMedianClogp": median([r["clogp"] for r in group_oral]),
            }
        )

    groups = {
        "Kinase inhibitors (ATC L01E)": [r for r in rows if r["kinase_inhibitor"]],
        "Macrocycles (ring of 12+ atoms)": [r for r in rows if r["macrocycle"]],
        "Peptides and oligosaccharides typed as non-small molecules": [
            r for r in rows if r["molecule_type"] != "Small molecule"
        ],
    }
    by_group = [{"group": name, **violation_summary(members), "oral": violation_summary([r for r in members if r["oral"]])} for name, members in groups.items()]

    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, **violation_summary(members)}
            for (code, label), members in by_class_rows.items()
            if len(members) >= MIN_CLASS_SIZE
        ),
        key=lambda item: -item["twoOrMorePct"],
    )

    recent_oral_breakers = sorted(
        (r for r in oral if r["ro5_alert"] and r["first_approval"] and int(r["first_approval"]) >= 2010),
        key=lambda r: (-int(r["first_approval"]), r["name"]),
    )

    def oral_between(first: int, last: int) -> list[dict]:
        return [r for r in oral if r["first_approval"] and first <= int(r["first_approval"]) <= last]

    # Windows matching published analyses, so the numbers can be compared.
    comparisons = {
        "oral1997AndEarlier": violation_summary(oral_between(0, 1997)),
        "oralAfter1997": violation_summary(oral_between(1998, 9999)),
        "oral2000To2022": violation_summary(oral_between(2000, 2022)),
        "oral2000To2022Rdkit": violation_summary(oral_between(2000, 2022), "ro5_violations_rdkit"),
        "oral2011To2022": violation_summary(oral_between(2011, 2022)),
        "oral2011To2022Rdkit": violation_summary(oral_between(2011, 2022), "ro5_violations_rdkit"),
    }

    chembl_null = sum(1 for r in rows if r["chembl_num_ro5_violations"] == "")
    chembl_known = [r for r in rows if r["chembl_num_ro5_violations"] != ""]
    return {
        "chemblVersion": chembl["status"]["chembl_db_version"],
        "chemblReleaseDate": chembl["status"]["chembl_release_date"],
        "rdkitVersion": rdkit.__version__,
        "approvedChemblRecords": len(chembl["approved"]),
        "parentCompounds": len(chembl["parent_ids"]),
        "excluded": dict(excluded),
        "drugs": len(rows),
        "withFirstApproval": sum(1 for r in rows if r["first_approval"]),
        "all": violation_summary(rows),
        "exactViolationCounts": {str(k): exact[k] for k in sorted(exact)},
        "allRdkitDonorsAcceptors": violation_summary(rows, "ro5_violations_rdkit"),
        "criteria": criteria,
        "criteriaPct": {key: pct(value, len(rows)) for key, value in criteria.items()},
        "oral": violation_summary(oral),
        "nonOral": violation_summary(non_oral),
        "parenteralOrTopicalOnly": violation_summary(other_route),
        "routeUnknown": violation_summary(route_unknown),
        "oralMedianMw": median([r["mw"] for r in oral]),
        "nonOralMedianMw": median([r["mw"] for r in non_oral]),
        "oralVeberPassPct": pct(sum(r["veber_pass"] for r in oral), len(oral)),
        "oralAlertVeberPassPct": pct(sum(r["veber_pass"] for r in oral if r["ro5_alert"]), sum(r["ro5_alert"] for r in oral)),
        "medianMw": median([r["mw"] for r in rows]),
        "medianClogp": median([r["clogp"] for r in rows]),
        "publishedWindowComparisons": comparisons,
        "byDecade": by_decade,
        "byGroup": by_group,
        "byTherapeuticClass": by_class,
        "recentOralRo5Alert": [
            {"name": r["name"], "chemblId": r["chembl_id"], "year": r["first_approval"], "mw": r["mw"], "violations": r["ro5_violations"]}
            for r in recent_oral_breakers
        ],
        "chemblFieldComparison": {
            "drugsWithoutChemblValue": chembl_null,
            "drugsWithoutChemblValueAndMwOver500": sum(1 for r in rows if r["chembl_num_ro5_violations"] == "" and r["mw"] > 500),
            "chemblFieldTwoOrMorePctOfKnown": pct(sum(1 for r in chembl_known if int(r["chembl_num_ro5_violations"]) >= 2), len(chembl_known)),
            "chemblFieldTwoOrMorePctOfAll": pct(sum(1 for r in chembl_known if int(r["chembl_num_ro5_violations"]) >= 2), len(rows)),
        },
    }


def check_against_library(rows: list[dict]) -> str:
    """Confirm the population matches ProteinIQ's approved drug library."""
    if not LIBRARY_TS.exists():
        return "library file not found; skipped"
    library_ids = set(re.findall(r'^\s*\["[^"]*","[^"]*","(CHEMBL\d+)"', LIBRARY_TS.read_text(encoding="utf-8"), re.M))
    ours = {row["chembl_id"] for row in rows}
    if library_ids != ours:
        return f"differs: {len(ours - library_ids)} only here, {len(library_ids - ours)} only in library"
    return f"same {len(ours)} ChEMBL IDs"


# --------------------------------------------------------------------------


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)

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

    with open(out / "approved-drugs-ro5.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")

    print(json.dumps({k: summary[k] for k in ("chemblVersion", "drugs", "all", "oral", "nonOral", "criteriaPct", "libraryCheck")}, indent=2))


if __name__ == "__main__":
    main()