# -*- coding: utf-8 -*-


import scipy
from scipy.optimize import curve_fit
from uncertainties import unumpy as unp
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors

dimension=1.0
T=1.0
mu=1.0
D=T*mu
k=1
g=mu*k
tau=1/k/mu
d=np.sqrt(D*tau)


plt.rcParams.update({
  "text.usetex": True,
})

size=20.0
plt.rcParams['figure.figsize'] = (6.4, 4.8) 
plt.rcParams['axes.labelsize'] =size
plt.rcParams['xtick.labelsize'] =size
plt.rcParams['ytick.labelsize'] =size
plt.rcParams['xtick.direction'] = "in"
plt.rcParams['ytick.direction'] = "in"
plt.rcParams['legend.fontsize'] =size
plt.rcParams['legend.title_fontsize'] =size
plt.rcParams['legend.frameon'] = False
plt.rcParams['legend.labelspacing'] = 0.05
plt.rcParams['lines.linewidth'] = 2.0
plt.rcParams['lines.markersize'] = 5.0
plt.rcParams['font.family'] = 'lmodern' 


def positiondata(string):
    f1 = open(string,'r')
    lines1 = f1.readlines()
    f1.close()
    t=[]
    y=[]
    y_nc=[]
    y_cc=[]    

    for line in lines1[1:28]:
        p=line.split()
        t.append(float(p[3]))
        y.append(float(p[4]))
        y_nc.append(float(p[5]))
        y_cc.append(float(p[6]))

    t_ret=[]

    y_ret=[]
    y_nc_ret=[]
    y_cc_ret=[]

    for i in range(0,len(t)):     
        if i%10!=0 or i==0:
            t_ret.append(t[i])
            #for the MSD
            y_ret.append(y[i])
            y_nc_ret.append(y_nc[i])
            y_cc_ret.append(y_cc[i])
            
    return t_ret,y_ret,y_nc_ret,y_cc_ret



num2=50
y=np.array([positiondata('../data/Data_HP/Data_280121_HO_dt_01/Data_'+str(i)+'/data_msd'+'.dat') for i in range(0,num2)])
t=y[0,0]

y_2_mean=np.mean( np.array([y[i][1] for i in range(0,len(y))]), axis=0 )
y_2_nc_mean=np.mean( np.array([y[i][2] for i in range(0,len(y))]), axis=0 )
y_2_cc_mean=np.mean( np.array([y[i][3] for i in range(0,len(y))]), axis=0 )

y_2_std=np.std( np.array([y[i][1] for i in range(0,len(y))]), axis=0 )/np.sqrt(num2)
y_2_nc_std=np.std( np.array([y[i][2] for i in range(0,len(y))]), axis=0 )/np.sqrt(num2)
y_2_cc_std=np.std( np.array([y[i][3] for i in range(0,len(y))]), axis=0 )/np.sqrt(num2)


y_2=unp.uarray(y_2_mean,y_2_std)
y_2_nc=unp.uarray(y_2_nc_mean,y_2_nc_std)
y_2_cc=unp.uarray(y_2_cc_mean,y_2_cc_std)


#-----------------------------------------------------------------------------------------
num3=50
t=y[0,0]
y=np.array([positiondata('../data/Data_HP/Data_280121_HO_dt_001/Data_'+str(i)+'/data_msd'+'.dat') for i in range(0,num3)])

y_3_mean=np.mean( np.array([y[i][1] for i in range(0,len(y))]), axis=0 )
y_3_nc_mean=np.mean( np.array([y[i][2] for i in range(0,len(y))]), axis=0 )
y_3_cc_mean=np.mean( np.array([y[i][3] for i in range(0,len(y))]), axis=0 )

y_3_std=np.std( np.array([y[i][1] for i in range(0,len(y))]), axis=0 )/np.sqrt(num3)
y_3_nc_std=np.std( np.array([y[i][2] for i in range(0,len(y))]), axis=0 )/np.sqrt(num3)
y_3_cc_std=np.std( np.array([y[i][3] for i in range(0,len(y))]), axis=0 )/np.sqrt(num3)


y_3=unp.uarray(y_3_mean,y_3_std)
y_3_nc=unp.uarray(y_3_nc_mean,y_3_nc_std)
y_3_cc=unp.uarray(y_3_cc_mean,y_3_cc_std)
#-----------------------------------------------------------------------------------------

num4=50
y=np.array([positiondata('../data/Data_HP/Data_280121_HO_dt_0001/Data_'+str(i)+'/data_msd'+'.dat') for i in range(0,num3)])

t=y[0,0]

y_4_mean=np.mean( np.array([y[i][1] for i in range(0,len(y))]), axis=0 )
y_4_nc_mean=np.mean( np.array([y[i][2] for i in range(0,len(y))]), axis=0 )
y_4_cc_mean=np.mean( np.array([y[i][3] for i in range(0,len(y))]), axis=0 )
y_4_std=np.std( np.array([y[i][1] for i in range(0,len(y))]), axis=0 )/np.sqrt(num4)
y_4_nc_std=np.std( np.array([y[i][2] for i in range(0,len(y))]), axis=0 )/np.sqrt(num4)
y_4_cc_std=np.std( np.array([y[i][3]for i in range(0,len(y))]), axis=0 )/np.sqrt(num4)

y_4=unp.uarray(y_3_mean,y_3_std)
y_4_nc=unp.uarray(y_3_nc_mean,y_3_nc_std)
y_4_cc=unp.uarray(y_3_cc_mean,y_3_cc_std)

#-----------------------------------------------------------------------------------------

#Functions

#------------------------functions without dt
def analytical_dt_MSD(t,dt):
    return 4*D/g/(g*dt-2)*((1-g*dt)**(t/dt)-1)

def analytical_dt_CC(t,dt):
    return (2*D*dt*(-1 + (1 - dt*g)**(t/dt)))/(-2 + dt*g)

def analytical_dt_MSD_NC(t,dt):
    return 2*D/g/(-2+ dt*g)*(-2*(1-g*dt)**(t/dt)+g*dt*(2*(1-g*dt)**(t/dt)+g*dt*(t/dt)-2*(t/dt)-2)+2)

#------------------------functions without dt
def funcii(t):
    return 2*D/g*(1-np.exp(-g*t))


#-----------------------------------------------------------------------------------------


#for the matrix legend
import string
from matplotlib.legend_handler import HandlerBase
from matplotlib.text import Text
from matplotlib.legend import Legend


class TextHandlerB(HandlerBase):
    def create_artists(self, legend, text ,xdescent, ydescent,
                        width, height, fontsize, trans):
        tx = Text(width/2.,height/2, text, fontsize=fontsize,
                  ha="center", va="center", fontweight="bold")
        return [tx]

Legend.update_default_handler_map({str : TextHandlerB()})
#-----------

import matplotlib.ticker as ticker
from matplotlib.legend_handler import HandlerTuple

fig, ax = plt.subplots()

# Set up plot
ax.set_xscale('log')
ax.set_yscale('log')
ax.set_xlabel(r'$t/\tau$')
ax.set_ylabel(r'$\Sigma (t) / D \tau $')

# Data to plot
msd_data = [
    (t, y_4_std, r"$1$", 'solid',"o", '#6BAED6','#6BAED6'), #leightblue
    (t, y_4_nc_std, r"$1$ NS" , 'solid', "s", '#FFA07A','#FFA07A' ), #lightorange
    (t, y_3_std,"$10$ ", '-', "^",  'blue','blue' ), #darkblue
    (t, y_3_nc_std, r"$10$ NS", '-',"D",'red','red' ), #darkorange
]

# Set the major and minor ticks on the y-axis
y_major_ticks = ticker.LogLocator(base=10, subs=[1.0], numticks=10)
y_minor_ticks = ticker.LogLocator(base=10, subs=np.arange(2, 10) * 0.1, numticks=10)
ax.yaxis.set_major_locator(y_major_ticks)
ax.yaxis.set_minor_locator(y_minor_ticks)


# Plot MSD data
for t, y, label, linestyle,marker, color,ecolor in msd_data:
    ax.errorbar(t/ tau, y/ d**2,linewidth=1, ecolor=ecolor,linestyle=linestyle, marker=marker,color=color, label=label)
 
p2, = ax.plot([0], [0], 'o', color="#6BAED6", linestyle='-') #leightblue
p1, = ax.plot([0], [0], 's', color='#FFA07A',  linestyle='-') #leightorange
p3, = ax.plot([0], [0], 'D', linestyle='-', color='red') #darkorange
p4, = ax.plot([0], [0], '^', linestyle='-', color='blue') #darkblue

legend3 =  ax.legend(handles=[(r'red \medspace std'), (p3, p4),(p1, p2)], labels=[r"$\Delta t/ \tau$" ,r'$10^{-3}$', r'$10^{-4}$'],handlelength=3.2,handler_map={tuple: HandlerTuple(ndivide=None)})
ax.add_artist(legend3)


ax.tick_params(axis='both', which='both', right='on',top='on',pad=6)
plt.tight_layout()
plt.savefig('C:/Users/Regin/OneDrive/nsp_algorithm/plots/Fig_3.pdf')

plt.show()

#https://stackoverflow.com/questions/27174425/how-to-add-a-string-as-the-artist-in-matplotlib-legend




#------------
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(7, 10)) 

left, bottom, width, height = [0.72, 0.59, 0.23, 0.13]
ax3 = fig.add_axes([left, bottom, width, height])
xxxxx=np.linspace(-1,1,10000)    

ax3.spines['bottom'].set_color('white')
ax3.spines['top'].set_color('white')
ax3.spines['left'].set_color('white')
ax3.spines['right'].set_color('white')
ax3.plot(xxxxx,xxxxx**2, color='black')
ax3.plot(0.25, 0.5, marker='o',markersize=12.0,color='b')
ax3.arrow(0.25, 0.5, 0.20, 0.0,head_width=0.1, head_length=0.08, width=0.01, fc='k', ec='k')
ax3.arrow(0.25, 0.5, -0.20, 0.0,head_width=0.10, head_length=0.08, width=0.01, fc='k', ec='k')




ax3.set_ylabel(r'$U(x)$')
ax3.set_xticks([], [])
ax3.set_yticks([], [])
ax3.tick_params(length=0)

#------------
# Set up plot
ax1.set_xscale('log')
ax1.set_yscale('log')
ax1.set_xlabel(r'$t/\tau$')


# Set the major and minor ticks on the y-axis
y_major_ticks = ticker.LogLocator(base=10, subs=[1.0], numticks=10)
y_minor_ticks = ticker.LogLocator(base=10, subs=np.arange(2, 10) * 0.1, numticks=10)
ax1.yaxis.set_major_locator(y_major_ticks)
ax1.yaxis.set_minor_locator(y_minor_ticks)

# Data to plot
msd_data = [
    (t, y_2_mean,y_2_std,"$10$ ", 'none', "^",  'blue','blue' ), #darkblue
    (t,y_2_nc_mean, y_2_nc_std, r"$10$ NS", 'none',"D",'red','red' ), #darkorange

]

# Plot MSD data
for t, y, yerr, label, linestyle,marker, color,ecolor in msd_data:
    ax1.errorbar(t/ tau, y/ d**2, yerr=yerr/ d**2, linewidth=1, ecolor=ecolor,linestyle=linestyle, marker=marker,color=color, label=label)
 
p3, = ax.plot([0], [0], 'D', linestyle='-', color='red') #darkorange
p4, = ax.plot([0], [0], '^', linestyle='-', color='blue') #darkblue
legend2 = ax1.legend([p3, p4], [r'$\langle\Delta x^{ \mathrm{red}}(t) ^2 \rangle / D \tau $', r'$\langle \Delta x(t) ^2 \rangle / D \tau $'], loc='upper left',handlelength=2.2,  handler_map={tuple: HandlerTuple(ndivide=None)})

xxxx=np.linspace(0.01,10,10000) 
ax1.plot(xxxx,2*xxxx,'--',color='grey',label='2Dt')
ax1.text(1,3.5,r'$2t/\tau$',rotation=23 ,fontsize=20, color="grey")
ax1.plot(xxxx,funcii(xxxx),label=r'analytics')
ax1.plot(xxxx,analytical_dt_MSD_NC(xxxx,0.01), color='red')
ax1.add_artist(legend2)
ax1.tick_params(axis='both', which='both', right='on',top='on',pad=6)

#------------

ax2.set_xscale('log')    
ax2.set_yscale('log')    
ax2.set_xlabel(r'$t/\tau$')
ax2.set_ylabel(r'$\langle \Delta x^{ \mathrm{red}}(t)  \Delta x(t)\rangle / D \tau $ ')
ax2.errorbar(t[0:28], y_2_cc_mean[0:28], yerr=y_2_cc_std[0:28], marker='o',color='#006203',linestyle="None")
ax2.errorbar(t[0:28], y_3_cc_mean[0:28], yerr=y_3_cc_std[0:28], marker='s',color='#0f9200',linestyle="None")
ax2.errorbar(t[0:22], y_4_cc_mean[0:22], yerr=y_4_cc_std[0:22], marker='D',color='#30cb00',linestyle="None")

xxxx=np.logspace(-2,1)
ax2.plot(xxxx,analytical_dt_CC(xxxx,0.01),color='#006203',label=r"$10^{-2}$")
ax2.plot(xxxx,analytical_dt_CC(xxxx,0.001),color='#0f9200',label=r"$10^{-3}$")
ax2.plot(xxxx,analytical_dt_CC(xxxx,0.0001),color='#30cb00',label=r"$10^{-4}$")
ax2.legend(loc=4, title=r"$\Delta t/\tau$")


p1, = ax.plot([0], [0], 'o', color="#006203", linestyle='-') 
p2, = ax.plot([0], [0], 's', color='#0f9200',  linestyle='-')
p3, = ax.plot([0], [0], 'D', color='#30cb00',linestyle='-')

legend2 = ax2.legend([p1, p2,p3], [r"$10^{-2}$",r"$10^{-3}$",r"$10^{-4}$"], loc='lower right',handlelength=2.2,  handler_map={tuple: HandlerTuple(ndivide=None)},title=r"$\Delta t/\tau$")
ax2.add_artist(legend2)



ax2.tick_params(axis='both', which='both', right='on',top='on',pad=6)
ax1.text(-0.15, 0.90, '(a)', transform=ax1.transAxes, fontsize=20,verticalalignment='top')
ax2.text(-0.15, 0.90, '(b)', transform=ax2.transAxes, fontsize=20,verticalalignment='top')
plt.tight_layout()
plt.savefig('C:/Users/Regin/OneDrive/nsp_algorithm/plots/Fig_2.pdf')

plt.show()




def numerisch(x, y, y_fehler, name,u,o):
    f4_der=[]
    x_f4=[]
    y_vacf = unp.uarray(y,y_fehler)
    for i in range(1,len(x)-1):
        if i%9!=0:
            h1 = x[i]-x[i-1]
            h2 = x[i+1]-x[i]
            f4_der.append((h1*y_vacf[i+1]-(h1+h2)*y_vacf[i]+h2*y_vacf[i-1])/(0.5*(h2*h1**2+h1*h2**2)))
            x_f4.append(x[i])

    x_f4=np.array(x_f4)
    f4_der=np.array(f4_der)

    plt.yscale('log')    
    plt.xscale('log')  
    plt.errorbar(x_f4, unp.nominal_values(f4_der)/2./dimension, yerr=unp.std_devs(f4_der)/2./dimension,marker='.',linestyle="None",label=name)
    
    return x_f4, unp.nominal_values(f4_der)/2./dimension, unp.std_devs(f4_der)/2./dimension



#errors
x_cc,y_cc,yerr_cc=numerisch(t,y_2_cc_mean, y_2_cc_std,'false',0,3)
x_nc,y_nc,yerr_nc=numerisch(t,y_2_nc_mean, y_2_nc_std,'false',0,3)
x,y,yerr=numerisch(t,y_2_mean, y_2_std,'false',0,3)
plt.show()

#-----------

plt.yscale('log')    
plt.xscale('log')  
plt.errorbar(x_nc, y_nc, yerr=yerr_nc,  color='#ff7f0e',marker=".",linestyle="None",label=r'$\mathrm{Z(t)}_{\mathrm{NC}}$')
plt.errorbar(x, -y, yerr=yerr, color='#1f77b4',marker=".",linestyle="None",label=r'Z(t)')
plt.errorbar(x_cc,-y_cc/0.005, yerr=yerr_cc, color='black',marker=".",linestyle="None",label=r'Frenkel')

#numerisch(t,y_3_nc_mean, y_3_nc_std, 'dsf')
#numerisch(t[0:],y_3_nc_mean[0:], y_3_nc_std[0:], 'dsf','true',0,3)
#plt.xlim(10**(-0.8),1)
plt.xlabel(r'$t/\tau$')
plt.ylabel(r'$-Z(t)\tau^2 /d^2$')

plt.xscale('log')    
xx=np.logspace(-1.8,0.765,100)
#plt.loglog(xx,-1/2./dimension*funct_VACF_HO(xx),label=r"$2D\gamma e^{-\gamma t}$")

#plt.xlim(0.19,11)
plt.legend()
#plt.annotate('', xy=(5.5, 0.0026), xytext=(2, 0.0022-0.0008), 
#            arrowprops=dict(facecolor='black', shrink=0.1),
#            )
plt.tight_layout()
#plt.savefig('C:/Users/Regina/OneDrive/9.Semester/Latex_MA/Teil2/t2_HO_VACF_nc.pdf')
#plt.savefig('C:/Users/Regina/OneDrive/9.Semester/Latex_MA/Teil2/t2_HO_VACF_nc_pres.pdf')
plt.show()



plt.yscale('log')    
plt.xscale('log')  
plt.plot(x_nc, yerr_nc,   color='#ff7f0e',marker=".",linestyle="None",label=r'$\mathrm{Z(t)}_{\mathrm{NC}}$')
plt.plot(x, yerr, color='#1f77b4',marker=".",linestyle="None",label=r'Z(t)')
plt.plot(x_cc,yerr_cc,  color='black',marker=".",linestyle="None",label=r'Frenkel')

#numerisch(t,y_3_nc_mean, y_3_nc_std, 'dsf')
#numerisch(t[0:],y_3_nc_mean[0:], y_3_nc_std[0:], 'dsf','true',0,3)
#plt.xlim(10**(-0.8),1)
plt.xlabel(r'$t/\tau$')
plt.ylabel(r'Fehler $-Z(t)\tau^2 /d^2$')

plt.xscale('log')    
xx=np.logspace(-1.8,0.765,100)
#plt.loglog(xx,-1/2./dimension*funct_VACF_HO(xx),label=r"$2D\gamma e^{-\gamma t}$")

#plt.xlim(0.19,11)
plt.legend()
#plt.annotate('', xy=(5.5, 0.0026), xytext=(2, 0.0022-0.0008), 
#            arrowprops=dict(facecolor='black', shrink=0.1),
#            )
plt.tight_layout()

plt.show()