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
#
start = timeit.default_timer()
#
#Define parameters
nbins = 50 #choose number of bins for histogram
onethird = 1.0/3.0
microm = 1.e-6
#
#----Obtain the number density of the particles
pdim = pc.read_pdim()
npar = pdim.npar
dim = pc.read_dim()
nx = dim.nx
ny = dim.ny
nz = dim.nz
param = pc.read_param(quiet=True)
np_swarm0 = param.np_swarm0
n0 = npar*np_swarm0/(nx*ny*nz)
print('n0=',n0)
#----normalization-----
#a0 = param.a0_initdist
a0 = param.ap0[0]
print('a0=',a0)
facSw = a0/n0
#
#
deltaivar = 1
param2 = pc.read_param(param2=True, quiet=True)
ivar0 = np.int(param2.tstart_droplet_coagulation/param2.dsnap)
print('ivar0=', ivar0)
#number density------------------------------------------------------------
ivar1 = len(glob.glob('./data/proc0/PVAR*'))
print('ivar1=', ivar1)
num =np.int((ivar1-ivar0)/deltaivar) 
print('num=', num)
radii_norm = np.zeros((nbins-1,num+1))
prob_norm = np.zeros((nbins-1,num+1))
#
pvar0 = pc.read_pvar('PVAR0')
logbin = np.logspace(np.log10(min(pvar0.ap)), np.log10(max(pvar0.ap)), nbins)
yy, xx = np.histogram(pvar0.ap, bins=logbin, weights=pvar0.npswarm)
norma = sum(yy)/n0
radii_norm[:,ivar0] = xx[1:]
prob_norm[:,ivar0] = yy/norma
del pvar0
#
for ivar in range(ivar0,ivar1,deltaivar):
    print('ivar=', ivar)
    pvar = pc.read_pvar('PVAR'+str(ivar))
    logbin = np.logspace(np.log10(min(pvar.ap)), np.log10(max(pvar.ap)), nbins)
    yy, xx = np.histogram(pvar.ap, bins=logbin, weights=pvar.npswarm)
    num_tmp = np.int((ivar-ivar0)/deltaivar) 
    radii_norm[:,num_tmp] = xx[1:]
    prob_norm[:,num_tmp] = yy/norma
    del pvar, yy, xx
#
with open('./data/radii_norm.json','w') as f:
    json.dump(radii_norm.tolist(), f, sort_keys = True, indent = 4, ensure_ascii = False)
with open('./data/prob_norm.json','w') as f:
    json.dump(prob_norm.tolist(), f, sort_keys = True, indent = 4, ensure_ascii = False)
#
##-----visualization------
#plt.close('all')
#for ivar in range(ivar0,ivar1,deltaivar):
#    num_tmp = np.int((ivar-ivar0)/deltaivar)
#    plt.semilogy(radii_norm[:,num_tmp]/microm, prob_norm[:,num_tmp])
#
stop = timeit.default_timer()
print('elapsed time=',stop-start)