import os 
import sys
import numpy as np
from IPython import embed
import pencil as pc
import timeit
import json
import math
import glob
from scipy.io.idl import readsav
from scipy.optimize import curve_fit
import matplotlib
import matplotlib.pyplot as plt
from matplotlib import rc
#matplotlib.rcParams['text.latex.preamble'] = [r'\boldmath']
plt.rc('text', usetex=True)
plt.ion()
import pencil as pc
#
start = timeit.default_timer()
#
#------curve fitting---------
def fitFunc(t, a, b):
    return a*np.exp(b*t)
#
#-----Define parameters-----
BM = 0.941 #Bernhard's correction factor
good_t = 5. #seconds
cutoff=3
siz_text = 16
#
#---fixed Re
run0='SW512condens0_coag_grav_turb_shima_f2em2_ngp_a24_L025'
run1='SW512condens0_coag_grav_turb_shima_f1p4em2_ngp_a24_L03'
run2='SW512condens0_coag_grav_turb_shima_f1em2_ngp_a24_L037b_random'
run3='SW512condens0_coag_grav_turb_shima_f7p2em3_ngp_a24_L044'
#
ts0 = pc.read_ts(datadir='../'+run0+'/data')
ts1 = pc.read_ts(datadir='../'+run1+'/data')
ts2 = pc.read_ts(datadir='../'+run2+'/data')
ts3 = pc.read_ts(datadir='../'+run3+'/data')
#
#-----Get n0 and a0 for normalization------------
pdim = pc.read_pdim(datadir='../'+run1+'/data')
npar=pdim.npar
dim = pc.read_dim(datadir='../'+run1+'/data')
nx=dim.nx
ny=dim.ny
nz=dim.nz
param1 = pc.read_param(datadir='../'+run1+'/data', quiet=True)
np_swarm0=param1.np_swarm0
a00=param1.a0_initdist
n0=npar*np_swarm0/(nx*ny*nz)
boxsize=abs(param1.xyz0*2)
boxsize=boxsize[0]
Ntot1=n0*boxsize**3
print,'n0=',n0
print,'a00=',a00
print,'Ntot1=',Ntot1
print,'Ntot/Ns=',Ntot1/npar
#
param1 = pc.read_param(datadir='../'+run0+'/data', quiet=True)
boxsize=abs(param1.xyz0*2)
boxsize=boxsize[0]
Ntot0=n0*boxsize**3
#
param1 = pc.read_param(datadir='../'+run1+'/data', quiet=True)
boxsize=abs(param1.xyz0*2)
boxsize=boxsize[0]
Ntot1=n0*boxsize**3
#

param1 = pc.read_param(datadir='../'+run2+'/data', quiet=True)
boxsize=abs(param1.xyz0*2)
boxsize=boxsize[0]
Ntot2=n0*boxsize**3
#
param1 = pc.read_param(datadir='../'+run3+'/data', quiet=True)
boxsize=abs(param1.xyz0*2)
boxsize=boxsize[0]
Ntot3=n0*boxsize**3
#
#
param = pc.read_param(datadir='../'+run1+'/data', param2=True, quiet=True)
nu=param.nu*1.0
#
#---get Kolmogorov scales----
good0 = np.where(ts0.t>good_t)[0][0]
good1 = np.where(ts1.t>good_t)[0][0]
good2 = np.where(ts2.t>good_t)[0][0]
good3 = np.where(ts3.t>good_t)[0][0]
#good4 = np.where(ts4.t>good_t)[0][0]
#
epsK_0 = np.mean(ts0.epsK[good0:])
epsK_1 = np.mean(ts1.epsK[good1:])
epsK_2 = np.mean(ts2.epsK[good2:])
epsK_3 = np.mean(ts3.epsK[good3:])
#epsk_4 = np.mean(ts4.epsK[good4:])
#
tauK_0=(nu/epsK_0)**(1./2)
tauK_1=(nu/epsK_1)**(1./2)
tauK_2=(nu/epsK_2)**(1./2)
tauK_3=(nu/epsK_3)**(1./2)
#tauK_4=(nu/epsK_4)**(1./2)
#
uK_0=(nu*epsK_0)**(1./4)
uK_1=(nu*epsK_1)**(1./4)
uK_2=(nu*epsK_2)**(1./4)
uK_3=(nu*epsK_3)**(1./4)
#uK_4=(nu*epsK_4)**(1./4)
#
ncoagpm0=(ts0.ncoagpm/ts0.dt)*npar/Ntot0
ncoagpm1=(ts1.ncoagpm/ts1.dt)*npar/Ntot1
ncoagpm2=(ts2.ncoagpm/ts2.dt)*npar/Ntot2
ncoagpm3=(ts3.ncoagpm/ts3.dt)*npar/Ntot3
#
colNorm_0=n0*(2*a00)**3*tauK_0**(-1)/BM
colNorm_1=n0*(2*a00)**3*tauK_1**(-1)/BM
colNorm_2=n0*(2*a00)**3*tauK_2**(-1)/BM
colNorm_3=n0*(2*a00)**3*tauK_3**(-1)/BM
#colNorm_0=1
#colNorm_1=1
#colNorm_2=1
#colNorm_3=1
#----remove the artificial factor from the volume, Ntot2 has the largest volume---
v_factor0=Ntot0/Ntot3
v_factor1=Ntot1/Ntot3
v_factor2=Ntot2/Ntot3
v_factor3=Ntot3/Ntot3
#--fixed Re---
siz = 15
text_siz = 20
siz_tit = 20
nevery = 100
plt.close()
plt.rc('font', family='Times', size=siz)
plt.rc('xtick', labelsize=siz)
plt.rc('ytick', labelsize=siz)
#plt.ylim([5.e-3,3.e-1])
#plt.ylim([5.e0,3.e2])
plt.ylim([5,2.e2])
plt.xlim([0,1.e1])
plt.semilogy(ts0.t-cutoff,ncoagpm0/colNorm_0*v_factor0,'-r', label=r'$\bar{\epsilon}=0.04\,\rm{m}^2\rm{s}^{-3}$')
plt.semilogy(ts1.t-cutoff,ncoagpm1/colNorm_1*v_factor1,'-g', label=r'$0.02$')
plt.semilogy(ts2.t-cutoff,ncoagpm2/colNorm_2*v_factor2,'-k', label=r'$0.01$')
plt.semilogy(ts3.t[ncoagpm3<10]-cutoff,ncoagpm3[ncoagpm3<10]/colNorm_3*v_factor3,'-b', label=r'$0.005$')
plt.xlabel(r'$t\,[\rm{s}]$', fontsize=siz_tit)
plt.ylabel(r'$\bar{R_c}\tau_{\eta}/\left[n_0\left<(r_i+r_j)^3\right>_{\rm ini}\right]$', fontsize=siz_tit)
leg = plt.legend(fontsize=siz_text, loc=2, bbox_to_anchor=(0.01,1.03), handletextpad=1., frameon=False)
def color_legend_texts(leg):
    """Color legend texts based on color of corresponding lines"""
    for line, txt in zip(leg.get_lines(), leg.get_texts()):
        txt.set_color(line.get_color())  

color_legend_texts(leg)
#
plt.tight_layout(pad=0)
plt.subplots_adjust(left=0.15, right=.98, top=.975, bottom=0.13)
#
plt.show()
filename=sys.argv[0]
print(filename)
basename=os.path.basename(filename)
base=basename.split('.')[0]
plt.savefig('python.eps',dpi=300)
print('!mv'+' '+'python.eps'+' '+'~/tex/nils/particles_coag/gravity/fig/'+base+'.eps')
#
stop = timeit.default_timer()
print('elapsed time=',stop-start)