import json
from pathlib import Path

import numpy as np

from logic import OPTIMISERS, Objective, Run, parse_function

ALGORITHMS = {
    spec["id"]: spec
    for spec in json.loads((Path(__file__).parent / "algorithms.json").read_text())["algorithms"]
}

# Iterations a run can take before stepping stops
MAX_ITERATIONS = 1000
# Sample points for the plotted curve (one variable) or surface (two, per side)
CURVE_POINTS = 401
GRID_POINTS = 101
# Gradient arrows per side
ARROW_POINTS = 15


class UserError(Exception):
    """A mistake the student can fix; its message is shown in the UI."""


def _json_values(values: np.ndarray) -> list:
    """Values as JSON numbers, with null where f has no finite value."""
    return [float(v) if np.isfinite(v) else None for v in values]


class Manager:
    def __init__(self):
        self._objective: Objective | None = None
        self._runs: list[Run] = []

    def handle_set_function(self, args: dict) -> None:
        """{variables: 1 | 2, expression}; clears the runs."""
        self._objective = None
        self._runs = []
        n = int(args["variables"])
        try:
            expr = parse_function(str(args["expression"]), n)
        except ValueError as e:
            raise UserError(str(e))
        try:
            self._objective = Objective(expr, n)
        except Exception:
            raise UserError("This function can't be differentiated here. Try writing it another way.")

    def handle_function_values(self, args: dict) -> dict | None:
        """f sampled over the plot's {x_range, y_range}; None with no function."""
        f = self._objective
        if f is None:
            return None
        with np.errstate(all="ignore"):
            if f.n == 1:
                x = np.linspace(*args["x_range"], CURVE_POINTS)
                return {"x": x.tolist(), "f": _json_values(f.values_on(x))}

            x = np.linspace(*args["x_range"], GRID_POINTS)
            y = np.linspace(*args["y_range"], GRID_POINTS)
            z = f.values_on(*np.meshgrid(x, y))
            ax, ay = np.meshgrid(np.linspace(*args["x_range"], ARROW_POINTS), np.linspace(*args["y_range"], ARROW_POINTS))
            u, v = f.gradient_on(ax, ay)
            return {
                "x": x.tolist(),
                "y": y.tolist(),
                # Rows follow y, columns follow x
                "f": [_json_values(row) for row in z],
                "arrows": {
                    "x": ax.ravel().tolist(),
                    "y": ay.ravel().tolist(),
                    "u": _json_values(u.ravel()),
                    "v": _json_values(v.ravel()),
                },
            }

    def handle_set_runs(self, args: dict) -> None:
        """{runs: [{algorithm, start, params}]}: the optimisers to step together."""
        self._runs = []
        f = self._objective
        if f is None:
            raise ValueError("Function not set")
        runs = []
        for run in args["runs"]:
            spec = ALGORITHMS[run["algorithm"]]
            given = dict(run.get("params", {}))
            params = {p["name"]: float(given.get(p["name"], p["default"])) for p in spec["params"]}
            start = np.array(run["start"], dtype=float)
            if start.shape != (f.n,):
                raise ValueError(f"Start point needs {f.n} coordinates, got {start.tolist()}")
            with np.errstate(all="ignore"):
                if not np.isfinite(f.value(start)):
                    point = ", ".join(f"{c:g}" for c in start)
                    raise UserError(f"The function has no finite value at the initial point ({point}).")
            runs.append(Run(f, OPTIMISERS[spec["id"]](f, **params), start))
        self._runs = runs

    def handle_state(self, args: dict) -> dict | None:
        """Every run after {iteration} iterations; None once all have stopped.

        "stopped" says why a run ended, at its last point and after it (where
        point is None). It is None while the run carries on.
        """
        k = int(args["iteration"])
        if not self._runs or k > MAX_ITERATIONS:
            return None
        points = [run.point(k) for run in self._runs]
        if all(p is None for p in points):
            return None
        # Look one iteration ahead to know whether this is a run's last point
        ended = [p is None or run.point(k + 1) is None for run, p in zip(self._runs, points)]
        return {
            "runs": [
                {"point": p, "stopped": run.stopped if end else None}
                for run, p, end in zip(self._runs, points, ended)
            ]
        }