"""Write src/sklearn_params.json: every setting of each model in models.json,
read from scikit-learn's parameter constraints and docstrings.

Run it with the scikit-learn version Pyodide ships (see the PYODIDE_SKLEARN
note in the README), from the backend folder:

    uv run --isolated --no-project --python 3.12 --with "scikit-learn==1.4.2" \\
        --with "numpy<2" python scripts/generate_sklearn_params.py
"""

import json
import re
import sys
from numbers import Integral
from pathlib import Path

import sklearn
from sklearn.utils._param_validation import Hidden, Interval, StrOptions

SRC = Path(__file__).resolve().parent.parent / "src"
sys.path.insert(0, str(SRC))
from logic import ESTIMATORS, MODEL_SPECS  # noqa: E402

# Settings that don't change what is learned, or are gone in newer scikit-learn
EXCLUDED = {"n_jobs", "verbose", "warm_start", "multi_class"}
# "precomputed" needs a matrix students can't give from the page
EXCLUDED_OPTIONS = {"precomputed"}


def describe(constraints: list) -> dict | None:
    """The control for a setting, or None if students can't set it from the page
    (e.g. it only takes callables, dicts or arrays)."""
    strings, none, boolean, intervals = [], False, False, []
    for c in constraints:
        if c is None:
            none = True
        elif c == "boolean":
            boolean = True
        elif c == "random_state":
            # A seed; scikit-learn also takes a RandomState object
            intervals.append(Interval(Integral, 0, 2**32 - 1, closed="both"))
            none = True
        elif isinstance(c, StrOptions):
            strings += sorted(c.options - c.deprecated - EXCLUDED_OPTIONS)
        elif isinstance(c, Interval):
            intervals.append(c)
        elif isinstance(c, Hidden):
            continue
        # Anything else (callable, dict, array-like, ...) has no control

    if boolean and not (strings or none or intervals):
        return {"kind": "boolean"}

    number = None
    if intervals:
        integer = all(i.type is Integral for i in intervals)
        number = {"integer": integer}
        lows = [i for i in intervals if i.left is not None]
        highs = [i for i in intervals if i.right is not None]
        # Unbounded if any interval is; open if the bounding interval is
        if len(lows) == len(intervals):
            low = min(lows, key=lambda i: i.left)
            number["min"] = low.left
            if low.closed in ("right", "neither"):
                number["minExclusive"] = True
        if len(highs) == len(intervals):
            high = max(highs, key=lambda i: i.right)
            number["max"] = high.right
            if high.closed in ("left", "neither"):
                number["maxExclusive"] = True

    if number and not strings:
        spec = {"kind": "integer" if number.pop("integer") else "number", **number}
        if none:
            spec["nullable"] = "None"
        return spec
    if strings and not number and not none:
        return {"kind": "choice", "options": [{"value": s, "label": s} for s in strings]}
    options = [{"value": s, "label": s} for s in strings] + ([{"value": None, "label": "None"}] if none else [])
    if len(options) + (number is not None) < 2:
        return None
    spec = {"kind": "mixed", "options": options}
    if number:
        spec["number"] = number
    return spec


def descriptions(estimator) -> dict[str, str]:
    """First sentence of each setting's description in the docstring."""
    doc = estimator.__doc__ or ""
    section = doc.split("Parameters\n    ----------\n", 1)[-1]
    section = re.split(r"\n    \w[\w ]*\n    -{3,}\n", section, maxsplit=1)[0]
    result = {}
    for match in re.finditer(r"^    (\w+) : .*\n((?:(?: {8}.*)?\n)+)", section, re.M):
        paragraph = " ".join(match.group(2).strip().split("\n\n")[0].split())
        sentence = re.split(r"(?<=\.)\s(?=[A-Z])", paragraph, maxsplit=1)[0]
        result[match.group(1)] = clean_rst(sentence)
    return result


def clean_rst(text: str) -> str:
    text = re.sub(r":\w+:`~?([^`<]*?)\s*(?:<[^>]*>)?`", lambda m: m.group(1).split(".")[-1], text)
    return text.replace("``", "").replace("`", "")


def generate() -> dict:
    models = {}
    for spec in MODEL_SPECS.values():
        estimator = ESTIMATORS[spec["estimator"]]
        defaults = estimator().get_params()
        texts = descriptions(estimator)
        params = []
        for name, constraints in estimator._parameter_constraints.items():
            if name in EXCLUDED:
                continue
            control = describe(constraints)
            if control is None:
                continue
            params.append({
                "name": name,
                "label": name if len(name) <= 2 else name.replace("_", " ").capitalize(),
                "description": texts.get(name, ""),
                **control,
                "default": defaults[name],
            })
        models[spec["id"]] = params
    return {"sklearn_version": sklearn.__version__, "models": models}


if __name__ == "__main__":
    out = SRC / "sklearn_params.json"
    out.write_text(json.dumps(generate(), indent=2) + "\n")
    print(f"Wrote {out} (scikit-learn {sklearn.__version__})")