"""Support for random optimizers, including the random-greedy path."""

import functools
import heapq
import math
import time
from collections import deque
from decimal import Decimal
from random import choices as random_choices
from random import seed as random_seed
from typing import Any, Dict, Generator, Iterable, List, Optional, Tuple, Union

from opt_einsum import helpers, paths
from opt_einsum.typing import ArrayIndexType, ArrayType, PathType

__all__ = ["RandomGreedy", "random_greedy", "random_greedy_128"]


class RandomOptimizer(paths.PathOptimizer):
    """Base class for running any random path finder that benefits
    from repeated calling, possibly in a parallel fashion. Custom random
    optimizers should subclass this, and the `setup` method should be
    implemented with the following signature:

    ```python
    def setup(self, inputs, output, size_dict):
        # custom preparation here ...
        return trial_fn, trial_args
    ```

    Where `trial_fn` itself should have the signature::

    ```python
    def trial_fn(r, *trial_args):
        # custom computation of path here
        return ssa_path, cost, size
    ```

    Where `r` is the run number and could for example be used to seed a
    random number generator. See `RandomGreedy` for an example.


    Parameters:
        max_repeats: The maximum number of repeat trials to have.
        max_time: The maximum amount of time to run the algorithm for.
        minimize:  Whether to favour paths that minimize the total estimated flop-count or
            the size of the largest intermediate created.
        parallel: Whether to parallelize the random trials, by default `False`. If
            `True`, use a `concurrent.futures.ProcessPoolExecutor` with the same
            number of processes as cores. If an integer is specified, use that many
            processes instead. Finally, you can supply a custom executor-pool which
            should have an API matching that of the python 3 standard library
            module `concurrent.futures`. Namely, a `submit` method that returns
            `Future` objects, themselves with `result` and `cancel` methods.
        pre_dispatch: If running in parallel, how many jobs to pre-dispatch so as to avoid
            submitting all jobs at once. Should also be more than twice the number
            of workers to avoid under-subscription. Default: 128.

    Attributes:
        path: The best path found so far.
        costs: The list of each trial's costs found so far.
        sizes: The list of each trial's largest intermediate size so far.
    """

    def __init__(
        self,
        max_repeats: int = 32,
        max_time: Optional[float] = None,
        minimize: str = "flops",
        parallel: Union[bool, Decimal, int] = False,
        pre_dispatch: int = 128,
    ):
        if minimize not in ("flops", "size"):
            raise ValueError("`minimize` should be one of {'flops', 'size'}.")

        self.max_repeats = max_repeats
        self.max_time = max_time
        self.minimize = minimize
        self.better = paths.get_better_fn(minimize)
        self._parallel: Union[bool, Decimal, int] = False
        self.parallel = parallel
        self.pre_dispatch = pre_dispatch

        self.costs: List[int] = []
        self.sizes: List[int] = []
        self.best: Dict[str, Any] = {"flops": float("inf"), "size": float("inf")}

        self._repeats_start = 0
        self._executor: Any
        self._futures: Any

    @property
    def path(self) -> PathType:
        """The best path found so far."""
        return paths.ssa_to_linear(self.best["ssa_path"])

    @property
    def parallel(self) -> Union[bool, Decimal, int]:
        return self._parallel

    @parallel.setter
    def parallel(self, parallel: Union[bool, Decimal, int]) -> None:
        # shutdown any previous executor if we are managing it
        if getattr(self, "_managing_executor", False):
            self._executor.shutdown()

        self._parallel = parallel
        self._managing_executor = False

        if parallel is False:
            self._executor = None
            return

        if parallel is True:
            from concurrent.futures import ProcessPoolExecutor

            self._executor = ProcessPoolExecutor()
            self._managing_executor = True
            return

        if isinstance(parallel, (int, Decimal)):
            from concurrent.futures import ProcessPoolExecutor

            self._executor = ProcessPoolExecutor(int(parallel))
            self._managing_executor = True
            return

        # assume a pool-executor has been supplied
        self._executor = parallel

    def _gen_results_parallel(self, repeats: Iterable[int], trial_fn: Any, args: Any) -> Generator[Any, None, None]:
        """Lazily generate results from an executor without submitting all jobs at once."""
        self._futures = deque()

        # the idea here is to submit at least ``pre_dispatch`` jobs *before* we
        # yield any results, then do both in tandem, before draining the queue
        for r in repeats:
            if len(self._futures) < self.pre_dispatch:
                self._futures.append(self._executor.submit(trial_fn, r, *args))
                continue
            yield self._futures.popleft().result()

        while self._futures:
            yield self._futures.popleft().result()

    def _cancel_futures(self) -> None:
        if self._executor is not None:
            for f in self._futures:
                f.cancel()

    def setup(
        self,
        inputs: List[ArrayIndexType],
        output: ArrayIndexType,
        size_dict: Dict[str, int],
    ) -> Tuple[Any, Any]:
        raise NotImplementedError

    def __call__(
        self,
        inputs: List[ArrayIndexType],
        output: ArrayIndexType,
        size_dict: Dict[str, int],
        memory_limit: Optional[int] = None,
    ) -> PathType:
        self._check_args_against_first_call(inputs, output, size_dict)

        # start a timer?
        if self.max_time is not None:
            t0 = time.time()

        trial_fn, trial_args = self.setup(inputs, output, size_dict)

        r_start = self._repeats_start + len(self.costs)
        r_stop = r_start + self.max_repeats
        repeats = range(r_start, r_stop)

        # create the trials lazily
        if self._executor is not None:
            trials = self._gen_results_parallel(repeats, trial_fn, trial_args)
        else:
            trials = (trial_fn(r, *trial_args) for r in repeats)

        # assess the trials
        for ssa_path, cost, size in trials:
            # keep track of all costs and sizes
            self.costs.append(cost)
            self.sizes.append(size)

            # check if we have found a new best
            found_new_best = self.better(cost, size, self.best["flops"], self.best["size"])

            if found_new_best:
                self.best["flops"] = cost
                self.best["size"] = size
                self.best["ssa_path"] = ssa_path

            # check if we have run out of time
            if (self.max_time is not None) and (time.time() > t0 + self.max_time):
                break

        self._cancel_futures()
        return self.path

    def __del__(self):
        # if we created the parallel pool-executor, shut it down
        if getattr(self, "_managing_executor", False):
            self._executor.shutdown()


def thermal_chooser(queue, remaining, nbranch=8, temperature=1, rel_temperature=True):
    """A contraction 'chooser' that weights possible contractions using a
    Boltzmann distribution. Explicitly, given costs `c_i` (with `c_0` the
    smallest), the relative weights, `w_i`, are computed as:

        $$w_i = exp( -(c_i - c_0) / temperature)$$

    Additionally, if `rel_temperature` is set, scale `temperature` by
    `abs(c_0)` to account for likely fluctuating cost magnitudes during the
    course of a contraction.

    Parameters:
        queue: The heapified list of candidate contractions.
        remaining: Mapping of remaining inputs' indices to the ssa id.
        temperature: When choosing a possible contraction, its relative probability will be
            proportional to `exp(-cost / temperature)`. Thus the larger
            `temperature` is, the further random paths will stray from the normal
            'greedy' path. Conversely, if set to zero, only paths with exactly the
            same cost as the best at each step will be explored.
        rel_temperature: Whether to normalize the `temperature` at each step to the scale of
            the best cost. This is generally beneficial as the magnitude of costs
            can vary significantly throughout a contraction.
        nbranch: How many potential paths to calculate probability for and choose from at each step.

    Returns:
        cost
        k1
        k2
        k3
    """
    n = 0
    choices = []
    while queue and n < nbranch:
        cost, k1, k2, k12 = heapq.heappop(queue)
        if k1 not in remaining or k2 not in remaining:
            continue  # candidate is obsolete
        choices.append((cost, k1, k2, k12))
        n += 1

    if n == 0:
        return None
    if n == 1:
        return choices[0]

    costs = [choice[0][0] for choice in choices]
    cmin = costs[0]

    # adjust by the overall scale to account for fluctuating absolute costs
    if rel_temperature:
        temperature *= max(1, abs(cmin))

    # compute relative probability for each potential contraction
    if temperature == 0.0:
        energies = [1 if c == cmin else 0 for c in costs]
    else:
        # shift by cmin for numerical reasons
        energies = [math.exp(-(c - cmin) / temperature) for c in costs]

    # randomly choose a contraction based on energies
    (chosen,) = random_choices(range(n), weights=energies)
    cost, k1, k2, k12 = choices.pop(chosen)

    # put the other choice back in the heap
    for other in choices:
        heapq.heappush(queue, other)

    return cost, k1, k2, k12


def ssa_path_compute_cost(
    ssa_path: PathType,
    inputs: List[ArrayIndexType],
    output: ArrayIndexType,
    size_dict: Dict[str, int],
) -> Tuple[int, int]:
    """Compute the flops and max size of an ssa path."""
    inputs = list(map(frozenset, inputs))
    output = frozenset(output)
    remaining = set(range(len(inputs)))
    total_cost = 0
    max_size = 0

    for i, j in ssa_path:
        k12, flops12 = paths.calc_k12_flops(inputs, output, remaining, i, j, size_dict)  # type: ignore
        remaining.discard(i)
        remaining.discard(j)
        remaining.add(len(inputs))
        inputs.append(k12)
        total_cost += flops12
        max_size = max(max_size, helpers.compute_size_by_dict(k12, size_dict))

    return total_cost, max_size


def _trial_greedy_ssa_path_and_cost(
    r: int,
    inputs: List[ArrayIndexType],
    output: ArrayIndexType,
    size_dict: Dict[str, int],
    choose_fn: Any,
    cost_fn: Any,
) -> Tuple[PathType, int, int]:
    """A single, repeatable, greedy trial run. **Returns:** ``ssa_path`` and cost."""
    if r == 0:
        # always start with the standard greedy approach
        choose_fn = None

    random_seed(r)

    ssa_path = paths.ssa_greedy_optimize(inputs, output, size_dict, choose_fn, cost_fn)
    cost, size = ssa_path_compute_cost(ssa_path, inputs, output, size_dict)

    return ssa_path, cost, size


class RandomGreedy(RandomOptimizer):
    def __init__(
        self,
        cost_fn: str = "memory-removed-jitter",
        temperature: float = 1.0,
        rel_temperature: bool = True,
        nbranch: int = 8,
        **kwargs: Any,
    ):
        """Parameters:
        cost_fn: A function that returns a heuristic 'cost' of a potential contraction
                with which to sort candidates. Should have signature
                `cost_fn(size12, size1, size2, k12, k1, k2)`.
        temperature: When choosing a possible contraction, its relative probability will be
                proportional to `exp(-cost / temperature)`. Thus the larger
                `temperature` is, the further random paths will stray from the normal
                'greedy' path. Conversely, if set to zero, only paths with exactly the
                same cost as the best at each step will be explored.
        rel_temperature: Whether to normalize the ``temperature`` at each step to the scale of
                the best cost. This is generally beneficial as the magnitude of costs
                can vary significantly throughout a contraction. If False, the
                algorithm will end up branching when the absolute cost is low, but
                stick to the 'greedy' path when the cost is high - this can also be
                beneficial.
        nbranch: How many potential paths to calculate probability for and choose from at each step.
        kwargs: Supplied to RandomOptimizer.
        """
        self.cost_fn = cost_fn
        self.temperature = temperature
        self.rel_temperature = rel_temperature
        self.nbranch = nbranch
        super().__init__(**kwargs)

    @property
    def choose_fn(self) -> Any:
        """The function that chooses which contraction to take - make this a
        property so that ``temperature`` and ``nbranch`` etc. can be updated
        between runs.
        """
        if self.nbranch == 1:
            return None

        return functools.partial(
            thermal_chooser,
            temperature=self.temperature,
            nbranch=self.nbranch,
            rel_temperature=self.rel_temperature,
        )

    def setup(
        self,
        inputs: List[ArrayIndexType],
        output: ArrayIndexType,
        size_dict: Dict[str, int],
    ) -> Tuple[Any, Any]:
        fn = _trial_greedy_ssa_path_and_cost
        args = (inputs, output, size_dict, self.choose_fn, self.cost_fn)
        return fn, args


def random_greedy(
    inputs: List[ArrayIndexType],
    output: ArrayIndexType,
    idx_dict: Dict[str, int],
    memory_limit: Optional[int] = None,
    **optimizer_kwargs: Any,
) -> ArrayType:
    """A simple wrapper around the `RandomGreedy` optimizer."""
    optimizer = RandomGreedy(**optimizer_kwargs)
    return optimizer(inputs, output, idx_dict, memory_limit)


random_greedy_128 = functools.partial(random_greedy, max_repeats=128)