'''
This script simulates g1 and g2 from static diffusive particle suspension and
plots the obtained results. Simulation parameters can be set below.    
'''
#%% Import packages
import Data_processing as fun # Import defined functions
import matplotlib.pylab as plt # Import matplotlib
import numpy as np # Import numpy
plt.close('all') # Close all plots

#%% Simulation parameters
T_C = 21 # Temperature
a = 50e-9 # Particle radius [nm]
eta = 1e-3 # Viscocity
n_kc = 1.33 # Refractive index
kB = 1.380649e-23 # Boltzman constant
T_K =T_C + 273.15 # Sample temperature in kelvin
kc = 7066688.789068299 # Wavenumber [1/m]7066688.7890683
D_c = kB*T_K/(6*np.pi*eta*a) # Diffusion coefficient [m^2/s]
N_b = 1000 # Number of time series
SNR = 100 # signal-to-noise ratio
N_a = 4096 # Time series length
dt = 1/5500 # Integration time [s]
N_fit = 100 # Number of fit points
N_p = 1000 # Number of simulated particles
N_mix = 100 # Number of autocorrelation functions for mixing

#%% Simulate and fit autocorrelations due to particle diffusion (no flow)
[standard_D, standard_A, mixed_D, mixed_A, mu_D, mu_A, g1_mu, g2_mu, e1, e2, \
     e1_mix, e2_mix, g1_var, g2_var, time] = fun.SimulationDiff(N_p, N_a, N_b,\
     N_fit, N_mix, D_c, dt, kc, n_kc, SNR)

#%% Plot simulated autocorrelation functions, variances and error correlations
fig, axes = plt.subplots(2,2, figsize=(16, 10)) 
axes[0,0].plot(time*1e3, g1_mu, color='blue', label='Average $\Re(g_1)$', linewidth=2)
axes[0,0].plot(time*1e3, g2_mu, color='red',label='Average $g_2$', linewidth=2)
axes[0,0].set_ylabel('Signal ACF', fontsize=18)
axes[0,0].set_xlabel('Lag time [ms]', fontsize=18)
axes[0,0].tick_params(axis='x', labelsize=18)
axes[0,0].tick_params(axis='y', labelsize=18) 
axes[0,0].legend(frameon=False, fontsize=18, loc='best')
axes[0,0].set_xlim(0, 4)
axes[0,0].set_ylim(-0.02, 1.02)
axes[0,1].plot(time*1e3, g1_var, color='blue', label='$\Re(g_1)$', linewidth=2)
axes[0,1].plot(time*1e3, g2_var, color='red',label='$g_2$', linewidth=2)
axes[0,1].set_ylabel('ACF variance', fontsize=18)
axes[0,1].set_xlabel('Lag time [ms]', fontsize=18)
axes[0,1].tick_params(axis='x', labelsize=18)
axes[0,1].tick_params(axis='y', labelsize=18) 
axes[0,1].legend(frameon=False, fontsize=18, loc='best')
axes[0,1].set_xlim(0, 4)
axes[0,1].set_ylim(0, 0.001)
axes[1,0].plot(time*1e3, e1, color='blue', label='Original $\Re(g_1)$', linewidth=2)
axes[1,0].plot(time*1e3, e2, color='red',label='Original $g_2$', linewidth=2)
axes[1,0].set_ylabel('Error ACF', fontsize=18)
axes[1,0].set_xlabel('Lag time [ms]', fontsize=18)
axes[1,0].tick_params(axis='x', labelsize=18)
axes[1,0].tick_params(axis='y', labelsize=18) 
axes[1,0].legend(frameon=False, fontsize=18, loc='best')
axes[1,0].set_xlim(0, 4)
axes[1,0].set_ylim(-0.02, 1.02)
axes[1,1].plot(time*1e3, e1_mix, color='blue', label='Mixed $\Re(g_1)$', linewidth=2)
axes[1,1].plot(time*1e3, e2_mix, color='red',label='Mixed $g_2$', linewidth=2)
axes[1,1].set_ylabel('Error ACF', fontsize=18)
axes[1,1].set_xlabel('Lag time [ms]', fontsize=18)
axes[1,1].tick_params(axis='x', labelsize=18)
axes[1,1].tick_params(axis='y', labelsize=18) 
axes[1,1].legend(frameon=False, fontsize=18, loc='best')
axes[1,1].set_xlim(0, 4)
axes[1,1].set_ylim(-0.02, 1)
plt.tight_layout()
#fig.savefig('Simulated autocorrelations, diffusion.pdf') # Save file

#%% Plot simulated and fitted diffusion coefficient distributions
fig, axes = plt.subplots(1,2, figsize=(16, 5)) 
axes[0].hist(mixed_D[:, 0]/D_c, color='red', density=True, label='Mixed $D$')
axes[0].hist(standard_D[:, 0]/D_c, color='blue', density=True, \
    alpha=0.6, label='Standard $D$')
axes[0].set_ylabel('PDF, $\Re(g_1)$', fontsize=18)
axes[0].set_xlabel('$D/D_0$', fontsize=18)
axes[0].tick_params(axis='x', labelsize=18)
axes[0].tick_params(axis='y', labelsize=18) 
axes[0].legend(frameon=False, fontsize=18, loc='best')
axes[1].hist(mixed_D[:, 1]/D_c, color='red', density=True, label='Mixed $D$')
axes[1].hist(standard_D[:, 1]/D_c, color='blue', density=True, \
    alpha=0.6, label='Standard $D$')
axes[1].set_ylabel('PDF, $g_2$', fontsize=18)
axes[1].set_xlabel('$D/D_0$', fontsize=18)
axes[1].tick_params(axis='x', labelsize=18)
axes[1].tick_params(axis='y', labelsize=18) 
axes[1].legend(frameon=False, fontsize=18, loc='best')
plt.tight_layout()
#fig.savefig('Simulated distributions, diffusion.pdf') # Save file