import numpy as np
import matplotlib.pyplot as plt
import matplotlib.pylab as pylab
import seaborn as sns
sns.set()
import matplotlib



if __name__ == "__main__":

    sns.set_style('whitegrid')
    sns.set_palette("bright")

    x = np.arange(-5, 10, 0.01)
    y1 = x**2
    y2 = (x - 4)**2    # 16
    y3 = (x - 1)**2     # 1/4

    a1 = 1
    a2 = 1
    sig_y1 = [np.exp(i*a1)/(1+np.exp(i*a1)) for i in y1]
    sig_y2 = [np.exp(i*a1)/(1+np.exp(i*a1)) for i in y2]
    sig_y3 = [np.exp(i * a1) / (1 + np.exp(i * a1)) for i in y3]
    final_y_one = [(i+j)/2 for (i,j) in zip(sig_y1, sig_y2)]
    sig_y2_1 = [np.exp(i * a2) / (1 + np.exp(i * a2)) for i in y3]
    final_y_two = [(i + j) / 2 for (i, j) in zip(sig_y1, sig_y2_1)]



    ftsize = 30
    littleft = 27
    matplotlib.rcParams['pdf.fonttype'] = 42
    matplotlib.rcParams['ps.fonttype'] = 42
    params = {'legend.fontsize': littleft,
              'axes.labelsize': ftsize,
              'axes.titlesize': ftsize,
              'xtick.labelsize': littleft,
              'ytick.labelsize': littleft}
    pylab.rcParams.update(params)
    lw = 5
    plt.rcParams["font.family"] = "Times New Roman"
    plt.rcParams["legend.handlelength"] = 0.8

    g, (ax1, ax2) = plt.subplots(1, 2, figsize=(13,5))
    types = ['initial',  'stay','sigm']
    linest = ['-.', '--', '-']
    colors = ['C0', 'C3','C2']
    colorid = ['k', 'r',  'g' ,'b']


    ax1.plot(x, y1, 'C4', linewidth = lw, ls='--',label = "     ")
    ax1.plot(x, y2, 'C1', linewidth=lw, ls=':',label = "     ")
    ax1.set_xlim([-6, 11.5])
    ax1.set_ylim([-0.2, 2.5])
    ax1.plot(x, final_y_one, 'C2', linewidth=lw, ls='-', label = "     ")

    ax2.plot(x, y1, 'C4', linewidth=lw, ls='--',label = "    ")
    ax2.plot(x, y3, 'C1', linewidth=lw, ls=':',label = "    ")
    ax2.set_xlim([-6, 11.5])
    ax2.set_ylim([-0.2, 2.5])

    ax2.plot(x, final_y_two, 'C2', linewidth=lw, ls='-', label = "     ")

    g.tight_layout(w_pad = 3)

    g.savefig('fig/sigmoid.pdf', bbox_inches='tight')
    plt.show()



