#%% Import packages
import numpy as np # Numpy
import DataProcessing as fun # Function package for analysis
import matplotlib.pylab as plt # Matplitlib

#%% Full file names with paths. Choose one '.data' type file for Chirp analysis
# and one '.oct' type file from a dataset for diffusion quantification.
file_chirp = 'C:/Users/konstantineche/surfdrive/Ganymede/Data/Article '\
'Measurements (B-scan)/Zenodo/Chirp.data' # Chirp file, same for all 
file_diff = 'C:/Users/konstantineche/surfdrive/Ganymede/Data/Article '\
'Measurements (B-scan)/16-04-2021/Raw Data/Diffusion, 2845.oct' # Diffusion 

#%% Setup parameters
dt = 1/36000 # Integration time [s] (fixed parameter)
n_kc = 1.38 # Refractive index [-] (fixed parameter)

#%% Load data, obtain correlation matrices, fit decay rates
Chirp = np.fromfile(file_chirp,np.float32) # Load chirp vector
# Load raw spectral file, process it, obtain correlation matrix g2
[g2, time, dopPhase, sigma, q, z, amp, N_a, _, _] = \
    fun.OCTgetCorrelation(file_diff, Chirp, dt, n_kc, 0, 0, 250)
G = fun.g2_fit(g2, time, 4, 100, 0, 0)[:, 1] # Fit decay rate to g2
D = G/(q**2) # Diffusion coefficient  

#%% Plotting diffusion coefficient (choose indices to cover the channel).
# Channel width is 0.2 mm, total imaging depth is around 1.8 mm in OPL.
i1 = 32 # Choose starting index, min index is 0
i2 = 175 # Choose ending index max index is 1023
fig, axes = plt.subplots(1,1, figsize=(12, 8)) # Plotting diffusion
axes.plot(z[i1:i2]/n_kc*1e3, D[i1:i2]*1e12, 'r', linewidth=2) # D vs. position
axes.set_xlabel('Position [mm]', fontsize=18,weight='bold') # X-label
axes.set_ylabel(r'$\mathbf{D\;[m^2/s \times 10^{-12}]}$', fontsize=18,weight='bold') # Y-label
axes.tick_params(axis='x', labelsize=18) # Tick font size
axes.tick_params(axis='y', labelsize=18) # Tick font size
labels = axes.get_xticklabels() + axes.get_yticklabels(); # Axis label size
[label.set_fontweight('bold') for label in labels]; # Axis label size
# Axis title, containing the mean diffusion coefficient over chosen range
axes.set_title('Diffusion coefficient, D = '+str(np.round(np.mean(D[i1:\
    i2])*1e14)/100)+r'$\mathbf{\,\times 10^{-12}\,mm^2/s}$', fontsize=18,weight='bold');

#%% Note
'''
Find the indices i1 and i2 which give consistent diffusion values
over the given positional range. Then note down the average diffusion 
coefficient and use it in beam waist, angle calibration and flow measurement
scripts as the calibration parameter. i1 and i2 are consistent for same set 
of measurements, i.e. measurements from the same date. They are same for 
A- and B-scan measurements.
'''