import numpy as np
import copy
import combo.misc
import cPickle as pickle
from results import history
from .. import utility
from ...variable import variable
from ..call_simulator import call_simulator
from ... import predictor
from ...gp import predictor as gp_predictor
from ...blm import predictor as blm_predictor
import combo.search.score
MAX_SEACH = int(20000)


class policy:
    def __init__(self, test_X, config=None):
        self.predictor = None
        self.training = variable()
        self.test = self._set_test(test_X)
        self.actions = np.arange(0, self.test.X.shape[0])
        self.history = history()
        self.config = self._set_config(config)

    def set_seed(self, seed):
        self.seed = seed
        np.random.seed(self.seed)

    def delete_actions(self, index, actions=None):
        actions = self._set_unchosed_actions(actions)
        return np.delete(actions, index)

    def write(self, action, t, X=None):
        if X is None:
            X = self.test.X[action, :]
            Z = self.test.Z[action, :] if self.test.Z is not None else None
        else:
            Z = self.predictor.get_basis(X) \
                if self.predictor is not None else None

        self.new_data = variable(X, t, Z)
        self.history.write(t, action)
        self.training.add(X=X, t=t, Z=Z)

    def random_search(self, max_num_probes, num_search_each_probe=1,
                      simulator=None, is_disp=True):

        N = int(num_search_each_probe)

        if int(max_num_probes) * N > len(self.actions):
            raise ValueError('max_num_probes * num_search_each_probe must \
                be smaller than the length of candidates')

        if is_disp:
            utility.show_interactive_mode(simulator, self.history)

        for n in xrange(0, max_num_probes):

            if is_disp and N > 1:
                utility.show_start_message_multi_search(self.history.num_runs)

            action = self.get_random_action(N)

            if simulator is None:
                return action

            t, X = call_simulator(simulator, action)

            self.write(action, t, X)

            if is_disp:
                utility.show_search_results(self.history, N)

        return copy.deepcopy(self.history)

    def bayes_search(self, training=None, max_num_probes=None,
                     num_search_each_probe=1,
                     predictor=None, is_disp=True,
                     simulator=None, score='TS', interval=0,
                     num_rand_basis=0):

        if max_num_probes is None:
            max_num_probes = 1
            simulator = None

        is_rand_expans = False if num_rand_basis == 0 else True

        self.training = self._set_training(training)

        if predictor is None:
            self.predictor = self._init_predictor(is_rand_expans)
        else:
            self.predictor = predictor

        N = int(num_search_each_probe)

        for n in xrange(max_num_probes):

            if utility.is_learning(n, interval):
                self.predictor.fit(self.training, num_rand_basis)
                self.test.Z = self.predictor.get_basis(self.test.X)
                self.training.Z = self.predictor.get_basis(self.training.X)
                self.predictor.prepare(self.training)
            else:
                try:
                    self.predictor.update(self.training, self.new_data)
                except:
                    self.predictor.prepare(self.training)

            if num_search_each_probe != 1:
                utility.show_start_message_multi_search(self.history.num_runs,
                                                        score)

            K = self.config.search.multi_probe_num_sampling
            alpha = self.config.search.alpha
            action = self.get_actions(score, N, K, alpha)

            if simulator is None:
                return action

            t, X = call_simulator(simulator, action)

            self.write(action, t, X)

            if is_disp:
                utility.show_search_results(self.history, N)

        return copy.deepcopy(self.history)

    def get_score(self, mode, predictor=None, training=None, alpha=1):
        self._set_training(training)
        self._set_predictor(predictor)
        actions = self.actions

        test = self.test.get_subset(actions)
        if mode == 'EI':
            f = combo.search.score.EI(predictor, training, test)
        elif mode == 'PI':
            f = combo.search.score.PI(predictor, training, test)
        elif mode == 'TS':
            f = combo.search.score.TS(predictor, training, test, alpha)
        else:
            raise NotImplementedError('mode must be EI, PI or TS.')
        return f

    def get_marginal_score(self, mode, chosed_actions, N, alpha):
        f = np.zeros((N, len(self.actions)))
        new_test = self.test.get_subset(chosed_actions)
        virtual_t \
            = self.predictor.get_predict_samples(self.training, new_test, N)

        for n in xrange(N):
            predictor = copy.deepcopy(self.predictor)
            train = copy.deepcopy(self.training)
            virtual_train = new_test
            virtual_train.t = virtual_t[n, :]

            if virtual_train.Z is None:
                train.add(virtual_train.X, virtual_train.t)
            else:
                train.add(virtual_train.X, virtual_train.t, virtual_train.Z)

            try:
                predictor.update(train, virtual_train)
            except:
                predictor.prepare(train)

            f[n, :] = self.get_score(mode, predictor, train)
        return f

    def get_actions(self, mode, N, K, alpha):
        f = self.get_score(mode, self.predictor, self.training, alpha)
        temp = np.argmax(f)
        action = self.actions[temp]
        self.actions = self.delete_actions(temp)

        chosed_actions = np.zeros(N, dtype=int)
        chosed_actions[0] = action

        for n in xrange(1, N):
            f = self.get_marginal_score(mode, chosed_actions[0:n], K, alpha)
            temp = np.argmax(np.mean(f, 0))
            chosed_actions[n] = self.actions[temp]
            self.actions = self.delete_actions(temp)

        return chosed_actions

    def get_random_action(self, N):
        random_index = np.random.permutation(xrange(self.actions.shape[0]))
        index = random_index[0:N]
        action = self.actions[index]
        self.actions = self.delete_actions(index)
        return action

    def load(self, file_history, file_training=None, file_predictor=None):
        self.history.load(file_history)

        if file_training is None:
            N = self.history.total_num_search
            X = self.test.X[self.history.chosed_actions[0:N], :]
            t = self.history.fx[0:N]
            self.training = variable(X=X, t=t)
        else:
            self.training = variable()
            self.training.load(file_training)

        if file_predictor is not None:
            with open(file_predictor) as f:
                self.predictor = pickle.load(f)

    def export_predictor(self):
        return self.predictor

    def export_training(self):
        return self.training

    def export_history(self):
        return self.history

    def _set_predictor(self, predictor=None):
        if predictor is None:
            predictor = self.predictor
        return predictor

    def _init_predictor(self, is_rand_expans, predictor=None):
        self.predictor = self._set_predictor(predictor)
        if self.predictor is None:
            if is_rand_expans:
                self.predictor = blm_predictor(self.config)
            else:
                self.predictor = gp_predictor(self.config)

        return self.predictor

    def _set_training(self, training=None):
        if training is None:
            training = self.training
        return training

    def _set_unchosed_actions(self, actions=None):
        if actions is None:
            actions = self.actions
        return actions

    def _set_test(self, test_X):
        if isinstance(test_X, np.ndarray):
            test = variable(X=test_X)
        elif isinstance(test_X, variable):
            test = test_X
        else:
            raise TypeError('The type of test_X must \
                             take ndarray or combo.variable')
        return test

    def _set_config(self, config=None):
        if config is None:
            config = combo.misc.set_config()
        return config

[((21, 23, 21, 57), 'numpy.arange', 'np.arange', ({(21, 33, 21, 34): '0', (21, 36, 21, 56): 'self.test.X.shape[0]'}, {}), '(0, self.test.X.shape[0])', True, 'import numpy as np\n'), ((22, 23, 22, 32), 'results.history', 'history', ({}, {}), '()', False, 'from results import history\n'), ((27, 8, 27, 33), 'numpy.random.seed', 'np.random.seed', ({(27, 23, 27, 32): 'self.seed'}, {}), '(self.seed)', True, 'import numpy as np\n'), ((31, 15, 31, 40), 'numpy.delete', 'np.delete', ({(31, 25, 31, 32): 'actions', (31, 34, 31, 39): 'index'}, {}), '(actions, index)', True, 'import numpy as np\n'), ((74, 15, 74, 42), 'copy.deepcopy', 'copy.deepcopy', ({(74, 29, 74, 41): 'self.history'}, {}), '(self.history)', False, 'import copy\n'), ((128, 15, 128, 42), 'copy.deepcopy', 'copy.deepcopy', ({(128, 29, 128, 41): 'self.history'}, {}), '(self.history)', False, 'import copy\n'), ((173, 15, 173, 27), 'numpy.argmax', 'np.argmax', ({(173, 25, 173, 26): 'f'}, {}), '(f)', True, 'import numpy as np\n'), ((177, 25, 177, 47), 'numpy.zeros', 'np.zeros', (), '', True, 'import numpy as np\n'), ((153, 24, 153, 53), 'copy.deepcopy', 'copy.deepcopy', ({(153, 38, 153, 52): 'self.predictor'}, {}), '(self.predictor)', False, 'import copy\n'), ((154, 20, 154, 48), 'copy.deepcopy', 'copy.deepcopy', ({(154, 34, 154, 47): 'self.training'}, {}), '(self.training)', False, 'import copy\n'), ((182, 29, 182, 42), 'numpy.mean', 'np.mean', ({(182, 37, 182, 38): 'f', (182, 40, 182, 41): '0'}, {}), '(f, 0)', True, 'import numpy as np\n'), ((209, 33, 209, 47), 'cPickle.load', 'pickle.load', ({(209, 45, 209, 46): 'f'}, {}), '(f)', True, 'import cPickle as pickle\n')]