from __future__ import annotations

from functools import partial
from typing import Any, Generic, Tuple, TypeVar

import jax.numpy as jnp
from jax.lax import while_loop

from ..annotations import BooleanNumeric, IntegralNumeric, PyTree, RealNumeric
from ..dataclasses import dataclass
from .augmented import AugmentedState, State
from .base import IteratedFunctionBase, Parameters, Trajectory

__all__ = ['IteratedFunction']


Comparand = TypeVar('Comparand', bound=PyTree)
TheAugmentedState = TypeVar('TheAugmentedState', bound=AugmentedState[Any])


@dataclass
class IteratedFunction(IteratedFunctionBase[Parameters, State, Trajectory, TheAugmentedState],
                       Generic[Parameters, State, Comparand, Trajectory, TheAugmentedState]):
    """
    An IteratedFunction object models an iterated function.

    It is a generic class in terms of five generic types, all of which are pytrees:
        * Parameters, which models the iteration parameters,
        * State, which is the state at each iteration,
        * Comparand, which is calculated from state and compared between iterations to check
          convergence,
        * Trajectory, which is the type produced by sample_trajectory, and
        * TheAugmentedState, which is typically defined by one of the tjax subclasses.

    The two main methods of IteratedFunction are
    * find_fixed_point, which iterates until the fixed point is reached or other stopping conditions
      are met, and
    * sample_trajectory, which iterates for a fixed number of iterations and returns a trajectory.

    Args:
        minimum_iterations: The minimum number of iterations for this fixed point solver.  This must
            be positive.
        maximum_iterations: The maximum number of iterations for this fixed point solver.  This must
            be positive.
        rtol: The relative tolerance for the comparison stopping condition.
        atol: The absolute tolerance for the comparison stopping condition.
    """
    minimum_iterations: IntegralNumeric
    maximum_iterations: IntegralNumeric
    rtol: RealNumeric
    atol: RealNumeric

    # New methods ----------------------------------------------------------------------------------
    def find_fixed_point(self,
                         theta: Parameters,
                         initial_state: State) -> TheAugmentedState:
        """
        Args:
            theta: The parameters for which gradients can be calculated.
            initial_state: An initial guess of the final state.
        Returns:
            state: The augmented state at the fixed point.
        """
        def f(augmented: TheAugmentedState) -> TheAugmentedState:
            new_state = self.sampled_state(theta, augmented.current_state)
            return self.iterate_augmented(new_state, augmented)

        return while_loop(partial(self.state_needs_iteration, theta),
                          f,
                          self.initial_augmented(initial_state))

    def debug_fixed_point(self, theta: Parameters, initial_state: State) -> TheAugmentedState:
        """
        This method is identical to find_fixed_point, but avoids using while_loop.
        """
        augmented = self.initial_augmented(initial_state)
        while self.state_needs_iteration(theta, augmented):
            new_state = self.sampled_state(theta, augmented.current_state)
            augmented = self.iterate_augmented(new_state, augmented)
        return augmented

    def state_needs_iteration(self, theta: Parameters, augmented: TheAugmentedState) -> bool:
        """
        Args:
            theta: The parameters.
            augmented: The state.
        Returns: True while iteration needs to continue.
        """
        enough_iterations = augmented.iterations >= self.minimum_iterations
        converged = self.converged(augmented)
        not_too_many_iterations = augmented.iterations < self.maximum_iterations
        return jnp.logical_and(not_too_many_iterations,
                               jnp.logical_not(jnp.logical_and(enough_iterations, converged)))

    # Abstract methods -----------------------------------------------------------------------------
    def expected_state(self, theta: Parameters, state: State) -> State:
        """
        Returns: The expected value of the next state given the old one.  This is used by the
            combinator.
        """
        raise NotImplementedError

    def converged(self, augmented: TheAugmentedState) -> BooleanNumeric:
        """
        Returns: A Boolean Array of shape () indicating whether the pytrees are close.
        """
        raise NotImplementedError

    def extract_comparand(self, state: State) -> Comparand:
        """
        Returns: The "comparand", which is a subset of the values in the state that are compared by
            subclasses to detect convergence.  In ComparingIteratedFunction, the comparand will be
            compared in successive states to detect convergence.  In StochasticIteratedFunction, the
            mean and variance of the comparand are used to detect convergence.
        """
        raise NotImplementedError

    def minimum_tolerances(self, augmented: TheAugmentedState) -> Tuple[RealNumeric, RealNumeric]:
        """
        Returns:
            The minimum value of atol that would lead to convergence now.
            The minimum value of rtol that would lead to convergence now.
        """
        raise NotImplementedError
