"""Parse a typed function into a computational graph, and run autodiff on it.

A graph is {"inputs": [...], "constants": {...}, "nodes": [...], "output": id}.
Nodes come in evaluation order, numbered w1, w2, ... with the inputs first. A
node's id is what labels the edges leaving it:
  {"id": "w1", "kind": "input", "name": "x"}
  {"id": "w2", "kind": "op", "op": "mul", "args": ["2", "w1"]}
An argument is a node's id or a key of constants, e.g. {"2": 2.0}. As in PyTorch
and JAX, constants are part of an operation rather than nodes of the graph, and
an expression of constants alone is worked out first: 2*3*x has one operation.
Identical subexpressions share one node. Python's ast is used rather than
sympy because sympy rewrites what was typed (x*x becomes x**2).

A run is a list of steps, one per iteration of the stepper:
  {"phase": "forward" | "backward", "node": id, "value": v, "derivative": d,
   "partials": [{"arg": id, "partial": p, "amount": a}]}
partial is the node's local derivative with respect to arg. In forward mode,
derivative is the node's tangent, the sum of the amounts (partial * arg's tangent).
In reverse mode's backward phase, derivative is the node's adjoint, and each
amount (adjoint * partial) is added to arg's adjoint. Steps in its forward phase
have no "derivative" key: Pyodide would turn a nested None into undefined, not null.
"""

import ast
import math
import re

FUNCTIONS = ("sin", "cos", "tan", "exp", "log", "sqrt")
CONSTANTS = {"pi": math.pi, "e": math.e}
BINARY_OPS = {ast.Add: "add", ast.Sub: "sub", ast.Mult: "mul", ast.Div: "div", ast.Pow: "pow"}

VALUES = {
    "add": lambda a, b: a + b,
    "sub": lambda a, b: a - b,
    "mul": lambda a, b: a * b,
    "div": lambda a, b: a / b,
    "pow": math.pow,
    "neg": lambda a: -a,
    "sin": math.sin,
    "cos": math.cos,
    "tan": math.tan,
    "exp": math.exp,
    "log": math.log,
    "sqrt": math.sqrt,
}


def _pow_by_base(a: float, b: float) -> float:
    return 0.0 if b == 0 else b * math.pow(a, b - 1)


def _pow_by_exponent(a: float, b: float) -> float:
    # The limit at a = 0, as PyTorch uses
    return 0.0 if a == 0 and b >= 0 else math.pow(a, b) * math.log(a)


# Each operation's derivative with respect to each argument
PARTIALS = {
    "add": (lambda a, b: 1.0, lambda a, b: 1.0),
    "sub": (lambda a, b: 1.0, lambda a, b: -1.0),
    "mul": (lambda a, b: b, lambda a, b: a),
    "div": (lambda a, b: 1 / b, lambda a, b: -a / b**2),
    "pow": (_pow_by_base, _pow_by_exponent),
    "neg": (lambda a: -1.0,),
    "sin": (math.cos,),
    "cos": (lambda a: -math.sin(a),),
    "tan": (lambda a: 1 / math.cos(a) ** 2,),
    "exp": (math.exp,),
    "log": (lambda a: 1 / a,),
    "sqrt": (lambda a: 1 / (2 * math.sqrt(a)),),
}

_INFIX = {"add": "{:g} + {:g}", "sub": "{:g} - {:g}", "mul": "{:g} * {:g}", "div": "{:g} / {:g}", "pow": "{:g} ^ {:g}", "neg": "-{:g}"}


def _formula(op: str, args: list[float]) -> str:
    return _INFIX.get(op, op + "({:g})").format(*args)


def _finite(fn, args: list[float], description: str) -> float:
    """fn(*args), or a student-facing ValueError if it has no finite value."""
    try:
        value = float(fn(*args))
    except (ValueError, ZeroDivisionError, OverflowError):
        value = math.nan
    if not math.isfinite(value):
        raise ValueError(f"{description} has no finite value. Try other values.")
    return value


def _value(op: str, args: list[float]) -> float:
    return _finite(VALUES[op], args, _formula(op, args))


_UNSUPPORTED = f"Use only numbers, variables, + - * / ^ and the functions {', '.join(FUNCTIONS)}."


class _GraphBuilder:
    def __init__(self):
        self.inputs = set()
        self.consts = {}
        self.ops = []
        # (op, args) -> a temporary id, so repeated subexpressions share a node
        self.op_ids = {}

    def number(self, value: object) -> str:
        if type(value) not in (int, float):
            raise ValueError(_UNSUPPORTED)
        try:
            number = float(value)
        except OverflowError:
            number = math.inf
        if not math.isfinite(number):
            raise ValueError(f"The number {value} is too large.")
        return self.const(repr(number).removesuffix(".0"), number)

    def const(self, node_id: str, value: float) -> str:
        self.consts[node_id] = value
        return node_id

    def op(self, node: ast.expr, op: str, *args: str) -> str:
        if all(arg in self.consts for arg in args):
            label = ast.unparse(node).replace("**", "^")
            return self.const(label, _value(op, [self.consts[arg] for arg in args]))
        key = (op, args)
        if key not in self.op_ids:
            self.op_ids[key] = f"#{len(self.ops)}"
            self.ops.append({"id": self.op_ids[key], "kind": "op", "op": op, "args": list(args)})
        return self.op_ids[key]

    def name(self, name: str) -> str:
        if name in CONSTANTS:
            return self.const(name, CONSTANTS[name])
        if name in FUNCTIONS:
            raise ValueError(f"{name} is a function: write it as {name}(x).")
        if re.fullmatch(r"w_?\d+", name):
            raise ValueError(f"{name} is the graph's name for an intermediate value: choose another variable name.")
        self.inputs.add(name)
        return name

    def visit(self, node: ast.expr) -> str:
        match node:
            case ast.Constant(value=value):
                return self.number(value)
            case ast.Name(id=name):
                return self.name(name)
            case ast.UnaryOp(op=ast.USub(), operand=ast.Constant(value=value)):
                return self.number(-value if type(value) in (int, float) else value)
            case ast.UnaryOp(op=ast.USub(), operand=operand):
                return self.op(node, "neg", self.visit(operand))
            case ast.UnaryOp(op=ast.UAdd(), operand=operand):
                return self.visit(operand)
            case ast.BinOp(left=left, op=op, right=right) if type(op) in BINARY_OPS:
                left_id = self.visit(left)
                return self.op(node, BINARY_OPS[type(op)], left_id, self.visit(right))
            case ast.Call(func=ast.Name(id=name), args=args, keywords=[]):
                if name not in FUNCTIONS:
                    raise ValueError(f"{name} isn't a function here. Use one of: {', '.join(FUNCTIONS)}.")
                if len(args) != 1 or isinstance(args[0], ast.Starred):
                    raise ValueError(f"{name} takes one argument.")
                return self.op(node, name, self.visit(args[0]))
        raise ValueError(_UNSUPPORTED)


def parse_function(text: str) -> dict:
    """The graph of a typed function. Raises ValueError with a student-facing message."""
    text = text.strip()
    if not text:
        raise ValueError("Enter a function, like x*y + sin(x).")
    builder = _GraphBuilder()
    try:
        # Python's own ^ binds looser than +, so it would read x^2 + 1 as x^(2 + 1)
        output = builder.visit(ast.parse(text.replace("^", "**"), mode="eval").body)
    except SyntaxError:
        raise ValueError("This function can't be read. Check the brackets, and write multiplication with *, like 2*x.")
    except (RecursionError, MemoryError):
        raise ValueError("This function is too long.")
    if not builder.inputs:
        raise ValueError("The function must use at least one variable, like x.")

    inputs = sorted(builder.inputs)
    ids = {name: f"w{i + 1}" for i, name in enumerate(inputs)}
    ids.update({op["id"]: f"w{len(inputs) + i + 1}" for i, op in enumerate(builder.ops)})
    ops = [{**op, "id": ids[op["id"]], "args": [ids.get(arg, arg) for arg in op["args"]]} for op in builder.ops]
    used = {arg for op in ops for arg in op["args"]}
    return {
        "inputs": inputs,
        "constants": {node_id: value for node_id, value in builder.consts.items() if node_id in used},
        "nodes": [*({"id": ids[name], "kind": "input", "name": name} for name in inputs), *ops],
        "output": ids[output],
    }


def _partials(node: dict, values: dict, constants: dict) -> list[tuple[str, float]]:
    """(arg, local derivative) for each argument that is a node; constants have no derivative."""
    args = [values[arg] for arg in node["args"]]
    description = f"The derivative of {_formula(node['op'], args)}"
    return [
        (arg, _finite(PARTIALS[node["op"]][i], args, description))
        for i, arg in enumerate(node["args"])
        if arg not in constants
    ]


def _step(phase: str, node_id: str, value: float, derivative: float | None, partials: list) -> dict:
    step = {"phase": phase, "node": node_id, "value": value, "partials": partials}
    if derivative is not None:
        step["derivative"] = derivative
    return step


def _amount(arg: str, partial: float, amount: float) -> dict:
    return {"arg": arg, "partial": partial, "amount": _finite(lambda: amount, [], "The derivative")}


def forward_mode(graph: dict, inputs: dict, seed: str) -> dict:
    """A run computing each operation's value and tangent, in one pass, with respect to seed."""
    values = dict(graph["constants"])
    tangents = {}
    start = {"values": {}, "derivatives": {}}
    steps = []
    for node in graph["nodes"]:
        node_id = node["id"]
        if node["kind"] == "input":
            values[node_id] = start["values"][node_id] = inputs[node["name"]]
            tangents[node_id] = start["derivatives"][node_id] = 1.0 if node["name"] == seed else 0.0
        else:
            values[node_id] = _value(node["op"], [values[arg] for arg in node["args"]])
            partials = [_amount(arg, p, p * tangents[arg]) for arg, p in _partials(node, values, graph["constants"])]
            tangents[node_id] = _finite(sum, [[p["amount"] for p in partials]], "The derivative")
            steps.append(_step("forward", node_id, values[node_id], tangents[node_id], partials))
    return {"start": start, "steps": steps}


def reverse_mode(graph: dict, inputs: dict) -> dict:
    """A run computing each operation's value going forward, then each adjoint going backward."""
    values = dict(graph["constants"])
    start = {"values": {}, "derivatives": {}}
    steps = []
    for node in graph["nodes"]:
        if node["kind"] == "input":
            values[node["id"]] = start["values"][node["id"]] = inputs[node["name"]]
        else:
            values[node["id"]] = _value(node["op"], [values[arg] for arg in node["args"]])
            steps.append(_step("forward", node["id"], values[node["id"]], None, []))

    output = graph["output"]
    adjoints = {output: 1.0}
    steps.append(_step("backward", output, values[output], 1.0, []))
    for node in reversed(graph["nodes"]):
        if node["kind"] == "op":
            adjoint = adjoints[node["id"]]
            partials = []
            for arg, p in _partials(node, values, graph["constants"]):
                partials.append(_amount(arg, p, adjoint * p))
                adjoints[arg] = _finite(VALUES["add"], [adjoints.get(arg, 0.0), partials[-1]["amount"]], "The derivative")
            steps.append(_step("backward", node["id"], values[node["id"]], adjoint, partials))
    return {"start": start, "steps": steps}


def state_after(run: dict, k: int) -> dict:
    """Every value and derivative known after k steps, and the latest step, if any."""
    values = dict(run["start"]["values"])
    derivatives = dict(run["start"]["derivatives"])
    for step in run["steps"][:k]:
        values[step["node"]] = step["value"]
        if "derivative" in step:
            derivatives[step["node"]] = step["derivative"]
        if step["phase"] == "backward":
            for p in step["partials"]:
                derivatives[p["arg"]] = derivatives.get(p["arg"], 0.0) + p["amount"]
    state = {"values": values, "derivatives": derivatives}
    if k > 0:
        state["step"] = run["steps"][k - 1]
    return state