import os 
import sys
import numpy as np
from IPython import embed
import pencil as pc
import timeit
import json
import math
import glob
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()
#
start = timeit.default_timer()
#
#Define parameters
ivar0 = 1
deltaivar = 2
microm = 1e-6
#
siz_text = 18
siz_tit = 25
siz_tic = 20 
#
speed = 1
#
run0='../SW512condens0_coag_grav0_turb_shima_f2em2_ngp_a24_L025_lucky2/data/'
run1='../SW512condens0_coag_grav0_turb_shima_f2em2_ngp_a24_L025_lucky2_fullTM/data/'
#
with open(run0+'radii_norm.json','rb') as f:
    radii_norm0 = np.array(json.loads(f.read()))
with open(run0+'prob_norm.json','rb') as f:
    prob_norm0 = np.array(json.loads(f.read()))
ivar1 = min(radii_norm0.shape)
#
with open(run1+'rbins_all.json','rb') as f:
    radii_norm1 = np.array(json.loads(f.read()))
with open(run1+'hist_all.json','rb') as f:
    prob_norm1 = np.array(json.loads(f.read()))
#
#----normalization------------
pdim = pc.read_pdim(datadir=run0)
npar = pdim.npar
dim = pc.read_dim(datadir=run0)
nx = dim.nx
ny = dim.ny
nz = dim.nz
param = pc.read_param(datadir=run0, quiet=True)
np_swarm0 = param.np_swarm0
n0 = npar*np_swarm0/(nx*ny*nz)
print('n0=',n0)
a0 = param.a0_initdist
n0_Sw = npar*np_swarm0/(nx*ny*nz)
print('n0_Sw=',n0_Sw)
print('a0=',a0)
facSw = a0/n0_Sw/microm
#
#Plot-----------------------------
plt.close('all')
plt.figure()
plt.xlabel(r'$r\,[\mu\rm{m}]$', fontsize=siz_tit)
plt.ylabel(r'$fr_{\rm ini}/n_0$', fontsize=siz_tit)
plt.ylim([1.e-6,1.e2])
plt.xlim([5.e0,6.5e1])
plt.rc('xtick',labelsize=siz_tic)
plt.rc('ytick',labelsize=siz_tic)
plt.rc('ytick',labelsize=siz_tic)
plt.rc('font', family='family')
plt.tight_layout(pad=0)
#plt.subplots_adjust(left=0.12, right=.98, top=.98, bottom=0.12)
plt.subplots_adjust(left=0.16, right=.98, top=.97, bottom=0.16)
colors = ['w','k', 'w', 'b', 'w', 'g', 'w', 'c', 'w', 'm']
for i in range(ivar0,ivar1,deltaivar):
  print(i)
  dlnad = math.log(radii_norm0[3,i])-math.log(radii_norm0[2,i])
  xx0 = radii_norm0[:,i]
  yy0 = facSw*prob_norm0[:,i]/(dlnad*xx0)
  plt.semilogy(xx0,yy0,'.', color=colors[i])
  #
  dlnad = math.log(radii_norm1[3,i])-math.log(radii_norm1[2,i])
  xx1 = radii_norm1[:,i]
  yy1 = facSw*prob_norm1[:,i]/(dlnad*xx1)
  plt.semilogy(xx1,yy1,':', color=colors[i], linewidth=2)
#
##Empty plot, just for labels
plt.semilogy(0,0,'.k',label=r'turbulence')
plt.semilogy(0,0,':k',label=r'turbulence, FullDev, well-mixed')
plt.text(8.e0, 1.e-5, r'$t=1\,\rm{s}$', color='k', fontsize=siz_text)
plt.text(2.2e1, 1.e-5, r'$3$', color='b', fontsize=siz_text)
plt.text(2.7e1, 1.e-5, r'$5$', color='g', fontsize=siz_text)
plt.text(5.e1, 1.e-5, r'$7$', color='c', fontsize=siz_text)
plt.text(6.e1, 1.e-5, r'$9$', color='m', fontsize=siz_text)
plt.legend(fontsize=siz_text, loc=1, bbox_to_anchor=(1.01,1.01))
#plt.yticks([1.e-6,1.e-4,1.e-2,1.e0,1.e2])
#
# save the figure
filename=sys.argv[0]
print(filename)
basename=os.path.basename(filename)
base=basename.split('.')[0]
current_folder_path, current_folder_name = os.path.split(os.getcwd())
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)