"""Functions to implement the randomized optimization and search algorithms."""

# Authors: Genevieve Hayes (modified by Andrew Rollings, Kyle Nakamura)
# License: BSD 3 clause

import numpy as np
from typing import Callable, Any
from mlrose_ky.decorators import short_name
from mlrose_ky.algorithms.decay import GeometricDecay


@short_name("sa")
def simulated_annealing(
    problem: Any,
    schedule: Any = GeometricDecay(),
    max_attempts: int = 10,
    max_iters: int = np.inf,
    init_state: np.ndarray = None,
    curve: bool = False,
    random_state: int = None,
    state_fitness_callback: Callable = None,
    callback_user_info: Any = None,
) -> tuple[np.ndarray, float, np.ndarray | None]:
    """Use simulated annealing to find the optimum for a given optimization problem.

    Parameters
    ----------
    problem: optimization object
        Object containing fitness function optimization problem to be solved.
        For example, :code:`DiscreteOpt()`, :code:`ContinuousOpt()` or
        :code:`TSPOpt()`.
    schedule: schedule object, default: :code:`mlrose_ky.GeometricDecay()`
        Schedule used to determine the value of the temperature parameter.
    max_attempts: int, default: 10
        Maximum number of attempts to find a better neighbor at each step.
    max_iters: int, default: np.inf
        Maximum number of iterations of the algorithm.
    init_state: np.ndarray, default: None
        1-D Numpy array containing starting state for algorithm.
        If :code:`None`, then a random state is used.
    curve: bool, default: False
        Boolean to keep fitness values for a curve.
        If :code:`False`, then no curve is stored.
        If :code:`True`, then a history of fitness values is provided as a
        third return value.
    random_state: int, default: None
        If random_state is a positive integer, random_state is the seed used
        by np.random.seed(); otherwise, the random seed is not set.
    state_fitness_callback: function taking five parameters, default: None
        If specified, this callback will be invoked once per iteration.
        Parameters are (iteration, max attempts reached?, current best state, current best fit, user callback data).
        Return true to continue iterating, or false to stop.
    callback_user_info: any, default: None
        User data passed as last parameter of callback.

    Returns
    -------
    best_state: np.ndarray
        Numpy array containing state that optimizes the fitness function.
    best_fitness: float
        Value of fitness function at best state.
    fitness_curve: np.ndarray
        Numpy array containing the fitness at every iteration.
        Only returned if input argument :code:`curve` is :code:`True`.

    References
    ----------
    Russell, S. and P. Norvig (2010). *Artificial Intelligence: A Modern Approach*, 3rd edition.
    Prentice Hall, New Jersey, USA.
    """
    if not isinstance(max_attempts, int) or max_attempts < 0:
        raise ValueError(f"max_attempts must be a positive integer. Got {max_attempts}")
    if not (isinstance(max_iters, int) or max_iters == np.inf) or max_iters < 0:
        raise ValueError(f"max_iters must be a positive integer or np.inf. Got {max_iters}")
    if init_state is not None and len(init_state) != problem.get_length():
        raise ValueError(
            f"init_state must have the same length as the problem. Expected length {problem.get_length()}, got {len(init_state)}"
        )

    # Set random seed
    if isinstance(random_state, int) and random_state > 0:
        np.random.seed(random_state)

    # Initialize problem
    fitness_curve = []
    if init_state is None:
        problem.reset()
    else:
        problem.set_state(init_state)
    if state_fitness_callback is not None:
        # initial call with base data
        state_fitness_callback(
            iteration=0,
            state=problem.get_state(),
            fitness=problem.get_adjusted_fitness(),
            fitness_evaluations=problem.fitness_evaluations,
            user_data=callback_user_info,
        )

    attempts = 0
    iters = 0
    continue_iterating = True
    while attempts < max_attempts and iters < max_iters:
        temp = schedule.evaluate(iters)
        iters += 1
        problem.current_iteration += 1

        if temp == 0:
            break
        else:
            # Find random neighbor and evaluate fitness
            next_state = problem.random_neighbor()
            next_fitness = problem.eval_fitness(next_state)

            # Calculate delta E and change prob
            current_fitness = problem.get_fitness()
            delta_e = next_fitness - current_fitness
            prob = np.exp(delta_e / temp)

            # If best neighbor is an improvement or random value is less
            # than prob, move to that state and reset attempts counter
            if delta_e > 0 or np.random.uniform() < prob:
                problem.set_state(next_state)
                attempts = 0
            else:
                attempts += 1

        if curve:
            fitness_curve.append((problem.get_adjusted_fitness(), problem.fitness_evaluations))

        # invoke callback
        if state_fitness_callback is not None:
            max_attempts_reached = attempts == max_attempts or iters == max_iters or problem.can_stop()
            continue_iterating = state_fitness_callback(
                iteration=iters,
                attempt=attempts + 1,
                done=max_attempts_reached,
                state=problem.get_state(),
                fitness=problem.get_adjusted_fitness(),
                fitness_evaluations=problem.fitness_evaluations,
                curve=np.asarray(fitness_curve) if curve else None,
                user_data=callback_user_info,
            )

        # break out if requested
        if not continue_iterating:
            break

    best_fitness = problem.get_maximize() * problem.get_fitness()
    best_state = problem.get_state()

    if curve:
        return best_state, best_fitness, np.asarray(fitness_curve)

    return best_state, best_fitness, None
