text: # -*- coding: utf-8 -*-
  # SPDX-License-Identifer: Apache-2.0
  
  """
  :Author: FMR LLC
  :Email: mabwiser@fmr.com
  
  This module defines the abstract base class for contextual multi-armed bandit algorithms.
  """
  
  import abc
  from itertools import chain
  from typing import Callable, Dict, List, NoReturn, Optional
  import multiprocessing as mp
  
  from joblib import Parallel, delayed
  import numpy as np
  
  from mabwiser.utils import Arm, Num
  
  
  class BaseMAB(metaclass=abc.ABCMeta):
      """Abstract base class for multi-armed bandits.
  
      This module is not intended to be used directly, instead it declares
      the basic skeleton of multi-armed bandits together with a set of parameters
      that are common to every bandit algorithm.
  
      It declares abstract methods that sub-classes can override to
      implement specific bandit policies using:
  
          - ``__init__`` constructor to initialize the bandit
          - ``add_arm`` method to add a new arm
          - ``fit`` method for training
          - ``partial_fit`` method for _online learning
          - ``predict_expectations`` method to retrieve the expectation of each arm
          - ``predict`` method for testing to retrieve the best arm based on the policy
  
      Attributes
      ----------
      rng:
          The random number generator.
      arms:
          The list of all arms.
      n_jobs:
          This is used to specify how many concurrent processes/threads should be used for parallelized routines.
          Default value is set to 1.
          If set to -1, all CPUs are used.
          If set to -2, all CPUs but one are used, and so on.
      backend: str, optional
          Specify a parallelization backend implementation supported in the joblib library. Supported options are:
          - “loky” used by default, can induce some communication and memory overhead when exchanging input and
            output data with the worker Python processes.
          - “multiprocessing” previous process-based backend based on multiprocessing.Pool. Less robust than loky.
          - “threading” is a very low-overhead backend but it suffers from the Python Global Interpreter Lock if the
            called function relies a lot on Python objects.
          Default value is None. In this case the default backend selected by joblib will be used.
      arm_to_expectation:
          The dictionary of arms (keys) to their expected rewards (values).
      """
  
      @abc.abstractmethod
      def __init__(self, rng: np.random.RandomState, arms: List[Arm], n_jobs: int, backend: str = None):
          """Abstract method.
  
          Creates a multi-armed bandit policy with the given arms.
          """
          self.rng: np.random.RandomState = rng
          self.arms: List[Arm] = arms
          self.n_jobs: int = n_jobs
          self.backend: str = backend
  
          self.arm_to_expectation: Dict[Arm, float] = dict.fromkeys(self.arms, 0)
  
      def add_arm(self, arm: Arm, binarizer: Callable = None, scaler: Callable = None) -> NoReturn:
          """Introduces a new arm to the bandit.
  
          Adds the new arm with zero expectations and
          calls the ``_uptake_new_arm()`` function of the sub-class.
          """
          self.arm_to_expectation[arm] = 0
          self._uptake_new_arm(arm, binarizer, scaler)
  
      @abc.abstractmethod
      def fit(self, decisions: np.ndarray, rewards: np.ndarray,
              contexts: Optional[np.ndarray] = None) -> NoReturn:
          """Abstract method.
  
          Fits the multi-armed bandit to the given
          decision and reward history and corresponding contexts if any.
          """
          pass
  
      @abc.abstractmethod
      def partial_fit(self, decisions: np.ndarray, rewards: np.ndarray,
                      contexts: Optional[np.ndarray] = None) -> NoReturn:
          """Abstract method.
  
          Updates the multi-armed bandit with the given
          decision and reward history and corresponding contexts if any.
          """
          pass
  
      @abc.abstractmethod
      def predict(self, contexts: Optional[np.ndarray] = None) -> Arm:
          """Abstract method.
  
          Returns the predicted arm.
          """
          pass
  
      @abc.abstractmethod
      def predict_expectations(self, contexts: Optional[np.ndarray] = None) -> Dict[Arm, Num]:
          """Abstract method.
  
          Returns a dictionary from arms (keys) to their expected rewards (values).
          """
          pass
  
      @abc.abstractmethod
      def _uptake_new_arm(self, arm: Arm, binarizer: Callable = None, scaler: Callable = None) -> NoReturn:
          """Abstract method.
  
          Updates the multi-armed bandit with the new arm.
          """
          pass
  
      @abc.abstractmethod
      def _fit_arm(self, arm: Arm, decisions: np.ndarray, rewards: np.ndarray,
                   contexts: Optional[np.ndarray] = None) -> NoReturn:
          """Abstract method.
  
          Fit operation for individual arm.
          """
          pass
  
      @abc.abstractmethod
      def _predict_contexts(self, contexts: np.ndarray, is_predict: bool,
                            seeds: Optional[np.ndarray] = None, start_index: Optional[int] = None) -> List:
          """Abstract method.
  
          Predict operation for set of contexts.
          """
          pass
  
      def _parallel_fit(self, decisions: np.ndarray, rewards: np.ndarray,
                        contexts: Optional[np.ndarray] = None):
  
          # Compute effective number of jobs
          n_jobs = self._effective_jobs(len(self.arms), self.n_jobs)
  
          # Perform parallel fit
          Parallel(n_jobs=n_jobs, require='sharedmem')(
                            delayed(self._fit_arm)(
                                arm, decisions, rewards, contexts)
                            for arm in self.arms)
  
      def _parallel_predict(self, contexts: np.ndarray, is_predict: bool):
  
          # Total number of contexts to predict
          n_contexts = len(contexts)
  
          # Partition contexts by job
          n_jobs, n_contexts, starts = self._partition_contexts(n_contexts)
          total_contexts = sum(n_contexts)
  
          # Get seed value for each context
          seeds = self.rng.randint(np.iinfo(np.int32).max, size=total_contexts)
  
          # Perform parallel predictions
          predictions = Parallel(n_jobs=n_jobs, backend=self.backend)(
                            delayed(self._predict_contexts)(
                                contexts[starts[i]:starts[i + 1]],
                                is_predict,
                                seeds[starts[i]:starts[i + 1]],
                                starts[i])
                            for i in range(n_jobs))
  
          # Reduce
          predictions = list(chain.from_iterable(t for t in predictions))
  
          return predictions if len(predictions) > 1 else predictions[0]
  
      def _partition_contexts(self, n_contexts: int):
  
          # Compute effective number of jobs
          n_jobs = self._effective_jobs(n_contexts, self.n_jobs)
  
          # Partition contexts between jobs
          n_contexts_per_job = np.full(n_jobs, n_contexts // n_jobs, dtype=np.int)
          n_contexts_per_job[:n_contexts % n_jobs] += 1
          starts = np.cumsum(n_contexts_per_job)
  
          return n_jobs, n_contexts_per_job.tolist(), [0] + starts.tolist()
  
      @staticmethod
      def _effective_jobs(size: int, n_jobs: int):
          if n_jobs < 0:
              n_jobs = max(mp.cpu_count() + 1 + n_jobs, 1)
          n_jobs = min(n_jobs, size)
          return n_jobs
  