# -*- coding: utf-8 -*-
# @Author: Aliaksei Halachkin
# @Date:   2016-11-15 18:48:22
# @Last Modified by:   Aliaksei Halachkin
# @Last Modified time: 2017-04-19 15:07:14


import numpy as np
import json
from exceptions import Exception
import warnings
import traceback
import time     
import abc
from itertools import product
import random


class OptimizerABC:
    """
    Function optimizer abstract base class 
    attributes:      
        xopt - optimal x
        fopt - func(xopt)
        n_evals - number of evaluation
        scores - computed function values
    """
    __metaclass__ = abc.ABCMeta

    def __init__(self):
        self._xopt = None
        self._fopt = None
        self.xs_ = []
        self.fs_ = []
        self._n_evals = []
        self._total_evals = 0
        self._func = None

    def _check_boundaires(self, x):
        if hasattr(self, 'xdomain'):
            if not self.xdomain is None:
                for idx in range(len(x)):
                    if x[idx] < self.xdomain[idx][0]:
                        x[idx] = self.xdomain[idx][0]
                    elif x[idx] > self.xdomain[idx][1]:
                        x[idx] = self.xdomain[idx][1]
                    else:
                        continue

    @property
    def xopt(self):
        return self._xopt

    @xopt.setter
    def xopt(self, value):
        self.xs_.append(value)
        self._xopt = value

    @property
    def fopt(self):
        return self._fopt

    @fopt.setter
    def fopt(self, value):
        self._n_evals.append(self._total_evals)
        self.fs_.append(value)
        self._fopt = value

    @property
    def n_evals(self):
        return self._n_evals

    @property
    def total_evals(self):
        return self._total_evals


    def func(self, *args, **kwargs):
        self._total_evals += 1
        return self._func(*args, **kwargs)

    def find_opt(self):
        self._fopt = np.min(self.fs_)
        self._xopt = self.xs_[np.argmin(self.fs_)]

    @abc.abstractmethod
    def run(self, func):
        self._func = func

    def json_dump(self, path):
        attributes = dict()

        for k in self.__dict__:
            if isinstance(self.__dict__[k], np.ndarray):
                jsonified = self.__dict__[k].tolist()

            elif isinstance(self.__dict__[k], list):
                jsonified = []
                for item in self.__dict__[k]:
                    if isinstance(item, np.ndarray):
                        jsonified.append(item.tolist())
                    else:
                        jsonified.append(item)
            else:
                jsonified = self.__dict__[k]
            try:
                json.dumps(jsonified)
            except:
                warnings.warn("cant jsonify " + str(k) + " attribute")
                jsonified = str(jsonified)
            finally:
                attributes[k] = jsonified

        attributes["object"] = type(self).__name__
        json.dump(attributes, open(path, mode="w"), indent=4)

                        
        
class BackwardElimination(OptimizerABC):
    def __init__(self, n_features, suppress_printing=True):
        OptimizerABC.__init__(self)
        self.xopt = np.ones(n_features)
        self.n_features = n_features
        self.n_selected_ = n_features
        self.suppress_printing = suppress_printing


    def _gen_xs(self):
        leaved = np.arange(self.n_features)[self.xopt > 0.0]
        xs = []
        for ind in leaved:
            x = self.xopt.copy()
            x[ind] = 0.0
            xs.append(x)
        return xs

    def run(self, func):
        self._func = func
        self.fopt = self.func(self.xopt)
        try:
            while(self.n_selected_ > 1):
                xs = self._gen_xs()
                scores = map(self.func, xs)
                if np.min(scores) < self.fopt:
                    self.xopt = xs[np.argmin(scores)]
                    self.fopt = np.min(scores)
                    self.n_selected_ -= 1
                    if not self.suppress_printing:
                        print self.fopt, self.xopt
                else:
                    break

        except KeyboardInterrupt:
            print "Stopped by user"

        except Exception as e:
            # print e.message
            traceback.print_exc()

        finally:
            pass


class GridSearch(OptimizerABC):
    def __init__(self, x_grid, suppress_printing=True):

        """
        x_grid [2d array like] : args to to test, will try every combination
                                 every row define search domine
                                 for each variable

        """
        OptimizerABC.__init__(self)
        self.x_grid = x_grid
        self.suppress_printing = suppress_printing
        self._x_to_test = [np.array(comb) for comb in product(*x_grid)]

        if not self.suppress_printing:
            print "Will test " + str(len(self._x_to_test)) + " combinations"
     

    def run(self, func):
        self._func = func

        try:
            start = time.time()
            for x in self._x_to_test:
                #see OptimizerABC for the details
                self.xopt = x
                self.fopt = self.func(self.xopt)

        except KeyboardInterrupt:
            print "Stopped by user"

        except Exception as e:
            print e.message

        finally:    

            end = time.time()        
            #find best combination
            self.find_opt()

            if not self.suppress_printing:
                msg = "Best score: " + str(self.fopt) 
                msg += ", with x " + str(self.xopt)
                msg +=  "time elapsed: " + str(end - start) + " s"
                print msg

            return self.xopt


#TODO: Make line search function support
class CoordinateDescent(OptimizerABC):
    def __init__(self, x_grid,
                       max_iter=1, 
                       order=None,
                       suppress_printing=True):
        OptimizerABC.__init__(self)

        self.x_grid = np.array(x_grid)
        if order is None:
            self.order = np.arange(len(x_grid))
        else:
            self.order = order

        self.max_iter = max_iter
        #init x
        self.xopt = np.array([x_grid[idx][0] for idx in self.order])
        self.suppress_printing = suppress_printing

    def run(self, func):
        self._func = func
        #init score
        self.fopt = self.func(self.xopt)
        try:
            for n in range(self.max_iter):
                for idx in self.order:
                    x_to_test = np.array([self.xopt]*len(self.x_grid[idx]))
                    x_to_test[:, idx] = self.x_grid[idx]
                    scores = map(self.func, x_to_test)
                    self.xopt = x_to_test[np.argmin(scores)]
                    self.fopt = np.min(scores)

        except KeyboardInterrupt:
            print "Stopped by user"

        except Exception as e:
            print e.message

        finally:    
            if not self.suppress_printing:
                msg = "Best score: " + str(self.fopt) 
                msg += ", with x " + str(self.xopt)
                print msg




def binary_search(func, a, b, tol):
    if (a > b):
        a, b = b, a
    mid = (b - a) / 2.0 + a
    if (b - a) < tol:
        return mid, func(mid)
    mid_left = (mid - a) / 2.0 + a
    mid_right =  (b - mid) / 2.0 + mid
    # print a, b, mid_left, mid_right
    fleft = func(mid_left )
    fright = func(mid_right )

    if fleft < fright:
        return binary_search(func, a, mid, tol)
    else:
        return binary_search(func, mid, b, tol)


class CoordinateDescentDichotomy(OptimizerABC):
    def __init__(self, xdomain,
                       max_iter=1, 
                       dichotomy_tol = 1e-6,
                       order=None,
                       suppress_printing=True):
        OptimizerABC.__init__(self)

        self.xdomain = xdomain
        self.max_iter = max_iter
        self.dichotomy_tol = dichotomy_tol
        if order is None:
            self.order = np.arange(len(xdomain))
        else:
            self.order = order
        #init x
        self.xopt = np.array([(b - a) / 2.0 + a for a, b in xdomain])
        self.suppress_printing = suppress_printing

    def _1d_func(self, x, idx, val):
        x_ = x.copy()
        x[idx] = val
        return self.func(x)

    def run(self, func):
        self._func = func
        #init score
        self.fopt = self.func(self.xopt)
        try:
            for n in range(self.max_iter):
                for idx in self.order:
                    f = lambda val : self._1d_func(self.xopt, idx, val)
                    new_x, fval = binary_search(f, self.xdomain[idx][0],
                                                   self.xdomain[idx][1],
                                                   self.dichotomy_tol)

                    self.xopt[idx] = new_x
                    self.fopt = fval

        except KeyboardInterrupt:
            print "Stopped by user"

        except Exception as e:
            print e.message

        finally:    
            if not self.suppress_printing:
                msg = "Best score: " + str(self.fopt) 
                msg += ", with x " + str(self.xopt)
                print msg



class HillClimb(OptimizerABC):

    def __init__(self, x_init, 
                 dx=0.1, 
                 xdomain=None,
                 max_iter=100,
                 ftol=1e-6,
                 xtol=1e-9, 
                 push_when_converged=False,
                 push_when_stucked=False,
                 push_factor=1.0,
                 n_candidates=3,
                 suppress_printing=True):


        OptimizerABC.__init__(self)
        self.max_iter = max_iter
        if hasattr(dx, '__len__') and hasattr(dx, '__getitem__'):
            if len(dx) != len(x_init):
                raise ValueError("dx should be a scalar or a list with \
                                 length equal to number of dimensions") 
            self._dx_is_arr = True
        else:
            self._dx_is_arr = False

        if hasattr(xdomain, '__len__') and hasattr(xdomain, '__getitem__'):
            if len(xdomain) != len(x_init):
                raise ValueError("xdomain should be a iterable of tuples with \
                                  length equal to number of dimensions. \
                                  Each tuple defines search domain (lower and \
                                  upper bound)")    
        self.xdomain = xdomain

        self.xopt = np.array(x_init)
        self.dx = dx
        self.suppress_printing = suppress_printing 
        self.ftol = ftol
        self.xtol = xtol
        self.push_when_converged = push_when_converged
        self.push_factor = push_factor
        self.n_candidates = n_candidates

    def _gen_neighbors(self, n_neighbors):
        x_neighbors = []
        ind = range(len(self.xopt))
        random.shuffle(ind)
        for idx in ind[:n_neighbors]:
            x_plus = self.xopt.copy()
            x_plus[idx] += self.dx[idx] if self._dx_is_arr else self.dx
            if x_plus[idx] < 0.0: x_plus[idx] = 0.0

            x_minus = self.xopt.copy()
            x_minus[idx] -= self.dx[idx] if self._dx_is_arr else self.dx
            if x_minus[idx] < 0.0: x_minus[idx] = 0.0

            self._check_boundaires(x_plus)
            self._check_boundaires(x_minus)

            x_neighbors += [x_plus, x_minus]

        return x_neighbors
       
    def _push(self):
        idx = np.random.randint(0, len(self.xopt))
        self.xopt[idx] += \
                      random.random() * 2 * self.push_factor - self.push_factor
        if self.xopt[idx] < 0 : 
            self.xopt[idx] = 0.0

    def run(self, func):
        self._func = func
        self.fopt = self.func(self.xopt)
        try:
            for n in range(self.max_iter):
                #generate neighbors
                x_neighbors = self._gen_neighbors(self.n_candidates)
                #evaluate neighbors
                scores = map(self.func, x_neighbors)
                #pick best
                idx = np.argmin(scores)
                self.xopt = x_neighbors[idx]
                self.fopt = scores[idx]

                if not self.suppress_printing:
                    print "current score: " + str(self.fopt),
                    print ", with weights" + str(self.xopt)

                #check for convergence
                if n > 1:
                    delta  = abs(self.fs_[-1] - self.fs_[-2])
                    if delta < self.ftol:
                        if not self.suppress_printing:
                            print "reached convergence criteria"
                        if self.push_when_converged:
                            self._push()
                        else:
                            break

                #check if similar point was already visited
                if n > 1:
                    delta = np.min([np.sum(np.abs(self.xopt - x)) 
                                            for x in self.xs_[:-1]])
                    if delta < self.xtol:
                        if not self.suppress_printing:
                            print "stucked"
                        if self.push_when_converged:
                            self._push()
                        else:
                            break 



        except KeyboardInterrupt:
            print "Stopped by user"

        except Exception as e:
            # print e.message
            traceback.print_exc()

        finally:
            self.find_opt()


class FirstChoiceHillClimb(OptimizerABC):
    def __init__(self, x_init, 
                       mutation_k = 1.0,
                       xdomain=None,
                       max_iter=100,
                       simulated_annealing=False,
                       temp_schedule=None,
                       suppress_printing=True):

        OptimizerABC.__init__(self)
        self.xdomain = xdomain
        self.xopt = np.array(x_init)
        self.mutation_k = mutation_k
        self.max_iter = max_iter
        self.suppress_printing = suppress_printing 

        self.simulated_annealing = simulated_annealing
        if simulated_annealing and temp_schedule is None:
            self.temp_schedule = [1/float(i) for i in range(5, self.max_iter+5)]
        else:
            self.temp_schedule = temp_schedule

        self.mutations_stats = np.zeros(10)


    def _mutate(self):
        k = self.mutation_k
        roll = np.random.rand()
        x = self.xopt.copy()

        if roll < 0.1:  
            x += np.random.rand(len(x)) * k / 5.0 - k / 10.0
            self.mutations_stats[0] += 1

        elif roll < 0.2:
            x += np.random.rand(len(x)) * k  -  k / 2.0
            self.mutations_stats[1] += 1

        elif roll < 0.3:
            x += np.random.rand(len(x)) * 5.0 * k  -   2.5 * k      
            self.mutations_stats[2] += 1

        elif roll < 0.4:
            idx = np.random.randint(0, len(x))
            x[idx] *= 2.0
            self.mutations_stats[3] += 1     

        elif roll < 0.5:
            idx = np.random.randint(0, len(x))
            x[idx] /= 2.0
            self.mutations_stats[4] += 1             

        elif roll < 0.6:
            ind = np.array(range(len(x)))
            random.shuffle(ind)
            n = random.randint(0, len(x))
            x[ind[:n]] += np.random.rand(n) * k / 5.0 - k / 10.0 
            self.mutations_stats[5] += 1

        elif roll < 0.7:
            ind = np.array(range(len(x)))
            random.shuffle(ind)
            n = random.randint(0, len(x))
            x[ind[:n]] += np.random.rand(n) * k  -  k / 2.0
            self.mutations_stats[6] += 1

        elif roll < 0.8:
            ind = np.array(range(len(x)))
            random.shuffle(ind)
            n = random.randint(0, len(x))
            x[ind[:n]] += np.random.rand(n) * k  -  k / 2.0
            self.mutations_stats[7] += 1

        elif roll < 0.9:
            ind = np.array(range(len(x)))
            random.shuffle(ind)
            n = random.randint(0, len(x))
            x[ind[:n]]  += np.random.rand(n) * 5.0 * k  -   2.5 * k 
            self.mutations_stats[8] += 1   

        else:
            idx = np.random.randint(0, len(x))
            x[idx] *= 2.0 
            self.mutations_stats[9] += 1


        self._check_boundaires(x)
        return x
           

    def run(self, func):
        self._func = func
        self.fopt = self.func(self.xopt)
        try:
            for n in range(self.max_iter):
                #generate new x
                x  = self._mutate()
                #evaluate 
                score = self.func(x)
                #if is better accept it
                if score < self.fopt or \
                    (self.simulated_annealing and 
                     np.random.rand() < self.temp_schedule[n]):

                    self.fopt = score
                    self.xopt = x

                    if not self.suppress_printing:
                        print "current score: " + str(self.fopt),
                        print ", with weights" + str(self.xopt)

        except KeyboardInterrupt:
            print "Stopped by user"

        except Exception as e:
            # print e.message
            traceback.print_exc()

        finally:
            self.find_opt()




if __name__ == '__main__':

    import matplotlib.pyplot as plt

    #Eggholder function
    # real min
    # print f([512, 404.2319])
    # -959.640662711
    def eggholder(x):
        """
        Eggholder function
        real min:
        >>print f([512, 404.2319])
        >>-959.640662711
        """
        a = -(x[1] + 47)*np.sin(np.sqrt(np.abs(x[0]/2.0 + (x[1] + 47))))  
        b = -x[0]*np.sin(np.sqrt(np.abs(x[0]-(x[1]+47))))
        return a + b


    def rosenbrock(x):
        """
        Rosenbrock function
        real min at f([1,1,1, ... 1]) = 0, len(x) > 1
        """
        return np.sum(100. * (x[1:] - x[:-1]**2)**2 + (x[:-1]-1)**2)

    # print rosenbrock(np.array([1,1,1,1]))

    def grid_search_example(func):
        x_grid = [ np.arange(-512, 513, 4), 
                   np.arange(-512, 513, 4)]

        gS =  GridSearch(x_grid,  suppress_printing=True)
        gS.run(func)
        print "Grid Search"
        print "xopt: " + str(gS.xopt)
        print "fopt: " + str(gS.fopt)
        print "func evals: " + str(gS.total_evals)

    def coordinate_descent_example(func):
        x_grid = [ np.arange(-512, 513, 4), 
                   np.arange(-512, 513, 4)]

        cD = CoordinateDescent(x_grid)
        cD.run(func)
        print "Coordinate Descent"
        print "xopt: " + str(cD.xopt)
        print "fopt: " + str(cD.fopt)
        print "func evals: " + str(cD.total_evals)


    def hill_climb_example(func):
        x_init = np.random.rand(2)*2
        x_init = np.array([3,-1.5])
        hC =  HillClimb(x_init, 
                       xdomain=[(-512,512),(-512,512)],
                       dx=0.005, 
                       max_iter=1000,
                       tol=1e-2, 
                       push_when_converged=True,
                       push_factor=0.2,
                       n_candidates=2,
                       suppress_printing=True)

        hC.run(func)
        print "Hill Climb"
        print "x init: " + str(x_init)
        print "xopt: " + str(hC.xopt)
        print "fopt: " + str(hC.fopt)
        print "func evals: " + str(hC.total_evals)

        plt.plot(hC.n_evals, hC.fs_); plt.grid();
        plt.xlabel("n_evals"); plt.ylabel("func values")
        plt.show()
           
    def first_choice_hill_climb_example(func):
        x_init = np.random.rand(10)*3
        fcHc = FirstChoiceHillClimb(x_init,  
                                     mutation_k=0.1,
                                     xdomain=None,
                                     max_iter=5000,
                                     simulated_annealing=False,
                                     temp_schedule=None,
                                     suppress_printing=True)

        fcHc.run(func)
        print "First Choice Hill Climb"
        print "x init: " + str(x_init)
        print "xopt: " + str(fcHc.xopt)
        print "fopt: " + str(fcHc.fopt)
        print "func evals: " + str(fcHc.total_evals)
        print len(fcHc.xs_)
        print fcHc.mutations_stats
        plt.plot(fcHc.n_evals, fcHc.fs_); plt.grid();
        plt.xlabel("n_evals"); plt.ylabel("func values")
        plt.show()
           

    def coordinate_descent_dichotomy_example(func):
        xdomain= [ [-512, 500], [-512, 500]]

        cDD = CoordinateDescentDichotomy(xdomain, dichotomy_tol=1e-6, max_iter=5)
        cDD.run(func)
        print "Coordinate Descent Dichotomy"
        print "xopt: " + str(cDD.xopt)
        print "fopt: " + str(cDD.fopt)
        print "func evals: " + str(cDD.total_evals)


    # grid_search_example(f)
    # coordinate_descent_example(f)
    # hill_climb_example(rosenbrock)


     
    # first_choice_hill_climb_example(rosenbrock)

    # print binary_search(lambda x: np.cos(x), 0, 3*np.pi / 2, 0.000001)

    coordinate_descent_dichotomy_example(eggholder)



    # x_init = np.random.rand(2)*500
    # x_init = [0.0, 0.0]
    # sA = SimulatedAnnealing(x_init, 
    #                         max_iter=1000,
    #                         tol=1e6,
    #                         suppress_printing=True)

    # sA.run(f)

    # print sA.xopt
    # print sA.fopt
    # plt.plot(sA.selected_scores); plt.grid()
    # plt.show()