import jax.numpy as jnp
from jax import lax

from simple_pytree import Pytree, dataclass

from typing import Callable


@dataclass
class FixedPointState(Pytree):
    """State information

    Attributes:
        iter_num (int): iteration number
        step_norm (jnp.ndarray): norm of step
        root_norm (jnp.ndarray): norm of root
    """

    iter_num: int
    step_norm: jnp.ndarray
    root_norm: jnp.ndarray


@dataclass
class OptStep(Pytree):
    """Current guess of the solution and state information

    Attributes:
        x (jnp.ndarray): current guess of the solution
        state (FixedPointState): state information
    """

    x: jnp.ndarray
    state: FixedPointState


@dataclass
class fxp_root(Pytree, mutable=True):
    """Initialize the solver state.

    Attributes:
        fun (Callable): system of fixed-point equations
        step_tol (float): tolerance level for step norm
        root_tol (float): tolerance level for root norm
        max_iter (int): maximum number of iterations
    """

    fun: Callable

    step_tol: float = 1e-8
    root_tol: float = 1e-6
    max_iter: int = 1000
    accelerator: str | None = None

    def _init_state(self) -> FixedPointState:
        """Initialize the solver state.

        Returns:
            state (FixedPointState)
        """
        return FixedPointState(
            iter_num=0,
            step_norm=jnp.asarray(jnp.inf),
            root_norm=jnp.asarray(jnp.inf),
        )

    def _update_none(self, fxp: Callable, step: OptStep, *args, **kwargs) -> OptStep:
        """Update fixed-point

        Args:
            fxp (Callable): system of fixed-point equations
            step (OptStep): current states

        Returns:
            step_next (OptStep): updated states
        """
        x_next, root = fxp(step.x, *args, **kwargs)

        next_state = FixedPointState(
            iter_num=step.state.iter_num + 1,
            step_norm=jnp.linalg.norm(x_next - step.x),
            root_norm=jnp.linalg.norm(root),
        )
        return OptStep(x=x_next, state=next_state)

    def _update_squarem(self, fxp: Callable, step: OptStep, *args, **kwargs) -> OptStep:
        """Update fixed-point by SQUAREM

        Args:
            fxp (Callable): system of fixed-point equations
            step (OptStep): current states

        Returns:
            step_next (OptStep): updated states
        """

        x1 = fxp(step.x, *args, **kwargs)[0]  # first fixed-point step
        x2 = fxp(x1, *args, **kwargs)[0]  # second fixed-point step

        # Accelerated step
        r = x1 - step.x  # change
        v = x2 - x1 - r  # curvature

        alpha = -jnp.sqrt(jnp.sum(r**2) / jnp.sum(v**2))

        x_next, root = fxp(step.x - 2 * alpha * r + (alpha**2) * v, *args, **kwargs)

        next_state = FixedPointState(
            iter_num=step.state.iter_num + 1,
            step_norm=jnp.linalg.norm(x_next - step.x),
            root_norm=jnp.linalg.norm(root),
        )
        return OptStep(x=x_next, state=next_state)

    def _update_squarem_r(
        self, fxp: Callable, step: OptStep, *args, **kwargs
    ) -> OptStep:
        """Update fixed-point by SQUAREM

        Args:
            fxp (Callable): system of fixed-point equations
            step (OptStep): current states

        Returns:
            step_next (OptStep): updated states
        """

        x1 = fxp(step.x, *args, **kwargs)[0]  # first fixed-point step
        x2 = fxp(x1, *args, **kwargs)[0]  # second fixed-point step

        q1 = x1 - step.x
        q2 = x2 - x1

        q2q1 = q2 - q1

        sr2 = jnp.matmul(q1.T, q1)
        sv2 = jnp.matmul(q2q1.T, q2q1)
        # srv = jnp.matmul(q1.T, q2q1)

        alpha = jnp.sqrt(sr2 / sv2)

        x_next, root = fxp(step.x - 2 * alpha * q1 + (alpha**2) * q2q1, *args, **kwargs)

        next_state = FixedPointState(
            iter_num=step.state.iter_num + 1,
            step_norm=jnp.linalg.norm(x_next - step.x),
            root_norm=jnp.linalg.norm(root),
        )
        return OptStep(x=x_next, state=next_state)

    def _condition(self, step: OptStep) -> jnp.ndarray:
        """Conditions for continuation of while loop

        Args:
            step (OptStep): current states

        Returns:
            condition (jnp.ndarray): boolean

        """
        # Evaluate stopping criterions
        cond1 = step.state.root_norm > self.root_tol  # stop if root is zero
        cond2 = step.state.step_norm > self.step_tol  # stop if step is zero
        cond3 = (
            step.state.iter_num < self.max_iter
        )  # stop if maximum iteration is reached
        cond4 = jnp.any(jnp.isnan(step.x)) == 0  # stop if any value is NaN

        cond_tol = jnp.logical_and(cond1, cond2)  # step or root is close to zero
        return jnp.all(jnp.logical_and(jnp.logical_and(cond_tol, cond3), cond4))

    def solve(self, guess: jnp.ndarray, *args, **kwargs) -> OptStep:
        """Solve fixed-point equation by fixed point iterations

        Args:
            guess (jnp.ndarray): initial guess

        Returns:
            state (OptStep): dataclass containing the solution to the fixed-point

        """

        if self.accelerator == "SQUAREM":
            # Set up accelerated fixed-point equations (SQUAREM)
            def _fun_squarem(step):
                return self._update_squarem(self.fun, step, *args, **kwargs)

            # Execute fixed-point iterations
            result = lax.while_loop(
                body_fun=_fun_squarem,
                cond_fun=self._condition,
                init_val=OptStep(x=guess, state=self._init_state()),
            )
        elif self.accelerator == "SQUAREM_R":
            # Set up accelerated fixed-point equations (SQUAREM)
            def _fun_squarem_r(step):
                return self._update_squarem_r(self.fun, step, *args, **kwargs)

            # Execute fixed-point iterations
            result = lax.while_loop(
                body_fun=_fun_squarem_r,
                cond_fun=self._condition,
                init_val=OptStep(x=guess, state=self._init_state()),
            )
        else:
            # Set up fixed-point equations
            def _fun_none(step):
                return self._update_none(self.fun, step, *args, **kwargs)

            # Execute fixed-point iterations
            result = lax.while_loop(
                body_fun=_fun_none,
                cond_fun=self._condition,
                init_val=OptStep(x=guess, state=self._init_state()),
            )

        # Print number of iterations and norm of root and step
        print(
            f"FXP output: iterations={result.state.iter_num}, root norm={result.state.root_norm}, step norm={result.state.step_norm}"
        )
        return result
