from logic import forward_mode, parse_function, reverse_mode, state_after


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


class Manager:
    def __init__(self):
        self._graph: dict | None = None
        self._run: dict | None = None

    def handle_set_function(self, args: dict) -> dict:
        """{expression}; returns the function's computational graph.

        A function that doesn't parse leaves the current one and its run in place.
        """
        try:
            self._graph = parse_function(str(args["expression"]))
        except ValueError as e:
            raise UserError(str(e))
        self._run = None
        return self._graph

    def handle_set_run(self, args: dict) -> dict:
        """{mode: "forward" | "reverse", inputs: {name: value}, seed}; seed is the input forward mode differentiates by.

        Returns {steps}, the run's number of steps.
        """
        self._run = None
        inputs = {name: float(value) for name, value in args["inputs"].items()}
        try:
            if args["mode"] == "forward":
                self._run = forward_mode(self._graph, inputs, args["seed"])
            else:
                self._run = reverse_mode(self._graph, inputs)
        except ValueError as e:
            raise UserError(str(e))
        return {"steps": len(self._run["steps"])}

    def handle_state(self, args: dict) -> dict | None:
        """The run after {iteration} steps; None past the end."""
        k = int(args["iteration"])
        if self._run is None or k > len(self._run["steps"]):
            return None
        return state_after(self._run, k)