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

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(Ors,ps,Ts):
    if len(Ors)==1:
        popt,pcov=curve_fit(fitfunc,ps,Ts,p0=np.array([1,pi/2]),\
        bounds=([0,-pi],[1,pi]))
        amps,offsets=popt[0],popt[1]
        err_amps,err_offsets=np.sqrt(pcov[0,0]),np.sqrt(pcov[1,1])

    else:
        amps,offsets,err_amps,err_offsets=np.zeros(len(Ors)),np.zeros(len(Ors)),np.zeros(len(Ors)),np.zeros(len(Ors))
        for i in range(len(Ors)):
            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(Ors,ps,Ts):
    Or_c=8
    if len(Ors)==1:
        if Ors>Or_c:
            bounds=([0,-pi,1],[1,pi,100])
            popt,pcov=curve_fit(fitfunc_skew1,ps,Ts,p0=np.array([1,pi/2,1]),\
            bounds=bounds)
            amps,offsets,zs=popt[0],popt[1],popt[2]
            err_amps,err_offsets,err_zs=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,p0=np.array([1,pi/2,0.05]),\
            bounds=bounds)
            amps,offsets,zs=popt[0],popt[1],1/popt[2]
            err_amps,err_offsets,err_zs=np.sqrt(pcov[0,0]),np.sqrt(pcov[1,1]),np.sqrt(pcov[2,2])*popt[2]**2
    else:
        amps,offsets,zs,err_amps,err_offsets,err_zs=np.zeros(len(Ors)),np.zeros(len(Ors)),np.zeros(len(Ors)),np.zeros(len(Ors)),np.zeros(len(Ors)),np.zeros(len(Ors))
        for i in range(len(Ors)):
            if Ors[i]>Or_c:
                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

b=2
l=np.sqrt(2)*b
def w_conv(Or):
    return (l/(2*Or))

labels=[]

z_inverse=True

phs_1_1=np.loadtxt('ir_3s2/T2s_1s/phs_1s.txt')
Ors_s=[]
for filename in os.listdir('ir_3s2/T2s_1s'):
    if len(filename)==8:
        Ors_s.append(filename[:4])
nOr_1_1=len(Ors_s)
#print(nOr_1)
Ors_1_1=np.zeros(nOr_1_1)
for j in range(nOr_1_1):
    Ors_1_1[j]=float(Ors_s[j])
Ors_1_1_temp=np.zeros(nOr_1_1)
inds=np.argsort(Ors_1_1)
for j in range(nOr_1_1):
    Ors_1_1_temp[j]=Ors_1_1[inds[j]]
Ors_1_1=Ors_1_1_temp
nph_1_1=len(phs_1_1)
T2s_1_1=np.zeros((nOr_1_1,nph_1_1))
for j in range(nOr_1_1):
    T2s_1_1[j,:]=np.loadtxt('ir_3s2/T2s_1s/%04.1f.txt'%(Ors_1_1[j]))

amps_1_1,offsets_1_1,zs_1_1,err_amps_1_1,err_offsets_1_1,err_zs_1_1=amps_func_skew(Ors_1_1,phs_1_1,T2s_1_1)
offsets_1_1=offsets_1_1/(2*pi)
err_offsets_1_1=err_offsets_1_1/(2*pi)
if z_inverse:
    zs_1_1=1/zs_1_1
    err_zs_1_1=err_zs_1_1*zs_1_1**2

nOr_1_g,nph_1_g=nOr_1_1,nph_1_1
Ors_1_g,phs_1_g=Ors_1_1,phs_1_1
gs=np.array([1,8,40])
ng=len(gs)
for i in range(1,ng):
    T2s=np.zeros((nOr_1_g,nph_1_g))
    for j in range(nOr_1_g):
        T2s[j,:]=np.loadtxt('ir_3s2/T2s_1sg/%04.1f_%d.txt'%(Ors_1_g[j],gs[i]))
    exec("T2s_1_%d=T2s"%(gs[i]))
    
    amps_1,offsets_1,zs_1,err_amps_1,err_offsets_1,err_zs_1=amps_func_skew(Ors_1_g,phs_1_g,T2s)
    if z_inverse:
        zs_1=1/zs_1
        err_zs_1=err_zs_1*zs_1**2
    exec("zs_1_%d=zs_1"%(gs[i]))
    exec("err_zs_1_%d=err_zs_1"%(gs[i]))
    offsets_1,err_offsets_1=offsets_1/(2*pi),err_offsets_1/(2*pi)
    exec("amps_1_%d=amps_1"%(gs[i]))
    exec("err_amps_1_%d=err_amps_1"%(gs[i]))
    exec("offsets_1_%d=offsets_1"%(gs[i]))
    exec("err_offsets_1_%d=err_offsets_1"%(gs[i]))

sep_o=True
if not sep_o:
    o,D=1e4,0e0#1e4 for 10,20,30,40,5e4 for 50,60,70
if sep_o:
    D=0
    o_bound=35
    ors=np.array([1e4,5e4])
for i in range(ng):
    folder_='ir_3s2'
    Ors_s=[]
    if not sep_o:
        for folder in os.listdir(folder_):
            if folder[5:]=='%05.2f_%7.2e'%(gs[i],o):
                if os.path.exists(folder_+'/%s_%05.2f_%7.2e/phs.txt'%(folder[:4],gs[i],o)):
                    Ors_s.append(folder[:4])
        nOr=len(Ors_s)
        Ors=np.zeros(nOr)
        a0s,T1s=np.zeros(nOr),np.zeros(nOr)
        amps,offsets,err_amps,err_offsets=np.zeros(nOr),np.zeros(nOr),np.zeros(nOr),np.zeros(nOr)
        for j in range(nOr):
            Ors[j]=float(Ors_s[j])
        Ors_temp=np.zeros(nOr)
        inds=np.argsort(Ors)
        for j in range(nOr):
            Ors_temp[j]=Ors[inds[j]]
        Ors=Ors_temp
        for j in range(nOr):
            a0s[j],T1s[j]=np.loadtxt(folder_+'/%04.1f_%05.2f_%7.2e/a0s.txt'%(Ors[j],gs[i],o))
        amps,offsets,zs,err_amps,err_offsets,err_zs=np.zeros(nOr),np.zeros(nOr),np.zeros(nOr),np.zeros(nOr),np.zeros(nOr),np.zeros(nOr)
        for j in range(nOr):
            a0s[j],T1s[j]=np.loadtxt(folder_+'/%04.1f_%05.2f_%7.2e/a0s.txt'%(Ors[j],gs[i],o))
            amps[j],offsets[j],zs[j],err_amps[j],err_offsets[j],err_zs[j]=amps_func_skew(np.array([Ors[j]]),np.loadtxt(folder_+'/%04.1f_%05.2f_%7.2e/phs.txt'%(Ors[j],gs[i],o)),np.loadtxt(folder_+'/%04.1f_%05.2f_%7.2e/T2s.txt'%(Ors[j],gs[i],o)))
        if z_inverse:
                zs=1/zs
                err_zs=err_zs*zs**2
    if sep_o:
        for folder in os.listdir(folder_):
            if len(folder)>8:
                if float(folder[:4])<o_bound:
                    o=ors[0]
                else:
                    o=ors[1]
            if folder[5:]=='%05.2f_%7.2e'%(gs[i],o):
                if os.path.exists(folder_+'/%s_%05.2f_%7.2e/phs.txt'%(folder[:4],gs[i],o)):
                    Ors_s.append(folder[:4])
        nOr=len(Ors_s)
        Ors=np.zeros(nOr)
        a0s,T1s=np.zeros(nOr),np.zeros(nOr)
        amps,offsets,err_amps,err_offsets=np.zeros(nOr),np.zeros(nOr),np.zeros(nOr),np.zeros(nOr)
        for j in range(nOr):
            Ors[j]=float(Ors_s[j])
        Ors_temp=np.zeros(nOr)
        inds=np.argsort(Ors)
        for j in range(nOr):
            Ors_temp[j]=Ors[inds[j]]
        Ors=Ors_temp
        for j in range(nOr):
            if Ors[j]<o_bound:
                o=ors[0]
            else:
                o=ors[1]
            a0s[j],T1s[j]=np.loadtxt(folder_+'/%04.1f_%05.2f_%7.2e/a0s.txt'%(Ors[j],gs[i],o))
        amps,offsets,zs,err_amps,err_offsets,err_zs=np.zeros(nOr),np.zeros(nOr),np.zeros(nOr),np.zeros(nOr),np.zeros(nOr),np.zeros(nOr)
        for j in range(nOr):
            if Ors[j]<o_bound:
                o=ors[0]
            else:
                o=ors[1]
            amps[j],offsets[j],zs[j],err_amps[j],err_offsets[j],err_zs[j]=amps_func_skew(np.array([Ors[j]]),np.loadtxt(folder_+'/%04.1f_%05.2f_%7.2e/phs.txt'%(Ors[j],gs[i],o)),np.loadtxt(folder_+'/%04.1f_%05.2f_%7.2e/T2s.txt'%(Ors[j],gs[i],o)))
        if z_inverse:
            zs=1/zs
            err_zs=err_zs*zs**2
        offsets,err_offsets=offsets/(2*pi),err_offsets/(2*pi)
        exec('amps_%d,offsets_%d,zs_%d,err_amps_%d,err_offsets_%d,err_zs_%d=amps,offsets,zs,err_amps,err_offsets,err_zs'%(gs[i],gs[i],gs[i],gs[i],gs[i],gs[i]))
        exec('Ors_%d=Ors'%(gs[i]))
        labels.append('$g=%d$'%gs[i])


fig=plt.figure(figsize=figsize)
heights,widths=[1,7,5,2.4,13],[10,1,10]
gspec=gridspec.GridSpec(ncols=3,nrows=5,width_ratios=widths,height_ratios=heights)
ax1=fig.add_subplot(gspec[4,2])
ax2=fig.add_subplot(gspec[4,0])
ax3=fig.add_subplot(gspec[1:3,2])
ax4a=fig.add_subplot(gspec[1,0])
ax4b=fig.add_subplot(gspec[2,0])
ax3.yaxis.set_ticks([0.2,0.4])
ssize=32
lwidth=2
csize=6
ax1.set_ylabel('$\\varepsilon/2\\pi$',rotation=90,labelpad=-5)
ax2.set_ylabel('$A$',rotation=90,labelpad=5)
ax3.set_ylabel('$z^{-1}$',rotation=90,labelpad=10)
emin,emax=0.23,0.3
amin,amax=0.92,1.02
zmin,zmax=0.0,0.45
xmin=8
xmax=74
ax1.set_ylim(emin,emax)
ax2.set_ylim(amin,amax)
ax3.set_ylim(zmin,zmax)
ax1.set_xlim(xmin=xmin,xmax=xmax)
ax2.set_xlim(xmin=xmin,xmax=xmax)
ax3.set_xlim(xmin=xmin,xmax=xmax)
colours=['blue','red','green']
colours2=['cornflowerblue','tomato','mediumspringgreen']
markers=['o','v','x','s','D','1','.','p','*','P']
ideal_lims=True
if ideal_lims:
    N=100
    ax1.plot(np.linspace(0,xmax,N),np.ones(N)*0.25,linestyle=':',color='black',zorder=1)
    ax2.plot(np.linspace(0,xmax,N),np.ones(N),linestyle=':',color='black',zorder=1)
    
ax1.set_xlabel('$\Omega_1/\Omega_0$',fontsize=48,labelpad=10)
ax2.set_xlabel('$\Omega_1/\Omega_0$',fontsize=48,labelpad=10)
ax3.set_xlabel('$\Omega_1/\Omega_0$',fontsize=48,labelpad=10)
ax1.xaxis.set_ticks([20,40,60])
ax2.xaxis.set_ticks([20,40,60])
ax3.xaxis.set_ticks([20,40,60])
zorders1=np.array([5,4,3])
zorders2=zorders1+4
for i in range(ng):
    exec("ax1.errorbar(Ors_%d,offsets_%d,yerr=err_offsets_%d,fmt='none',capsize=csize,color=colours[i],linewidth=None,zorder=zorders1[i])"%(gs[i],gs[i],gs[i]))
    exec("ax2.errorbar(Ors_%d,amps_%d,yerr=err_amps_%d,fmt='none',capsize=csize,color=colours[i],linewidth=None,zorder=zorders1[i])"%(gs[i],gs[i],gs[i]))
    exec("ax3.errorbar(Ors_%d,zs_%d,yerr=err_zs_%d,fmt='none',capsize=csize,color=colours[i],linewidth=None,zorder=zorders1[i])"%(gs[i],gs[i],gs[i]))
    exec("ax1.scatter(Ors_%d,offsets_%d,color=colours[i],marker=markers[i],s=ssize,label='%s',zorder=zorders2[i])"%(gs[i],gs[i],labels[i]))
    exec("ax2.scatter(Ors_%d,amps_%d,color=colours[i],marker=markers[i],s=ssize,zorder=zorders2[i])"%(gs[i],gs[i]))
    exec("ax3.scatter(Ors_%d,zs_%d,color=colours[i],marker=markers[i],s=ssize,zorder=zorders2[i])"%(gs[i],gs[i]))
Nc=1000
Ors_1c=np.linspace(Ors_1[0],Ors_1[-1],Nc)
epsc=np.linspace(emin,emax,Nc)
ampsc=np.linspace(amin,amax,Nc)
zsc=np.linspace(zmin,zmax,Nc)
for i in range(ng):
    exec("amps_1=amps_1_%d"%(gs[i]))
    exec("err_amps_1=err_amps_1_%d"%(gs[i]))
    exec("offsets_1=offsets_1_%d"%(gs[i]))
    exec("err_offsets_1=err_offsets_1_%d"%(gs[i]))
    exec("zs_1=zs_1_%d"%(gs[i]))
    exec("err_zs_1=err_zs_1_%d"%(gs[i]))
    fe_min,fa_min,fz_min=interp1d(Ors_1_1,offsets_1-err_offsets_1),interp1d(Ors_1_1,amps_1-err_amps_1),interp1d(Ors_1_1,zs_1-err_zs_1)
    fe_max,fa_max,fz_max=interp1d(Ors_1_1,offsets_1+err_offsets_1),interp1d(Ors_1_1,amps_1+err_amps_1),interp1d(Ors_1_1,zs_1+err_zs_1)
    fe_me,fa_me,fz_me=interp1d(Ors_1_1,offsets_1),interp1d(Ors_1_1,amps_1),interp1d(Ors_1_1,zs_1)
    e_min,e_max,e_me=fe_min(Ors_1c),fe_max(Ors_1c),fe_me(Ors_1c)
    a_min,a_max,a_me=fa_min(Ors_1c),fa_max(Ors_1c),fa_me(Ors_1c)
    z_min,z_max,z_me=fz_min(Ors_1c),fz_max(Ors_1c),fz_me(Ors_1c)
    ax1.fill_between(Ors_1c,e_min,e_max,color=colours2[i],zorder=1)
    ax2.fill_between(Ors_1c,a_min,a_max,color=colours2[i],zorder=1)
    ax3.fill_between(Ors_1c,z_min,z_max,color=colours2[i],zorder=1)
    ax1.plot(Ors_1c,e_me,color=colours[i],zorder=2)
    ax2.plot(Ors_1c,a_me,color=colours[i],zorder=2)
    ax3.plot(Ors_1c,z_me,color=colours[i],zorder=2)
fig.legend(ncol=ng,loc=9,markerscale=3,fontsize=48)
g=40
ax4b.set_xlabel('$\\theta/(2\\pi)$',labelpad=0)
ax4a.set_xticklabels([])
ax4a.set_ylabel('$T_2$',rotation=90,labelpad=10)
ax4b.set_ylabel('$\\zeta$',rotation=90,labelpad=10)
ax4a.set_ylim(0,1)
ax4b.tick_params(axis='x',pad=10)
if z_inverse:
    z_label='$z^{-1}$'
else:
    z_label='$z$'
exec("Ors=Ors_%d"%g)
Ors=Ors[1::2]
nOr=len(Ors)
exec("amps=amps_%d"%g)
exec("err_amps=err_amps_%d"%g)
exec("offsets=offsets_%d"%g)
exec("err_offsets=err_offsets_%d"%g)
exec("zs=zs_%d"%g)
exec("err_zs=err_zs_%d"%g)
amps=amps[1::2]
err_amps=err_amps[1::2]
offsets=offsets[1::2]
err_offsets=err_offsets[1::2]
zs=zs[1::2]
err_zs=err_zs[1::2]
for i in range(nOr):
    if sep_o:
        if Ors[i]<o_bound:
            o=ors[0]
        else:
            o=ors[1]
    T2s=np.loadtxt(folder_+'/%04.1f_%05.2f_%7.2e/T2s.txt'%(Ors[i],g,o))
    exec("T2s_%d=T2s"%i)
    phs=np.loadtxt(folder_+'/%04.1f_%05.2f_%7.2e/phs.txt'%(Ors[i],g,o))
    exec("phs_%d=phs"%i)
zorders=[6,5,4,3,2,1,0]
markers=['P','v','s','x','D','1','.','p','*']
ssize2=64
N_ph_fit=200
phs_fit=np.linspace(0,2*pi,N_ph_fit)
for i in range(nOr):
    exec('ax4a.scatter(phs_%d/(2*pi),T2s_%d,label="$\\Omega_1/\\Omega_0=%04.1f$",marker=markers[i],s=ssize2,zorder=zorders[i])'%(i,i,Ors[i]))  
    exec('ax4b.scatter(phs_%d/(2*pi),(T2s_%d-fitfunc(phs_%d,1,pi/2)),marker=markers[i],s=ssize2,zorder=zorders[i])'%(i,i,i))  
    if z_inverse:
        ax4a.plot(phs_fit/(2*pi),fitfunc_skew2(phs_fit,amps[i],offsets[i]*(2*pi),zs[i]),label="$\\Omega_1/\\Omega_0=%04.1f$"%(Ors[i]),linewidth=lwidth,zorder=zorders[i])
        ax4b.plot(phs_fit/(2*pi),(fitfunc_skew2(phs_fit,amps[i],offsets[i]*(2*pi),zs[i])-fitfunc(phs_fit,1,pi/2)),linewidth=lwidth,zorder=zorders[i])
    else:
        ax4a.plot(phs_fit/(2*pi),fitfunc_skew1(phs_fit,amps[i],offsets[i]*(2*pi),zs[i]),label="$\\Omega_1/\\Omega_0=%04.1f$"%(Ors[i]),linewidth=lwidth,zorder=zorders[i])
        ax4b.plot(phs_fit/(2*pi),(fitfunc_skew1(phs_fit,amps[i],offsets[i]*(2*pi),zs[i])-fitfunc(phs_fit,1,pi/2)),linewidth=lwidth,zorder=zorders[i])

ax4b.set_xlim(-0.01,1.01)
ax4b.set_ylim(-0.18,0.18)
      
letters=['a)(i','b','c','d']
positionxs,positionys=[0.135,0.575],[0.79,0.37]
positions=[positionxs[0],positionys[0],positionxs[1],positionys[0],positionxs[0],positionys[1],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)
for i in range(3):
    exec("ax%d.tick_params(length=6,width=2)"%(i+1))
ax4a.tick_params(length=6,width=2)
ax4b.tick_params(length=6,width=2)
fig.text(positionxs[0],0.6,'(a)(ii)',fontsize=48)

plt.savefig('fig_5.eps')
