# Copyright 2018 D-Wave Systems Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS F ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""A dimod :term:`sampler` that uses the MST2 multistart tabu search algorithm."""

import numpy as np
import dimod

from typing import Optional

from dimod.core.initialized import InitialStateGenerator
from dwave.samplers.tabu import TabuSearch

__all__ = ["TabuSampler"]


class TabuSampler(dimod.Sampler, dimod.Initialized):
    """A tabu-search sampler.

    `Tabu search <https://en.wikipedia.org/wiki/Tabu_search>`_ is a heuristic that
    employs local search and can escape local minima by maintaining a "tabu list" of
    recently explored states that it does not revisit. This sampler implements the
    `MST2 multistart tabu search algorithm.
    <https://link.springer.com/article/10.1023/B:ANOR.0000039522.58036.68>`_
    for quadratic unconstrained binary optimization (QUBO) problems.

    Examples:
        This example solves a two-variable Ising model.

        >>> from dwave.samplers import TabuSampler
        >>> samples = TabuSampler().sample_ising({'a': -0.5, 'b': 1.0}, {'ab': -1})
        >>> list(samples.data()) # doctest: +SKIP
        [Sample(sample={'a': -1, 'b': -1}, energy=-1.5, num_occurrences=1)]
        >>> samples.first.energy
        -1.5

    """

    properties = None
    parameters = None

    def __init__(self):
        self.parameters = {
            'initial_states': [],
            'initial_states_generator': [],
            'num_reads': [],
            'seed': [],
            'tenure': [],
            'timeout': [],
            'num_restarts': [],
            'energy_threshold': [],
            'coefficient_z_first': [],
            'coefficient_z_restart': [],
            'lower_bound_z': [],
        }
        self.properties = {}

    def sample(self, bqm: dimod.BinaryQuadraticModel,
               initial_states: Optional[dimod.typing.SamplesLike] = None,
               initial_states_generator: InitialStateGenerator = 'random',
               num_reads: Optional[int] = None,
               seed: Optional[int] = None,
               tenure: Optional[int] = None,
               timeout: int = 20,
               num_restarts: int = 1000000,
               energy_threshold: Optional[float] = None,
               coefficient_z_first: Optional[int] = None,
               coefficient_z_restart: Optional[int] = None,
               lower_bound_z: Optional[int] = None,
               **kwargs) -> dimod.SampleSet:
        """Run a multistart tabu search on a given binary quadratic model.

        Args:
            bqm: Binary quadratic model (BQM) to be sampled.

            initial_states:
                One or more samples, each defining an initial state for all the
                problem variables. Initial states are given one per read, but
                if fewer than ``num_reads`` initial states are defined, additional
                values are generated as specified by ``initial_states_generator``.

            initial_states_generator:
                Defines the expansion of ``initial_states`` if fewer than
                ``num_reads`` are specified:

                * "none":
                    If the number of initial states specified is smaller than
                    ``num_reads``, raises ``ValueError``.

                * "tile":
                    Reuses the specified initial states if fewer than ``num_reads``
                    or truncates if greater.

                * "random":
                    Expands the specified initial states with randomly generated
                    states if fewer than ``num_reads`` or truncates if greater.

            num_reads:
                Number of reads. Each read is generated by one run of the tabu
                algorithm. If ``num_reads`` is not explicitly given, it is selected
                to match the number of initial states given. If initial states
                are not provided, only one read is performed.

            seed:
                32-bit unsigned integer seed to use for the PRNG. If the ``timeout``
                parameter is not None, results from the same seed may not be
                identical between runs due to finite clock resolution.

            tenure:
                Tabu tenure, which is the length of the tabu list, or number of
                recently explored solutions kept in memory.
                Default is a quarter of the number of problem variables up to
                a maximum value of 20.

            timeout:
                Maximum running time per read in milliseconds.

            num_restarts:
                Maximum number of tabu search restarts per read. Setting this value
                to zero results in a simple tabu search. 

            energy_threshold:
                Terminate when an energy lower than or equal to ``energy_threshold`` is found.

            coefficient_z_first:
                :code:`max(bqm.num_variables*coefficient_z_first, lower_bound_z)`
                bounds the number of variable updates considered in
                the first simple tabu search (STS). Variable updates arising from 
                the STS greedy-descent subroutine, invoked upon discovery of new 
                global optima, are excluded from the count. The 
                coefficient defaults to 10000 for small problems (up to 500 
                variables) and 25000 for larger problems.
                
            coefficient_z_restart:
                Controls the number of variable updates on restarted simple tabu
                search stages, matching the description for ``coefficient_z_first``. 
                The coefficient defaults to :code:`coefficient_z_first/4`
                
            lower_bound_z:
                Sets a minimum number of variable updates on all simple tabu
                searches, see ``coefficient_z_first``. The bound defaults to 
                500000. 

        Examples:
            This example samples a simple two-variable Ising model.

            >>> import dimod
            >>> bqm = dimod.BQM.from_ising({}, {'ab': 1})

            >>> from dwave.samplers import TabuSampler
            >>> sampler = TabuSampler()

            >>> samples = sampler.sample(bqm)
            >>> samples.first.energy
            -1.0
        """

        if not bqm:
            return dimod.SampleSet.from_samples([], energy=0, vartype=bqm.vartype)

        if tenure is None:
            tenure = min(20, len(bqm) // 4)
        elif not isinstance(tenure, int):
            raise TypeError("'tenure' should be an integer in range [0, num_vars - 1]")
        elif not 0 <= tenure < len(bqm):
            raise ValueError("'tenure' should be an integer in range [0, num_vars - 1]")

        # Get initial_states in binary form
        parsed = self.parse_initial_states(bqm.binary,
                                           initial_states=initial_states,
                                           initial_states_generator=initial_states_generator,
                                           num_reads=num_reads,
                                           seed=seed)

        parsed_initial_states = np.ascontiguousarray(parsed.initial_states.record.sample)

        qubo, varorder = self._bqm_to_tabu_qubo(bqm.binary)

        if timeout is None:
            timeout = -1    # Using negative timeout to mean ignore timeout parameter

        # run Tabu search
        samples = np.empty((parsed.num_reads, len(bqm)), dtype=np.int8)

        rng = np.random.default_rng(seed)

        restarts = []
        for ni, initial_state in enumerate(parsed_initial_states):
            seed_per_read = rng.integers(2**32, dtype=np.uint32)
            r = TabuSearch(qubo, initial_state, tenure, timeout, num_restarts, seed_per_read, energy_threshold, coefficient_z_first, coefficient_z_restart, lower_bound_z)
            samples[ni, :] = r.bestSolution()
            restarts.append(r.numRestarts())

        # we received samples in binary form, so convert if needed
        if bqm.vartype is dimod.SPIN:
            samples *= 2
            samples -= 1
        elif bqm.vartype is not dimod.BINARY:
            # sanity check
            raise ValueError("unknown vartype")

        return dimod.SampleSet.from_samples_bqm((samples, varorder), bqm=bqm, num_restarts=restarts)

    @staticmethod
    def _bqm_to_tabu_qubo(bqm):
        # construct dense matrix representation
        ldata, (irow, icol, qdata), offset, varorder = bqm.binary.to_numpy_vectors(return_labels=True)
        ud = np.zeros((len(bqm), len(bqm)), dtype=np.double)
        ud[np.diag_indices(len(bqm), 2)] = ldata
        ud[irow, icol] = qdata

        # Note: normally, conversion would be: `ud + ud.T - np.diag(np.diag(ud))`,
        # but the Tabu solver we're using requires slightly different qubo matrix.
        ud *= .5
        symm = ud + ud.T
        return symm, varorder
