import numpy as np
from matplotlib.lines import Line2D
from matplotlib import animation
import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib.colors import LogNorm

class OptimizationVisualizer(object):
    def __init__(self,name='himmelblau',jac=False):
        self.jac=jac
        self.name = name
        options = ['himmelblau','rosenbrock']
        if name not in options:
            raise ValueError('Unkown name {}. Possible options {}'.format(name,options))
        self.samples = []
        self.values = []
        self.num_evals = 0
    def objective(self,x):
        if self.name == 'rosenbrock': 
            return scipy.optimize.rosen(x)
        if self.name == 'himmelblau':
            X = x[0]
            Y = x[1]
            a = X*X + Y - 11
            b = X + Y*Y - 7
            return a*a + b*b
    def jacobian(self,x):
        if self.name == 'rosenbrock': 
            return scipy.optimize.rosen_der(x)
        if self.name == 'himmelblau':
            X = x[0]
            Y = x[1]
            a = X*X + Y - 11
            b = X + Y*Y - 7
            da = np.array([2*X,1])
            db = np.array([1,2*Y])
            return 2*a*da + 2*b*db
        
    def __call__(self,x):
        self.samples.append(x.tolist())
        val = self.objective(x)
        self.values.append(val)
        self.num_evals += 1
        if not self.jac: return val
        jac = self.jacobian(x)
        return val,jac  
    
    def init(self): return self.line,self.scatter

    def update_annotation(self,num_evals):
        min_val = min(self.values[:num_evals])
        if min_val < 0.1:
            template = 'evaluations: {}\nminimum: {:.3e}'
        else:
            template = 'evaluations: {}\nminimum: {:.2f}'
            
        self.annotation.set_text(template.format(
            num_evals,min_val))
        
    def update_simplex(self,step):
        num_dim = len(self.samples[0])
        #number of evals: (total evals - left samples) + simplex + trial
        num_current_evals = self.num_evals - len(self.samples)+num_dim+2
        self.update_annotation(num_current_evals)
        if len(self.samples) < num_dim + 2: return self.line,self.scatter,self.annotation
        simplex = np.array(self.samples[:num_dim+1])
        simp_values = np.array(self.values[:num_dim+1])
        idx_sorted = np.argsort(simp_values)
        simplex = simplex[idx_sorted]
        simp_values = simp_values[idx_sorted]
        #show simplex
        
        line_points = np.append(simplex,[simplex[0]],axis=0)
        self.line.set_data(line_points[:,0],line_points[:,1])
        #points shown in the graph
        points = simplex
        
        idx_worst = idx_sorted[-1]
        idx_best = idx_sorted[0]
        
        #reflection
        sample1 = self.samples[num_dim+1]
        val1 = self.values[num_dim+1]
        points = np.append(points,[sample1],axis=0)
        self.scatter.set_offsets(points)
        
        if simp_values[0] <= val1 < simp_values[-2]:
            #delete worst sample
            del self.samples[idx_worst]                
            del self.values[idx_worst]
            return self.line,self.scatter,self.annotation
        
        #else a second sample is taken, either expansion or reflection
        sample2 = self.samples[num_dim+2]
        val2 = self.values[num_dim+2]
        points = np.append(points,[sample2],axis=0)
        self.scatter.set_offsets(points)
        num_current_evals += 1
        self.update_annotation(num_current_evals)
        #expansion
        if (val1<simp_values[0]):
            if val2 < val1:
                del self.samples[num_dim+1]                
                del self.values[num_dim+1]
            else:
                del self.samples[num_dim+2]
                del self.values[num_dim+2]
            del self.samples[idx_worst]                
            del self.values[idx_worst]
            return self.line,self.scatter,self.annotation
        
        #contraction
        if (val1>=simp_values[-2]):
            if val2 < simp_values[-1]:
                del self.samples[num_dim+1]                
                del self.values[num_dim+1]
                del self.samples[idx_worst]                
                del self.values[idx_worst]
            else: #shrink, i.e. delete all but best point
                num_current_evals += self.num_dim
                self.update_annotation(num_current_evals)
                for idx in range(num_dim+2,0,-1):
                    if idx != idx_best: del self.samples[idx]                
            return self.line,self.scatter,self.annotation

    def update_population(self,step):
        pop = self.population_size
        samples = np.array(self.samples[step*pop:(step+1)*pop])
        self.scatter.set_offsets(samples)
        self.update_annotation((step+1)*pop)
        return self.line,self.scatter,self.annotation
        
    def update(self,step):

        if self.show_simplex:
            return self.update_simplex(step)
                                    
        if self.population_size > 1:
            return self.update_population(step)
            
        samples = np.array(self.samples[:step+1])
        values = np.array(self.values[:step+1])
        self.line.set_data(samples[:,0],samples[:,1])
        self.scatter.set_offsets(samples)
        sizes = np.array(range(len(values)))/(len(values))
        sizes[np.where(sizes<0.2)] = 0.2
        self.scatter.set_array(sizes)
        if len(sizes) > 1:
            self.scatter.set_clim(np.min(sizes),np.max(sizes))
        self.update_annotation(step+1)
                             
        return self.line,self.scatter,self.annotation
        
    def animate(self,
                x = np.arange(-8,8,0.1), 
                y = np.arange(-8,8,0.1),
                show_simplex=False,
                population_size=1,
    ):

        self.show_simplex = show_simplex
        self.population_size = population_size
        
        X, Y = np.meshgrid(x, y)
        Z = np.array([[self.objective(np.array([a,b])) for a in x] for b in y])
        fig, ax = plt.subplots()
        self.scatter = ax.scatter(x,y,marker='o',vmin=0, vmax=1,
                                    cmap=cm.brg,zorder=10, edgecolor="k")
        self.line = Line2D([],[], marker=None, ls='-', lw=0.8, ms=4, color='black',
                           alpha=0.75)
        self.annotation = plt.annotate('evaluations: 0\n ', xy=(-7,6))

        ax.add_line(self.line)
        cs = ax.contour(X,Y,Z,np.logspace(0,10,21,base=10),norm = LogNorm(),
                        cmap = cm.jet, alpha=0.5)
        ax.clabel(cs,levels=[10,100,1000],fmt='%1.0f')
        
        anim = animation.FuncAnimation(fig, self.update,
                               frames=range(int(len(self.samples)/population_size)), 
                               interval=1000, init_func=self.init, blit=True)
        return anim
