import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation, gridspec
from matplotlib.lines import Line2D
from IPython.display import HTML, clear_output
import scipy
from scipy.special import erf, erfc
from .gaussian_process import GaussianProcess


class BO_visualization(object):

    def __init__(self,use_derivative=False,acquisition_function='expected improvement'):
        """
        Input:
           use_derivative: If true, derivative observations are used to minimize the function
           acquisition_function: Can be one of 'expected improvement', 
                                 'lower confidence bound', or 
                                 'probability of improvement'
        """
        self.use_derivative = use_derivative
        self.acquisition_function = acquisition_function
        
        self.gp = GaussianProcess(sigma=1,l=3)
        self.f = lambda x: 0.5*np.sin(0.8*x+0.5)+0.015*x**2-0.5
        self.fp = lambda x: 0.8*0.5*np.cos(0.8*x+0.5)+ 2*0.015*x
        self.Xs = np.arange(-10,10,0.05)
    
        #current/initial sampling point
        self.sample = -5.2
        self.fmin = self.f(self.sample)
        self.samples = []
        self.fig = plt.figure(figsize=(10, 4),clear=True) 
        gs = gridspec.GridSpec(2, 1, height_ratios=[2, 1]) 
        self.ax1 = self.fig.add_subplot(gs[0])
        plt.setp(self.ax1.get_xticklabels(), visible=False)
        self.ax2 = self.fig.add_subplot(gs[1], sharex=self.ax1)
    
        self.ax1.plot(self.Xs,self.f(self.Xs),color='darkred',label='function $f$')
        self.ax1.plot(self.Xs[np.argmin(self.f(self.Xs))],np.min(self.f(self.Xs)),marker='x',ls=None,color='darkred',markersize=10.0,label='global minimum')

        self.ax1_lines = [
            Line2D([],[],marker='o',ls='None',label='observations of $f$'),
            Line2D([],[],ls='--',linewidth=2.0,color='black',label='mean'),
            Line2D([],[],ls='-',linewidth=2.0,color='black',label='mean $\pm$ standard deviation',alpha=0.4),
            Line2D([],[],ls='-',linewidth=2.0,color='black',alpha=0.4)]
        self.ax2_lines = [
            Line2D([],[],ls='-',color='blue',label=acquisition_function),
            Line2D([],[],marker='x',ls='None',color='blue',label='maximum')]
        for l in self.ax1_lines: self.ax1.add_line(l)
        for l in self.ax2_lines: self.ax2.add_line(l)
        
    def init(self):
        self.ax1.set_xlim([-10,10])
        self.ax1.set_ylim([-2,2])
        self.ax1.legend(loc='lower left')
        self.ax2.legend(loc='lower left')
        self.ax2.set_xlim([-10,10])
        self.ax2.set_ylim([1e-4,10])
        self.ax2.set_yscale('log')
        self.ax2.set_xlabel('parameter value', fontsize=12)
        return self.ax1_lines+self.ax2_lines
        
    def update(self,step):
        self.gp.add(self.sample,self.f(self.sample),0)    
        #uncomment, if you want to see whats happens if you supply derivative information
        if self.use_derivative: self.gp.add(self.sample,self.fp(self.sample),1)  
        
        try: self.gp.prepare()
        except np.linalg.LinAlgError: return self.ax1_lines+self.ax2_lines
        
        predictions = np.array([self.gp.predict(x) for x in self.Xs])
        mu = predictions[:,0]
        sdev = predictions[:,1]
        #avoid division by zero
        sdev[np.where(sdev<1e-12)] = 1e-12
        
        self.samples.append(self.sample)
        u = (self.fmin-mu)/(np.sqrt(2)*sdev)
        if self.acquisition_function == 'expected improvement':
            acqf = sdev * np.exp(-u**2) / np.sqrt(2*np.pi) + (self.fmin - mu) * 0.5 * (1 + erf(u))
        elif self.acquisition_function == 'probability of improvement':
            acqf = 0.5*(1+erf(u))   
        elif self.acquisition_function == 'lower confidence bound':
            acqf = -(mu-sdev)
        elif self.acquisition_function == 'test':
            acqf = -mu
        else: raise ValueError('Unkown acquisition function {}'.format(self.acquisition_function))
            
        acqf = np.nan_to_num(acqf)
        
        xMax = self.Xs[np.argmax(acqf)]
        self.sample = xMax
        self.fmin = min(self.fmin,self.f(self.sample))
            
        self.ax1.collections.clear()
        self.ax1_lines[0].set_data(self.samples,self.f(np.array(self.samples)))
        self.ax1_lines[1].set_data(self.Xs, mu)
        self.ax1_lines[2].set_data(self.Xs, mu-sdev)
        self.ax1_lines[3].set_data(self.Xs, mu+sdev)
        self.ax1.fill_between(self.Xs, mu-sdev, mu+sdev,color='lightgrey')
        
        self.ax2_lines[0].set_data(self.Xs,acqf)
        self.ax2_lines[1].set_data([xMax],[max(acqf)])
        return self.ax1_lines+self.ax2_lines

    def animation(self):    
        return animation.FuncAnimation(self.fig, self.update, frames=range(8),
                                           interval=1000, init_func=self.init, blit=True)

    def show(self):
        anim = self.animation()
        clear_output()
        return HTML(anim.to_jshtml())
