import numpy as np
from rl.hyperoptimizer.base_hyperoptimizer import HyperOptimizer


class RandomSearch(HyperOptimizer):

    '''
    Random Search by sampling on hysphere around a search path
    algo:
    1. init x a random position in space
    2. until termination (max_eval or fitness, e.g. solved all), do:
        2.1 sample new pos some radius away: next_x = x + r
        2.2 if f(next_x) > f(x) then set x = next_x

    Extra search memory units:
    - search_path
    - best_point

    to save an restore experiment:
    - all searched points
    - all updated points and their fitness score, just use performance_score, current pointer of x as last
    careful with param_search_list and next_trial_num
    need to get stuff from self.experiment_data
    path:
    '''

    def set_keys(self, **kwargs):
        self.REQUIRED_GLOBAL_VARS = [
            'experiment_spec',
            'experiment_id_override',
            'times',
            'max_evals'
        ]
        super(RandomSearch, self).set_keys(**kwargs)

    # calculate the constant radius needed to traverse unit cube
    def cube_traversal_radius(self):
        traversal_diameter = 1/np.power(self.max_evals,
                                        1/self.search_dim)
        traversal_radius = traversal_diameter/2
        return traversal_radius

    def decay_radius(self):
        '''
        future implementation, start of half cube for diameter
        (so 1/4 for radius), then decay
        '''
        return

    @classmethod
    def sample_hypersphere(cls, dim, r=1):
        '''Marsaglia algo for sampling uniformly on a hypersphere'''
        v = np.random.randn(dim)
        v = v * r / np.linalg.norm(v)
        return v

    def sample_cube(self):
        return np.random.rand(self.search_dim)

    def sample_r(self):
        return self.sample_hypersphere(
            self.search_dim, self.search_radius)

    # biject [0, 1] to [x_min, x_max]
    @classmethod
    def biject_continuous(cls, norm_val, x_min, x_max):
        return norm_val*(x_max - x_min) + x_min

    # biject [0, 1] to x_list = [a, b, c, ...] by binning
    def biject_discrete(self, norm_val, x_list):
        list_len = len(x_list)
        inds = np.arange(list_len)
        cont_val = self.biject_continuous(norm_val, 0, list_len)
        ind = np.digitize(cont_val, inds) - 1
        return x_list[ind]

    def biject_dim(self, norm_val, dim_spec):
        if isinstance(dim_spec, list):  # discrete
            return self.biject_discrete(norm_val, dim_spec)
        else:  # cont
            return self.biject_continuous(
                norm_val, dim_spec['min'], dim_spec['max'])
        return

    # biject a vector on unit cube into a param in param_space
    def biject_param(self, v):
        param = {}
        for i, param_key in enumerate(self.param_range_keys):
            dim_spec = self.param_range[param_key]
            param[param_key] = self.biject_dim(v[i], dim_spec)
        return param

    def init_search(self):
        '''
        Initialize the random search internal variables
        '''
        self.num_of_trials = self.max_evals
        self.search_dim = len(self.param_range_keys)
        self.search_radius = self.cube_traversal_radius()
        self.search_path = []
        self.best_point = {
            'trial_num': None,
            'param': None,
            'x': self.sample_cube(),
            'fitness_score': float('-inf'),
        }

    def search(self):
        '''
        algo step 2.1 sample new pos some radius away: next_x = x + r
        update search_path and param_search_list
        '''
        next_x = self.best_point['x'] + self.sample_r()
        next_param = self.biject_param(next_x)
        self.search_path.append(next_x)
        self.param_search_list.append(next_param)
        print('next param')
        print('next param')
        print('next param')
        print(next_x)
        print(next_param)

    def update_search(self):
        '''
        algo step 2.2 if f(next_x) > f(x) then set x = next_x
        invoked right after the latest run_trial()
        update self.best_point
        '''
        if self.next_trial_num < self.PARALLEL_PROCESS_NUM:  # first runs yet
            return
        assert len(self.experiment_data) > 0, \
            'self.experiment_data must not be empty for update_search'
        x = self.search_path[-1]
        trial_data = self.experiment_data[-1]
        trial_num, param, fitness_score = self.get_fitness(trial_data)

        if fitness_score > self.best_point['fitness_score']:
            self.best_point = {
                'trial_num': trial_num,
                'param': param,
                'x': x,
                'fitness_score': fitness_score,
            }

    def satisfy_fitness(self):
        '''use performance score, solved ratio, solved mean reward'''
        # ideal_fitness_score = util.PROBLEMS:
        # SOLVED_MEAN_REWARD/MAX_EPISODES/2
        return False

    def to_terminate(self):
        return (self.next_trial_num >= self.max_evals or
                self.satisfy_fitness())

@@ -77,16 +77,45 @@
 timizer%0A
+from rl.util import PROBLEMS%0A
 %0A%0Aclass 
@@ -524,26 +524,16 @@
     
-to 
 save 
-an rest
+f
 or
-e
  exp
@@ -543,254 +543,75 @@
 ment
+ resume
 :%0A    -
- all
  search
-ed points%0A    - all updated points and their fitness score, just use performance_score, current pointer of x as last%0A    careful with param_search_list and next_trial_num%0A    need to get stuff from self.experiment_data%0A    path:
+_path%0A    - best_point%0A    - param_search_list
 %0A   
@@ -2214,16 +2214,72 @@
 t%5Bind%5D%0A%0A
+    # biject one dimension: %5B0, 1%5D to a param_range val%0A
     def 
@@ -3339,32 +3339,205 @@
 inf'),%0A        %7D
+%0A        problem = PROBLEMS.get(self.experiment_spec%5B'problem'%5D)%0A        self.ideal_fitness_score = 0.5 * %5C%0A            problem%5B'SOLVED_MEAN_REWARD'%5D/problem%5B'MAX_EPISODES'%5D
 %0A%0A    def search
@@ -3883,140 +3883,8 @@
 ram)
-%0A        print('next param')%0A        print('next param')%0A        print('next param')%0A        print(next_x)%0A        print(next_param)
 %0A%0A  
@@ -4743,177 +4743,240 @@
  '''
-use performance score, solved ratio, solved mean reward'''%0A        # ideal_fitness_score = util.PROBLEMS:%0A        # SOLVED_MEAN_REWARD/MAX_EPISODES/2%0A        return Fals
+%0A        break on the first strong solution%0A        '''%0A        if self.next_trial_num %3C self.PARALLEL_PROCESS_NUM:%0A            return False%0A        else:%0A            return self.best_point%5B'fitness_score'%5D %3E self.ideal_fitness_scor
 e%0A%0A 

add satisfy_fitness function