from __future__ import division
import numpy as np
from numpy import pi
from scipy.optimize import curve_fit
import sys,os
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from matplotlib import ticker
from matplotlib import rc
import warnings
figsize=(20,14)
plt.rc('text', usetex=True)
plt.rc('font', family='serif',serif='Times',size=48)

folder='ir'
folderdf='ir_df'
inds=[1,2,4]
for i in range(len(inds)):
    exec("T2s_%d=np.loadtxt(folder+'/T2s_%d.txt')"%(inds[i],inds[i]))
phs=np.loadtxt(folder+'/phs.txt')
nph=len(phs)
ws=np.loadtxt(folder+'/ws.txt')
df=True
if df:
    a0s_df=np.loadtxt(folderdf+'/a0s.txt')
    wsdf_s=[]
    for filename in os.listdir(folderdf):
        if not filename=='a0s.txt':
            if os.path.exists(folderdf+'/'+filename+'/T2s.txt'):
                wsdf_s.append(filename)
    nwdf=len(wsdf_s)
    wsdf=np.zeros(nwdf)
    phsdf=np.loadtxt(folderdf+'/'+wsdf_s[0]+'/phs.txt')
    nphdf=len(phsdf)
    T2s_df=np.zeros((nwdf,nphdf))
    for i in range(nwdf):
        wsdf[i]=float(wsdf_s[i])
        T2s_df[i,:]=np.loadtxt(folderdf+'/'+wsdf_s[i]+'/T2s.txt')
    wsdf_temp=np.zeros(nwdf)
    T2s_df_temp=np.zeros((nwdf,nphdf))
    inds_dfsort=np.argsort(wsdf)
    for i in range(nwdf):
        wsdf_temp[i]=wsdf[inds_dfsort[i]]
        T2s_df_temp[i,:]=T2s_df[inds_dfsort[i],:]
    wsdf=wsdf_temp
    T2s_df=T2s_df_temp


for i in range(len(inds)):
    a0s=np.zeros(len(ws))
    for j in range(len(ws)):
        a0s[j]=np.loadtxt(folder+'/a0s/%3.2f_%d.txt'%(ws[j],inds[i]))[0]
    exec("a0s_%d=a0s"%(inds[i]))

def fitfunc(theta,A,e):
    T=0.5*(1+A*np.sin(theta-e))
    return T

def clausen(theta,z):
    J1=200
    S=0
    for j in range(1,J1):
        S=S+np.sin(j*theta)/j**z
    return S

def fitfunc_skew1(thetas,A,e,z):
    return 0.5*(1+A*clausen(thetas-e,z))

def fitfunc_skew2(thetas,A,e,z):
    return 0.5*(1+A*clausen(thetas-e,1/z))

def amps_func(ws,ps,Ts):
    amps,offsets,err_amps,err_offsets=np.zeros(len(ws)),np.zeros(len(ws)),np.zeros(len(ws)),np.zeros(len(ws))
    for i in range(len(ws)):
        popt,pcov=curve_fit(fitfunc,ps,Ts[i,:],p0=np.array([1,pi/2]),\
        bounds=([0,-pi],[1,pi]))
        amps[i],offsets[i]=popt[0],popt[1]
        err_amps[i],err_offsets[i]=np.sqrt(pcov[0,0]),np.sqrt(pcov[1,1])
    return amps,offsets,err_amps,err_offsets

def amps_func_skew(ws,ps,Ts):
    amps,offsets,zs,err_amps,err_offsets,err_zs=np.zeros(len(ws)),np.zeros(len(ws)),np.zeros(len(ws)),np.zeros(len(ws)),np.zeros(len(ws)),np.zeros(len(ws))
    for i in range(len(ws)):
        if ws[i]>0.3:
            bounds=([0,-pi,1],[1,pi,100])
            popt,pcov=curve_fit(fitfunc_skew1,ps,Ts[i,:],p0=np.array([1,pi/2,1]),\
            bounds=bounds)
            amps[i],offsets[i],zs[i]=popt[0],popt[1],popt[2]
            err_amps[i],err_offsets[i],err_zs[i]=np.sqrt(pcov[0,0]),np.sqrt(pcov[1,1]),np.sqrt(pcov[2,2])

        else:
            bounds=([0,-pi,0],[1,pi,1])
            popt,pcov=curve_fit(fitfunc_skew2,ps,Ts[i,:],p0=np.array([1,pi/2,0.05]),\
            bounds=bounds)
            amps[i],offsets[i],zs[i]=popt[0],popt[1],1/popt[2]
            err_amps[i],err_offsets[i],err_zs[i]=np.sqrt(pcov[0,0]),np.sqrt(pcov[1,1]),np.sqrt(pcov[2,2])*popt[2]**2
    return amps,offsets,zs,err_amps,err_offsets,err_zs


fig=plt.figure(figsize=figsize)
heights,widths=[1,14,2,14],[10,1,10]
gspec=gridspec.GridSpec(ncols=3,nrows=4,width_ratios=widths,height_ratios=heights)
ax1=fig.add_subplot(gspec[3,2])
ax2=fig.add_subplot(gspec[3,0])
ax3=fig.add_subplot(gspec[1,2])
ax4=fig.add_subplot(gspec[1,0])
ax3.yaxis.set_ticks([0.2,0.4])
ax4.set_xlim(-0.01,1.01)
ssize=32
lwidth=2
csize=6
w_cut=True
if w_cut:
    indsc=[0,2,4,6,8,10,12,14,16,18,20,22,24,26,28,30,31,33,35,37,39,41,43,45,47,49,51,53,55,57,59,61]
    wsc=np.zeros(len(indsc))
    for k in range(len(inds)):
        exec('T2s_%dc=np.zeros((len(indsc),nph))'%inds[k])
    for j in range(len(indsc)):
        wsc[j]=ws[indsc[j]]
        for k in range(len(inds)):
            exec('T2s_%dc[j]=T2s_%d[indsc[j],:]'%(inds[k],inds[k]))
else:
    wsc=ws
for i in range(3):
    exec("ax%d.set_xlabel('$w$',labelpad=0)"%(i+1))
lim=True
if lim:
    ax1.set_ylim(0.17,0.27)
    ax1.set_xlim(xmax=0.64)
    ax2.set_ylim(0.1,1.05)
    ax2.set_xlim(xmax=0.64)
    ax3.set_xlim(xmax=0.64)
ax1.set_ylabel('$\\varepsilon/2\\pi$',rotation=90,labelpad=5)
ax2.set_ylabel('$A$',rotation=90,labelpad=20)
z_inverse=True
if z_inverse:
    ax3.set_ylabel('$z^{-1}$',rotation=90,labelpad=15)
    ax3.set_ylim(ymin=0,ymax=0.5)
else:
    ax3.set_ylabel('$z$',rotation=90,labelpad=15)
    ax3.set_ylim(ymin=1,ymax=52)

colours=['red','green','cyan','black']
zorders1=np.array([4,2,1,3])
zorders2=zorders1+4
markers=['P','v','s','x','D','1','.','p','*']
labels=['$V_h(x)$','$V_{\\mathrm{RM}}(x)$','$V_{RMW}(x)$','$V_{\\mathrm{G}}(x,\\sigma_1=0.2)$','$V_{G}(x,\\sigma_2=2)$','$V_{\\delta}(x)$']
for i in range(len(inds)):
    exec('amps,offsets,zs,err_amps,err_offsets,err_zs=amps_func_skew(wsc,phs,T2s_%dc)'%(inds[i]))
    offsets,err_offsets=offsets/(2*pi),err_offsets/(2*pi)
    if z_inverse:
        zs=1/zs
        err_zs=err_zs*zs**2
    ax1.errorbar(wsc,offsets,yerr=err_offsets,color=colours[i],zorder=zorders2[i],fmt='none',capsize=csize)
    ax2.errorbar(wsc,amps,yerr=err_amps,color=colours[i],zorder=zorders2[i],fmt='none',capsize=csize)
    ax3.errorbar(wsc,zs,yerr=err_zs,color=colours[i],zorder=zorders2[i],fmt='none',capsize=csize)
    ax1.scatter(wsc,offsets,color=colours[i],zorder=zorders1[i],label=labels[inds[i]-1],marker=markers[i],s=ssize)
    ax2.scatter(wsc,amps,color=colours[i],zorder=zorders1[i],marker=markers[i],s=ssize)
    ax3.scatter(wsc,zs,color=colours[i],zorder=zorders1[i],marker=markers[i],s=ssize)
if df:
    amps,offsets,zs,err_amps,err_offsets,err_zs=amps_func_skew(wsdf,phsdf,T2s_df)
    offsets,err_offsets=offsets/(2*pi),err_offsets/(2*pi)
    if z_inverse:
        zs=1/zs
        err_zs=err_zs*zs**2
    ax1.errorbar(wsdf,offsets,yerr=err_offsets,color=colours[-1],zorder=zorders2[-1],fmt='none',capsize=csize)
    ax1.scatter(wsdf,offsets,color=colours[-1],zorder=zorders1[-1],label=labels[-1],marker=markers[-3],s=ssize)
    ax2.errorbar(wsdf,amps,yerr=err_amps,color=colours[-1],zorder=zorders2[-1],fmt='none',capsize=csize)
    ax2.scatter(wsdf,amps,color=colours[-1],zorder=zorders1[-1],marker=markers[-3],s=ssize)
    ax3.errorbar(wsdf,zs,yerr=err_zs,color=colours[-1],zorder=zorders2[-1],fmt='none',capsize=csize)
    ax3.scatter(wsdf,zs,color=colours[-1],zorder=zorders1[-1],marker=markers[-3],s=ssize)
fig.legend(loc='upper left',ncol=len(inds)+1,fontsize=48,labelspacing=2,columnspacing=1,markerscale=4,handletextpad=0.2)
ax4.set_xlabel('$\\theta/(2\\pi)$',labelpad=0)
barrier_type=1
if barrier_type=='df':
    phs,ws=phsdf,wsdf
ax4.set_ylabel('$T_2$',rotation=90,labelpad=10)
ax4.set_ylim(0,1)
ax4.tick_params(axis='x',pad=10)
z_inverse=True
if z_inverse:
    z_label='$z^{-1}$'
else:
    z_label='$z$'
exec('amps,offsets,zs,err_amps,err_offsets,err_zs=amps_func_skew(ws,phs,T2s_%s)'%(barrier_type))
inds2=[0,9,19,38,58]
zorders=[4,3,2,1,0]
ssize2=64
N_ph_fit=200
phs_fit=np.linspace(phs[0],phs[-1],N_ph_fit)
for i in range(len(inds2)):
    exec('ax4.scatter(phs/(2*pi),T2s_%s[inds2[i],:],label="w=%3.2f",marker=markers[i],s=ssize2,zorder=zorders[i])'%(barrier_type,ws[inds2[i]]))    
    if z_inverse:
        ax4.plot(phs_fit/(2*pi),fitfunc_skew1(phs_fit,amps[inds2[i]],offsets[inds2[i]],zs[inds2[i]]),label="%s=%3.2f"%(z_label,1/zs[inds2[i]]),linewidth=lwidth,zorder=zorders[i])
    else:
        ax4.plot(phs_fit/(2*pi),fitfunc_skew1(phs_fit,amps[inds2[i]],offsets[inds2[i]],zs[inds2[i]]),label="%s=%3.2f"%(z_label,zs[inds2[i]]),linewidth=lwidth,zorder=zorders[i])
letters=['a','b','c','d']
positionxs,positionys=[0.14,0.58],[0.78,0.37]
positions=[positionxs[0],positionys[0],positionxs[1],positionys[0],positionxs[0],positionys[1]-0.05,positionxs[1],positionys[1]]
for i in range(4):
    positionx,positiony=positions[2*i:2*i+2]
    fig.text(positionx,positiony,'(%s)'%letters[i],fontsize=48)
    exec("ax%d.tick_params(length=6,width=2)"%(i+1))

plt.savefig('fig_4.eps')

