'''
This script analyses raw interferrometric file generated by Thorlabs Ganymede
OCT system. OCT data is obtained from a static (diffusing) particle suspension.
N_a is the time series length (A-scans), N_b is the number of repeats (B-scans).
Autocorrelations are computed along the A-scan direction. Averaging, mixing, 
and statistical analysis is performed along the B-scan direction. Results are
plotted at the end. No calibration measurements are required. Only the 
measurement (.oct) and chirp interpolation (Chirp.data) files are required.
The beam remains stationary during the measurement.   
'''
#%% Import packages
from joblib import Parallel, delayed # Package for parallel processing
import matplotlib.pylab as plt # Import Matplitlib
from functools import partial # Package for partial function declaration
import Data_processing as fun # Function package for analysis
import multiprocessing # Package for parallel processing
import numpy as np # Import Numpy

#%% Load files
file_path = '//tudelft.net/staff-umbrella/kote/03-03-2023/' # Path
file_name = 'Diffusion_03032023.oct' # Diffusion filename
file_chirp = 'Chirp.data' # Chirp file
Chirp = np.fromfile(file_path+file_chirp,np.float32) # Load chirp vector

#%% Setup parameters
T_C = 21 # Temperature in Celsius
dt = 1/5500 # Integration time [s] 
N_a = 4096 # Time series length (A-scans)
T_K = T_C + 273.15 # Sample temperature in kelvin
kB = 1.380649e-23 # Boltzman constant
N_max = 50 # Time delays for fitting
n_kc = 1.33 # Refractive index 
N_b = 1100 # Number of repeats (B-scans)
    
#%% Run parallel, diffusion measurement
inputs = range(N_b) # Number of repeats (B-scans)
time = np.linspace(0, (N_max)*dt, N_max+1) # Time lag vector for g1 and g2
g1 = np.zeros((N_b, N_max+1, 1024)) # g1  preallocation
g2 = np.zeros((N_b, N_max+1, 1024)) # g2 preallocation
g1_mix = np.zeros_like(g1) # Mixed g1 preallocation
g2_mix = np.zeros_like(g2) # Mixed g2 preallocation
num_cores =  multiprocessing.cpu_count() # Number of cores 
# Run multicore processing for obtaining autocorrelation functions
processInput = partial(fun.OCTgetACF, file_path+file_name, Chirp, 300, N_max)
output = Parallel(n_jobs=num_cores)(delayed(processInput)(i) for i in inputs)
for i in range(N_b): # Fill autocorrelation arrays
    g1[i, :, :] = np.real(output[i][0]) # Take real part
    g2[i, :, :] = np.real(output[i][1]) # Take real part
for i in range(N_max+1): # Mix autocorrelation functions (fully)
    g1_mix[:, i, :] = np.roll(g1[:, i, :], i, axis=0)
for i in range(N_max+1):
    g2_mix[:, i, :] = np.roll(g2[:, i, :], i, axis=0)
g1_mu = np.mean(g1, axis=0) # Average g1
g2_mu = np.mean(g2, axis=0) # Average g2
g1_var = np.var(g1, axis=0) # Variance in g1
g2_var = np.var(g2, axis=0) # Variance in g2
kc = output[0][2] # Central spectral wavenumber
z = output[0][3] # OPL vector
q = 2*kc*n_kc # Central scattering wavenumber
x = z/n_kc # Depth  vector    

#%% Fit parameters to average autocorrelation functions
fit_par = fun.OCTfitDiff(g1_mu[None, :, :], g2_mu[None, :, :], time, N_max, 0)
D1_mu = fit_par[:, 1]/q**2 # D from average g1
D2_mu = fit_par[:, 3]/q**2 # D from average g2
A1_mu = fit_par[:, 0] # A from average g1
A2_mu = fit_par[:, 2] # A from average g2

#%% Fit parameters in parallel (standard ACF)
num_cores = multiprocessing.cpu_count() # Number of cores 
processInput = partial(fun.OCTfitDiff, g1, g2, time, N_max) # Multicore fit
output = Parallel(n_jobs=num_cores)(delayed(processInput)(i) for i in inputs)
A1_standard = np.zeros((N_b, 1024)) # Fit amplitude for standard g1
A2_standard = np.zeros((N_b, 1024)) # Fit amplitude for standard g2
D1_standard = np.zeros((N_b, 1024)) # Diffusion coefficient for standard g1
D2_standard = np.zeros((N_b, 1024)) # Diffusion coefficient for standard g2
for i in range(N_b): # Obtain parameters from multicore processing
    A1_standard[i, :] = output[i][:, 0]
    A2_standard[i, :] = output[i][:, 2]
    D1_standard[i, :] = output[i][:, 1]/q**2
    D2_standard[i, :] = output[i][:, 3]/q**2

#%% Fit parameters in parallel (mixed ACF)
num_cores = multiprocessing.cpu_count() # Number of cores 
processInput = partial(fun.OCTfitDiff, g1_mix, g2_mix, time, N_max) # Multicore fit
output = Parallel(n_jobs=num_cores)(delayed(processInput)(i) for i in inputs)
A1_mix = np.zeros((N_b, 1024)) # Fit amplitude for standard g1
A2_mix = np.zeros((N_b, 1024)) # Fit amplitude for standard g2
D1_mix = np.zeros((N_b, 1024)) # Diffusion coefficient for standard g1
D2_mix = np.zeros((N_b, 1024)) # Diffusion coefficient for standard g2
for i in range(N_b): # Obtain parameters from multicore processing
    A1_mix[i, :] = output[i][:, 0]
    A2_mix[i, :] = output[i][:, 2]
    D1_mix[i, :] = output[i][:, 1]/q**2
    D2_mix[i, :] = output[i][:, 3]/q**2
    
#%% Calculate standard deviations and average coefficients
sigma_D1_standard = np.std(D1_standard, axis=0) # sigma_D, standard g1
sigma_D2_standard = np.std(D2_standard, axis=0) # sigma_D, standard g2
sigma_D1_mix = np.std(D1_mix, axis=0) # sigma_D, mixed g1
sigma_D2_mix = np.std(D2_mix, axis=0) # sigma_D, mixed g2
mean_D1_standard = np.mean(D1_standard, axis=0) # mean D, standard g1
mean_D2_standard = np.mean(D2_standard, axis=0) # mean D, standard g2
mean_D1_mix = np.mean(D1_mix, axis=0) # mean D, mixed g1
mean_D2_mix = np.mean(D2_mix, axis=0) # mean D, mixed g2

#%% Plotting mean diffusion coefficients, standard deviations, fit amplitudes
fig, axes = plt.subplots(1,3, figsize=(24, 5)) 
axes[0].plot(x*1e3, mean_D1_standard*1e12, color='blue', linewidth=2, \
    label='Standard $\Re(g_1)$', zorder=0)
axes[0].plot(x*1e3, D1_mu*1e12, color='cyan', linewidth=2, \
    label='Average $\Re(g_1)$', zorder=0)
axes[0].plot(x*1e3, mean_D1_mix*1e12, color='orange', linewidth=2, \
    label='Mixed $\Re(g_1)$', zorder=0)
axes[0].plot(x*1e3, mean_D2_standard*1e12, color='red', linewidth=2, \
    label='Standard $g_2$', zorder=-1)
axes[0].plot(x*1e3, D2_mu*1e12, color='magenta', linewidth=2, \
    label='Average $g_2$', zorder=-1)
axes[0].plot(x*1e3, mean_D2_mix*1e12, color='green', linewidth=2, \
    label='Mixed $g_2$', zorder=-1)
axes[0].set_ylabel('Average $D$ [$\mu$m$^2$/s]', fontsize=18) 
axes[0].set_xlabel('Depth [mm]', fontsize=18)
axes[0].tick_params(axis='x', labelsize=18)
axes[0].tick_params(axis='y', labelsize=18)  
axes[0].set_xlim([x[0]*1e3, x[-1]*1e3]) 
axes[0].set_ylim([3.8, 4.8])
axes[0].legend(frameon=False, fontsize=18, loc='best', labelspacing=0.25)
axes[1].plot(x*1e3, sigma_D1_standard*1e12, color='blue', linewidth=2, \
    label='Standard $\Re(g_1)$')
axes[1].plot(x*1e3, sigma_D2_standard*1e12, color='red', linewidth=2, \
    label='Standard $g_2$')
axes[1].plot(x*1e3, sigma_D1_mix*1e12, color='orange', linewidth=2, \
    label='Mixed $\Re(g_1)$')
axes[1].plot(x*1e3, sigma_D2_mix*1e12, color='green', linewidth=2, \
    label='Mixed $g_2$')
axes[1].set_ylabel('$\sigma_D$ [$\mu$m$^2$/s]', fontsize=18) 
axes[1].set_xlabel('Depth [mm]', fontsize=18)
axes[1].tick_params(axis='x', labelsize=18)
axes[1].tick_params(axis='y', labelsize=18)  
axes[1].set_xlim([x[0]*1e3, x[-1]*1e3]) 
axes[1].set_ylim([0, 2.0])
axes[1].legend(frameon=False, fontsize=18, loc='best',  labelspacing=0.25)
axes[2].plot(x*1e3, A1_mu, color='blue', linewidth=2, \
    label='Average $\Re(g_1)$', zorder=0)
axes[2].plot(x*1e3, A2_mu, color='red', linewidth=2, \
    label='Average $g_2$', zorder=-1)
axes[2].set_ylabel('ACF amplitude', fontsize=18) 
axes[2].set_xlabel('Depth [mm]', fontsize=18)
axes[2].tick_params(axis='x', labelsize=18)
axes[2].tick_params(axis='y', labelsize=18)  
axes[2].set_xlim([x[0]*1e3, x[-1]*1e3]) 
axes[2].set_ylim([0, 1.2])
axes[2].legend(frameon=False, fontsize=18, loc='best')
plt.tight_layout()
#fig.savefig('Processed data, diffusion.pdf') # Save file