"""Clashscore percentiles by resolution from wwPDB validation data.

Queries the RCSB PDB Search API for entries in a resolution window and
counts their wwPDB validation clashscores in narrow bins (range facets on
pdbx_vrpt_summary_geometry.clashscore), then interpolates the 10th, 25th,
50th, 75th and 90th percentiles. Entries without a clashscore are skipped.
Writes clashscore_by_resolution.csv next to this script.

Usage: python clashscore_by_resolution.py
"""

import csv
import json
import urllib.request
from pathlib import Path

URL = "https://search.rcsb.org/rcsbsearch/v2/query"
EDGES = [round(i * 0.25, 2) for i in range(0, 160)] + [40.0 + 2.5 * i for i in range(0, 25)]  # 0.25 steps to 40, 2.5 steps to 100, then 100+
PERCENTILES = [10, 25, 50, 75, 90]

WINDOWS = [
    # (method label, exptl.method value, resolution from, to, release-date filter);
    # NMR entries have no resolution, so that window takes every solution NMR entry
    ("X-ray", "X-RAY DIFFRACTION", 0.9, 1.1, None),
    ("X-ray", "X-RAY DIFFRACTION", 1.4, 1.6, None),
    ("X-ray", "X-RAY DIFFRACTION", 1.9, 2.1, None),
    ("X-ray", "X-RAY DIFFRACTION", 2.4, 2.6, None),
    ("X-ray", "X-RAY DIFFRACTION", 2.9, 3.1, None),
    ("X-ray", "X-RAY DIFFRACTION", 3.4, 3.6, None),
    ("Cryo-EM", "ELECTRON MICROSCOPY", 2.4, 2.6, None),
    ("Cryo-EM", "ELECTRON MICROSCOPY", 2.9, 3.1, None),
    ("Cryo-EM", "ELECTRON MICROSCOPY", 3.4, 3.6, None),
    ("Cryo-EM", "ELECTRON MICROSCOPY", 3.9, 4.1, None),
    ("NMR", "SOLUTION NMR", None, None, None),
    ("X-ray", "X-RAY DIFFRACTION", 1.9, 2.1, ("1900-01-01", "2001-12-31")),
    ("X-ray", "X-RAY DIFFRACTION", 1.9, 2.1, ("2020-01-01", "2026-12-31")),
]


def terminal(attribute, operator, value):
    return {"type": "terminal", "service": "text",
            "parameters": {"attribute": attribute, "operator": operator, "value": value}}


def histogram(method, lo, hi, released):
    nodes = [terminal("exptl.method", "exact_match", method)]
    if lo is not None:
        nodes.append(terminal("rcsb_entry_info.resolution_combined", "range",
                              {"from": lo, "to": hi, "include_lower": True, "include_upper": False}))
    if released:
        nodes.append(terminal("rcsb_accession_info.initial_release_date", "range",
                              {"from": released[0], "to": released[1],
                               "include_lower": True, "include_upper": True}))
    ranges = [{"from": a, "to": b} for a, b in zip(EDGES, EDGES[1:])] + [{"from": EDGES[-1]}]
    body = {"query": {"type": "group", "logical_operator": "and", "nodes": nodes},
            "return_type": "entry",
            "request_options": {"paginate": {"start": 0, "rows": 0},
                                "facets": [{"name": "cs", "aggregation_type": "range",
                                            "attribute": "pdbx_vrpt_summary_geometry.clashscore",
                                            "ranges": ranges}]}}
    req = urllib.request.Request(URL, data=json.dumps(body).encode(),
                                 headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req) as resp:
        data = json.load(resp)
    counts = {b["label"]: b["population"] for b in data["facets"][0]["buckets"]}
    return [bucket_count(counts, a, b) for a, b in zip(EDGES, EDGES[1:])] + [bucket_count(counts, EDGES[-1], None)]


def bucket_count(counts, a, b):
    for label, n in counts.items():
        lo, _, hi = label.partition("-")
        if abs(float(lo) - a) < 1e-9 and ((b is None and hi == "*") or (b is not None and hi != "*" and abs(float(hi) - b) < 1e-9)):
            return n
    return 0


def percentile(counts, p):
    total = sum(counts)
    target = total * p / 100
    running = 0
    for i, n in enumerate(counts):
        if running + n >= target and n > 0:
            lo = EDGES[i]
            hi = EDGES[i + 1] if i + 1 < len(EDGES) else None
            if hi is None:
                return float("nan")
            return lo + (hi - lo) * (target - running) / n
        running += n
    return float("nan")


def main():
    out = Path(__file__).with_name("clashscore_by_resolution.csv")
    with out.open("w", newline="", encoding="utf-8") as fh:
        writer = csv.writer(fh)
        writer.writerow(["method", "resolution_from", "resolution_to", "released", "entries_with_clashscore"]
                        + [f"p{p}" for p in PERCENTILES] + ["share_above_20"])
        for label, method, lo, hi, released in WINDOWS:
            counts = histogram(method, lo, hi, released)
            total = sum(counts)
            above_20 = sum(n for e, n in zip(EDGES, counts) if e >= 20) / total
            row = [label, lo, hi, f"{released[0]}..{released[1]}" if released else "all", total]
            row += [f"{percentile(counts, p):.1f}" for p in PERCENTILES] + [f"{above_20:.3f}"]
            writer.writerow(row)
            print(row)


if __name__ == "__main__":
    main()