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 scipy.interpolate import interp1d
from matplotlib import rc
plt.rc('text', usetex=True)
plt.rc('font', family='serif',serif='Times',size=48)
figsize=(20,14)

folder_='dyna1'


b=2
l=2*b

def O_ratio(b,w):
    return (b/w)

def int_t(int_na,dt):
    PT=np.zeros(3)
    N=len(int_na[0,:])
    for i in range(3):
        for j in range(N-1):
            PT[i]=PT[i]+int_na[i,j]+int_na[i,j+1]
        PT[i]=float(PT[i]*dt/2)
    return PT

colours=['blue','red','green','black']
linestyles=['solid','dotted','dashed','dashdot']
scale_ee=True
labelsa=['$P_{g_1}$','$P_{g_2}$','$P_{e}\\times10^{3}$']

g=1
a_rr,a_rg=g,g
O=1e4
D=0e0
ws_s=[]
for folder in os.listdir(folder_):
    if folder[5:]=='%05.2f_%05.2f_%7.2e_%7.2e'%(a_rr,a_rg,O,D):
        if os.path.exists(folder_+'/%s_%05.2f_%05.2f_%7.2e_%7.2e/int_na.txt'%(folder[:4],a_rr,a_rg,O,D)):
            ws_s.append(folder[:4])
nw=len(ws_s)
ws=np.zeros(nw)
a0s,T1s,dts=np.zeros(nw),np.zeros(nw),np.zeros(nw)
for i in range(nw):
    ws[i]=float(ws_s[i])
    a0s[i],T1s[i],dts[i]=np.loadtxt(folder_+'/%3.2f_%05.2f_%05.2f_%7.2e_%7.2e/a0s.txt'%(ws[i],a_rr,a_rg,O,D))
    
ws_temp,T1s_temp,dts_temp=np.zeros(nw),np.zeros((nw)),np.zeros(nw)
inds=np.argsort(ws)
for i in range(nw):
    ws_temp[i]=ws[inds[i]]
    dts_temp[i]=dts[inds[i]]
    T1s_temp[i]=T1s[inds[i]]
ws,T1s,dts=ws_temp,T1s_temp,dts_temp
for i in range(nw):
    exec("int_na_%d=np.loadtxt(folder_+'/%3.2f_%05.2f_%05.2f_%7.2e_%7.2e/int_na.txt')[:,:-1]"%(i,ws[i],a_rr,a_rg,O,D))
    exec("int_na_%d[[0,1],:]=int_na_%d[[1,0],:]"%(i,i))
    exec("ts_%d=np.arange(len(int_na_%d[0,:]))*dts[i]"%(i,i))
    exec("taus_%d=np.linspace(0,1,len(int_na_%d[0,:]))"%(i,i))
    if scale_ee:
        exec("int_na_%d[2,:]=int_na_%d[2,:]*1e3"%(i,i))


fig=plt.figure(figsize=figsize)
heights,widths=[14,3,14,3,14],[20,1,20]
gspec=gridspec.GridSpec(ncols=3,nrows=5,width_ratios=widths,height_ratios=heights)
ax4=fig.add_subplot(gspec[2,2])
ax2=fig.add_subplot(gspec[2,0])
ax1=fig.add_subplot(gspec[0,2])
ax3=fig.add_subplot(gspec[0,0])
pn=4
labelsb=[]
handles=[]
lsize=48
lwidth=2
rota=90
for i in range(3):
    exec("ax%d.set_ylabel(labelsa[i],rotation=rota,labelpad=10,fontsize=lsize)"%(i+1))
    for j in range(nw):
        exec("l%d,=ax%d.plot(taus_%d[::pn],int_na_%d[i,::pn],color=colours[j],linewidth=lwidth,linestyle=linestyles[j])"%(j,i+1,j,j))
ax1.set_xlabel('$\\tau$',labelpad=-5,fontsize=lsize)
ax2.set_xlabel('$\\tau$',labelpad=-5,fontsize=lsize)
ax3.set_xlabel('$\\tau$',labelpad=-5,fontsize=lsize)
ax1.yaxis.set_ticks([0.8,0.9,1])
ax1.set_ylim(0.7,1.07)
ax2.yaxis.set_ticks([0,0.1,0.2])
ax4.yaxis.set_ticks([0,20,40])
ax4.set_xlim(0,0.45)
for j in range(nw):
    labelsb.append('$w=%3.2f$'%ws[j])
    exec("handles.append(l%d)"%j)
PTs=np.zeros((3,nw))
for i in range(nw):
    exec("PTs[:,i]=int_t(int_na_%d,dts[i])"%i)
#print(PTs[0,:],ws)
ax4.set_ylabel('$\\int_0^\\tau P_{g_2} \\mathrm{d}t$',rotation=rota,fontsize=lsize,labelpad=10)
lim=False
if lim:
    ax4.set_ylim(0,1)
    ax4.set_xlim(xmin=0.0,xmax=Os[-1])
msize=100
ax4.set_xlabel('$w$',labelpad=-5,fontsize=lsize)
ax4.xaxis.set_ticks([0,0.2,0.4])
ax4.scatter(ws,PTs[0,:],color=colours[0],s=msize)
ax4.set_xscale("linear")
letters=['a','b','c','d']
positionxs,positionys=[0.14,0.58],[0.84,0.55]
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]
    if i==1:
        positiony=positiony-0.04
    fig.text(positionx,positiony,'(%s)'%letters[i],fontsize=44)
fig.legend(loc=9,ncol=nw,handles=handles,labels=labelsb,fontsize=48,handlelength=1,labelspacing=0,handletextpad=0.45,columnspacing=0.9)
    

labels=[]

o,D=1e6,0e0
gs=np.array([-8,4,40,-40,-4,8,-20,0,20])
inds_g=[3,6,0,4,7,1,5,8,2]
ng=len(gs)
folder_='gt'
ws_s=[]
T0_remove=True
for folder in os.listdir(folder_):
    if os.path.exists(folder_+'/%s/T1s.txt'%(folder[:4])):
        if T0_remove:
            if abs(np.loadtxt(folder_+'/%s/T1s0.txt'%(folder[:4]))-5e-1)<2e-2:
                ws_s.append(folder[:4])
            else:
                print(folder[:4])
        else:
            ws_s.append(folder[:4])
nw=len(ws_s)
ws=np.zeros(nw)
T1s=np.zeros((nw,ng))
T1s0=np.zeros((nw))
for j in range(nw):
    ws[j]=float(ws_s[j])
    T1s0[j]=np.loadtxt(folder_+'/%3.2f/T1s0.txt'%(ws[j]))
    for k in range(ng):
        T1s[j,inds_g[k]]=np.loadtxt(folder_+'/%3.2f/T1s.txt'%(ws[j]))[k]

ws_temp,T1s_temp,T1s0_temp=np.zeros(nw),np.zeros((nw,ng)),np.zeros(nw)
inds=np.argsort(ws)
for i in range(nw):
    ws_temp[i]=ws[inds[i]]
    T1s_temp[i,:]=T1s[inds[i],:]
    T1s0_temp[i]=T1s0[inds[i]]
ws,T1s,T1s0=ws_temp,T1s_temp,T1s0_temp

cut=True
ind_cut=np.argmin(abs(ws-0.18))
if cut:
    ws=ws[:ind_cut]
    T1s0=T1s0[:ind_cut]
    T1s=T1s[:ind_cut,:]

folder1='ir'
folder2='ir3'
ws_1=np.loadtxt(folder2+'/ws.txt')
nw_1=len(ws_1)
T1s_1_1=np.zeros((nw_1))
for j in range(nw_1):
    T1s_1_1[j]=np.loadtxt(folder1+'/a0s/%3.2f_1.txt'%(ws_1[j]))[1]
T1s_1_g=np.zeros((nw_1,ng))
for i in range(len(gs)):
    for j in range(nw_1):
        T1s_1_g[j,i]=np.loadtxt(folder2+'/T1s/%3.2f_%d.txt'%(ws_1[j],gs[i]))

adjust_cond=False
if adjust_cond:
    deltas=T1s0-0.5*np.ones(nw)
    T1s0=T1s0-deltas
    for i in range(ng):
        T1s[:,i]=T1s[:,i]-deltas

ax5=fig.add_subplot(gspec[4,0])
scatter=True
N=40
linewidth=2.5

markers=['P','v','s','x','D','1','.','p','*','^']
if scatter:
    Nc=200
    ws_in=np.linspace(ws_1[0],ws_1[-1],Nc)
    f0=interp1d(ws_1,T1s_1_1)
    for j in range(ng):
        exec("f%d=interp1d(ws_1,T1s_1_g[:,j])"%(j+1))
    msize=130
    ax5.scatter(ws,T1s0,marker=markers[0],label='$g=1$',color='black',s=msize)
    ax5.plot(ws_in,f0(ws_in),color='black',linewidth=linewidth)
    for i in range(ng):
        ax5.scatter(ws,T1s[:,i],marker=markers[i+1],s=msize,label='$g=%d$'%gs[i])
        exec("ax5.plot(ws_in,f%d(ws_in),linewidth=linewidth)"%(i+1))
        
ax5.set_xlabel('$w$',labelpad=0)
ax5.set_ylabel('$T_1$',rotation=90,labelpad=10)
ax5.set_ylim(0,1)
fig.legend(ncol=3,loc=(0.46,0.11),markerscale=1.5,fontsize=38,labelspacing=0.35,columnspacing=0.18,handletextpad=0.10,borderpad=0,edgecolor=None,frameon=None,framealpha=0)
ax5.set_xlim(0.01,0.18)
ax5.xaxis.set_ticks([0.02,0.08,0.14])
fig.text(positionxs[0],0.26,'(e)',fontsize=44)

for i in range(5):
    exec("ax%d.tick_params(length=6,width=2)"%(i+1))

plt.savefig('fig_3.eps')
