"""Objective functions and optimisers, in one or two variables.

Positions, gradients and Hessians are NumPy arrays of length 1 or 2, so each
optimiser is written once for both.
"""

import numpy as np
from sympy import Abs, Expr, acos, asin, atan, cos, exp, lambdify, log, pi, sin, sqrt, symbols, tan
from sympy.parsing.sympy_parser import (
    convert_xor,
    implicit_multiplication_application,
    parse_expr,
    standard_transformations,
)

# Real, so that e.g. d/dx |x| is sign(x)
VARIABLES = symbols("x y", real=True)

ALLOWED_NAMES = {
    "sin": sin, "cos": cos, "tan": tan, "asin": asin, "acos": acos, "atan": atan,
    "log": log, "ln": log, "sqrt": sqrt, "abs": Abs, "exp": exp,
    "e": exp(1), "pi": pi, "π": pi,
}


def parse_function(expression: str, n_variables: int) -> Expr:
    """Parse e.g. "x^2 + 3y^2"; raises ValueError with a message for students."""
    if not expression.strip():
        raise ValueError("Enter a function.")
    variables = VARIABLES[:n_variables]
    names = {str(v): v for v in VARIABLES} | ALLOWED_NAMES
    try:
        expr = parse_expr(
            expression,
            local_dict=names,
            transformations=standard_transformations + (implicit_multiplication_application, convert_xor),
        )
    except Exception:
        raise ValueError("This isn't a function sympy can read. Check the brackets and operators.")
    if not isinstance(expr, Expr):
        raise ValueError("This isn't a function sympy can read. Check the brackets and operators.")
    unknown = expr.free_symbols - set(variables)
    if unknown:
        allowed = " and ".join(str(v) for v in variables)
        raise ValueError(f"Unknown name {', '.join(sorted(map(str, unknown)))}: use only {allowed}.")
    return expr


class Objective:
    """f with its gradient and Hessian, evaluated at a point or on a grid."""

    def __init__(self, expr: Expr, n_variables: int):
        variables = VARIABLES[:n_variables]
        self.n = n_variables
        self._f = lambdify(variables, expr, modules="numpy")
        self._gradient = [lambdify(variables, expr.diff(v), modules="numpy") for v in variables]
        self._hessian = [[lambdify(variables, expr.diff(a, b), modules="numpy") for b in variables] for a in variables]

    def value(self, p: np.ndarray) -> float:
        return float(self._f(*p))

    def gradient(self, p: np.ndarray) -> np.ndarray:
        return np.array([float(g(*p)) for g in self._gradient])

    def hessian(self, p: np.ndarray) -> np.ndarray:
        return np.array([[float(h(*p)) for h in row] for row in self._hessian])

    def values_on(self, *grid: np.ndarray) -> np.ndarray:
        # Constant expressions give a scalar; spread it over the grid
        return np.broadcast_to(np.asarray(self._f(*grid), dtype=float), grid[0].shape)

    def gradient_on(self, *grid: np.ndarray) -> list[np.ndarray]:
        return [np.broadcast_to(np.asarray(g(*grid), dtype=float), grid[0].shape) for g in self._gradient]


class SingularHessian(Exception):
    pass


# Each optimiser keeps its own state between steps; step(p) returns the next position.

class GradientDescent:
    """Heavy-ball momentum: v ← μv − η∇f(x), x ← x + v."""

    def __init__(self, objective: Objective, learning_rate: float, momentum: float):
        self.f, self.lr, self.momentum = objective, learning_rate, momentum
        self.v = 0.0

    def step(self, p: np.ndarray) -> np.ndarray:
        self.v = self.momentum * self.v - self.lr * self.f.gradient(p)
        return p + self.v


class Nesterov:
    """Gradient taken at the look-ahead point x + μv."""

    def __init__(self, objective: Objective, learning_rate: float, momentum: float):
        self.f, self.lr, self.momentum = objective, learning_rate, momentum
        self.v = 0.0

    def step(self, p: np.ndarray) -> np.ndarray:
        self.v = self.momentum * self.v - self.lr * self.f.gradient(p + self.momentum * self.v)
        return p + self.v


class Newton:
    """x ← x − H⁻¹∇f(x)."""

    def __init__(self, objective: Objective):
        self.f = objective

    def step(self, p: np.ndarray) -> np.ndarray:
        h, g = self.f.hessian(p), self.f.gradient(p)
        try:
            delta = np.linalg.solve(h, g)
        except np.linalg.LinAlgError:
            raise SingularHessian()
        # Pyodide's NumPy can't see floating-point errors, so solve returns inf or nan instead of raising
        if np.all(np.isfinite(h)) and np.all(np.isfinite(g)) and not np.all(np.isfinite(delta)):
            raise SingularHessian()
        return p - delta


class Adagrad:
    def __init__(self, objective: Objective, learning_rate: float, epsilon: float):
        self.f, self.lr, self.eps = objective, learning_rate, epsilon
        self.v = 0.0  # sum of squared gradients

    def step(self, p: np.ndarray) -> np.ndarray:
        g = self.f.gradient(p)
        self.v = self.v + g**2
        return p - self.lr / np.sqrt(self.v + self.eps) * g


class RMSProp:
    def __init__(self, objective: Objective, learning_rate: float, beta: float, epsilon: float):
        self.f, self.lr, self.beta, self.eps = objective, learning_rate, beta, epsilon
        self.v = 0.0  # moving average of squared gradients

    def step(self, p: np.ndarray) -> np.ndarray:
        g = self.f.gradient(p)
        self.v = self.beta * self.v + (1 - self.beta) * g**2
        return p - self.lr / np.sqrt(self.v + self.eps) * g


class Adadelta:
    def __init__(self, objective: Objective, beta: float, epsilon: float):
        self.f, self.beta, self.eps = objective, beta, epsilon
        self.v = 0.0  # moving average of squared gradients
        self.s = 0.0  # moving average of squared updates

    def step(self, p: np.ndarray) -> np.ndarray:
        g = self.f.gradient(p)
        self.v = self.beta * self.v + (1 - self.beta) * g**2
        delta = -np.sqrt(self.s + self.eps) / np.sqrt(self.v + self.eps) * g
        self.s = self.beta * self.s + (1 - self.beta) * delta**2
        return p + delta


class Adam:
    def __init__(self, objective: Objective, learning_rate: float, beta1: float, beta2: float, epsilon: float):
        self.f, self.lr, self.beta1, self.beta2, self.eps = objective, learning_rate, beta1, beta2, epsilon
        self.m = 0.0  # first moment
        self.v = 0.0  # second moment
        self.t = 0

    def step(self, p: np.ndarray) -> np.ndarray:
        g = self.f.gradient(p)
        self.t += 1
        self.m = self.beta1 * self.m + (1 - self.beta1) * g
        self.v = self.beta2 * self.v + (1 - self.beta2) * g**2
        m_hat = self.m / (1 - self.beta1**self.t)
        v_hat = self.v / (1 - self.beta2**self.t)
        return p - self.lr / (np.sqrt(v_hat) + self.eps) * m_hat


# Ids match backend/src/algorithms.json
OPTIMISERS = {
    "gradient_descent": GradientDescent,
    "nesterov": Nesterov,
    "newton": Newton,
    "adam": Adam,
    "adagrad": Adagrad,
    "rmsprop": RMSProp,
    "adadelta": Adadelta,
}

# A run that gets further than this from the origin has diverged
DIVERGENCE_LIMIT = 1e6

DIVERGED = "The next point has no finite value or is beyond 1,000,000: the run diverged or left the function's domain."
SINGULAR = "The Hessian is singular here, so Newton's method can't take a step."


class Run:
    """One optimiser's trajectory, computed as far as it has been asked for."""

    def __init__(self, objective: Objective, optimiser, start: np.ndarray):
        self.f = objective
        self.optimiser = optimiser
        self.points = [self._describe(start)]
        self.position = start
        # Why the run ended, once it has
        self.stopped: str | None = None

    def point(self, k: int) -> dict | None:
        """The point after k iterations, or None if the run stopped before then."""
        with np.errstate(all="ignore"):
            while len(self.points) <= k and self.stopped is None:
                self._advance()
        return self.points[k] if k < len(self.points) else None

    def _advance(self) -> None:
        try:
            p = self.optimiser.step(self.position)
        except SingularHessian:
            self.stopped = SINGULAR
            return
        if not np.all(np.abs(p) <= DIVERGENCE_LIMIT) or not np.isfinite(self.f.value(p)):
            self.stopped = DIVERGED
            return
        self.position = p
        self.points.append(self._describe(p))

    def _describe(self, p: np.ndarray) -> dict:
        return {
            "position": p.tolist(),
            "value": self.f.value(p),
            "gradient": self.f.gradient(p).tolist(),
            "hessian": self.f.hessian(p).tolist(),
        }