import os.path
import _pickle as pickle
import random
import re
from functools import partial
from typing import Union, Sequence, Optional, Callable

from mabwiser.mab import MAB, LearningPolicy
from scipy.special import softmax

from golem.core.dag.graph import Graph
from golem.core.dag.graph_node import GraphNode
from golem.core.optimisers.adaptive.operator_agent import OperatorAgent, ActType, ObsType, ExperienceBuffer
from golem.core.optimisers.adaptive.reward_agent import FitnessRateRankRewardTransformer
from golem.core.optimisers.adaptive.utils import get_callable_name
from golem.core.paths import default_data_dir


class MultiArmedBanditAgent(OperatorAgent):
    def __init__(self,
                 actions: Sequence[ActType],
                 n_jobs: int = 1,
                 enable_logging: bool = True,
                 decaying_factor: float = 1.0,
                 path_to_save: Optional[str] = None,
                 is_initial_fit: bool = True):
        super().__init__(actions=actions, enable_logging=enable_logging)
        self.actions = list(actions)
        self._indices = list(range(len(actions)))
        # str because parent operator for mutation is stored as string for custom mutations serialisation
        self._arm_by_action = dict(map(lambda x, y: (get_callable_name(x), y), actions, self._indices))
        self._agent = MAB(arms=self._indices,
                          learning_policy=LearningPolicy.EpsilonGreedy(epsilon=0.4),
                          n_jobs=n_jobs)
        self._reward_agent = FitnessRateRankRewardTransformer(decaying_factor=decaying_factor)
        if is_initial_fit:
            self._initial_fit()
        self._path_to_save = path_to_save

    def _initial_fit(self):
        n = len(self.actions)
        uniform_rewards = [1. / n] * n
        self._agent.fit(decisions=self._indices, rewards=uniform_rewards)

    def choose_action(self, obs: ObsType) -> ActType:
        arm = self._agent.predict()
        action = self.actions[arm]
        return action

    def get_action_values(self, obs: Optional[ObsType] = None) -> Sequence[float]:
        prob_dict = self._agent.predict_expectations()
        prob_list = [prob_dict[i] for i in range(len(prob_dict))]
        return prob_list

    def get_action_probs(self, obs: Optional[ObsType] = None) -> Sequence[float]:
        return softmax(self.get_action_values())

    def choose_nodes(self, graph: Graph, num_nodes: int = 1) -> Union[GraphNode, Sequence[GraphNode]]:
        subject_nodes = random.sample(graph.nodes, k=num_nodes)
        return subject_nodes[0] if num_nodes == 1 else subject_nodes

    def partial_fit(self, experience: ExperienceBuffer):
        """Continues learning of underlying agent with new experience."""
        _, arms, processed_rewards = self._get_experience(experience)
        self._agent.partial_fit(decisions=arms, rewards=processed_rewards)

    def _get_experience(self, experience: ExperienceBuffer):
        """ Get experience from ExperienceBuffer, process rewards and log. """
        obs, actions, rewards = experience.retrieve_experience()
        arms = [self._arm_by_action[action] for action in actions]
        processed_rewards = self._reward_agent.get_rewards_for_arms(rewards, arms)
        self._dbg_log(obs, actions, processed_rewards)
        return obs, arms, processed_rewards

    def save(self, path_to_save: Optional[str] = None):
        """ Saves bandit to specified file. """

        if not path_to_save:
            path_to_save = self._path_to_save

        # if path was not specified at all
        if not path_to_save:
            path_to_save = os.path.join(default_data_dir(), 'MAB')

        if not path_to_save.endswith('.pkl'):
            os.makedirs(path_to_save, exist_ok=True)
            mabs_num = [int(name.split('_')[0]) for name in os.listdir(path_to_save)
                        if re.fullmatch(r'\d_mab.pkl', name)]
            if not mabs_num:
                max_saved_mab = 0
            else:
                max_saved_mab = max(mabs_num) + 1
            path_to_file = os.path.join(path_to_save, f'{max_saved_mab}_mab.pkl')
        else:
            path_to_dir = os.path.dirname(path_to_save)
            os.makedirs(path_to_dir, exist_ok=True)
            path_to_file = path_to_save
        with open(path_to_file, 'wb') as f:
            pickle.dump(self, f)
        self._log.info(f"MAB was saved to {path_to_file}")

    @staticmethod
    def load(path: str):
        """ Loads bandit from the specified file. """
        with open(path, 'rb') as f:
            mab = pickle.load(f)
        return mab
