#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Jul  9 22:17:37 2019

@author: ken
"""

import matplotlib.pyplot as plt 
import pandas as pd
import numpy as np
from scipy.signal import savgol_filter
import QCM_functions as qcm
import scipy.optimize as optimize

water = {'drho':np.inf, 'grho3':1e8, 'phi':90}
boundary = 'water'

def df_bulk(n, grho3, phi):
    film={'drho':np.inf, 'grho3':grho3, 'phi':phi}
    layers={'film':film}
    
    return qcm.calc_delfstar(n, layers, 'SLA')

def phi(grho3):
    # phi as a functiono f grho3 for vgp plot
    log_grhomin = np.log(grhomin)
    log_grhomax = np.log(grhomax)
    log_grho3 = np.log(grho3)
    return 90*((log_grhomax-log_grho3)/(log_grhomax-log_grhomin))


def calc_df(n, drho, grho3, phi):
    # convert varibles to scalars that have been somehow converted to 
    # length one vectors
    if not np.isscalar(drho):
        drho = drho[0]
    if not np.isscalar(grho3):
        grho3 = grho3[0]
    if not np.isscalar(phi):
        phi = phi[0]
    film={'drho':drho, 'grho3':grho3, 'phi':phi}
    layers={'film':film}
    if boundary == 'water':
        layers['overlayer']=water
    return qcm.calc_delfstar(n, layers, 'SLA')

def calc_df_diff(n1, n2, drho, grho3, phi):
    # difference between df/n at n2 n1
    return (calc_df(n2, drho, grho3, phi).real/n2 - 
               calc_df(n1, drho, grho3, phi).real/n1)
    
    
grhomin = 1e8 # grho3 where phi=90
grhomax = 1e12 # grho3 where phi=0

# put the excel data into dataframes.  The read process is a bit slow, so
# you don't want run this code section every time you change something
# in the analysis
dataframe={}  
dataframe['good'] = pd.read_excel('20190522_JoVE_collagenflow.xlsx')
dataframe['bad'] = pd.read_excel('20190709_JoVE_collagenflow_baddata.xlsx')


#%%  Figure 3
# bulid a 'data' dictionary that contains the data
# in a form that is more easy for us to manipulate
data = {}; fig = {}; ax = {}; axD = {}
tref = {'good':0.6, 'bad':0.18}
plt.close ('all')

dataset=['good', 'bad']
titles = [['A) equilibrated', 'B) equilibrated'], ['C) non-equilibrated',
    'D) non-equilibrated']]
fig, ax = plt.subplots(2,2, figsize=(7,7))
for k in [0, 1]:
    data[dataset[k]] = {'t':0, 'delfstar':{}, 'delfstar_smooth':{}}   
    # convert time to hours
    data[dataset[k]]['t'] = np.array(dataframe[dataset[k]]['Time_1 (s)']/3600)
    for n in [1, 3, 5, 7, 11, 13]:
        fstr = 'f' + str(n) + '_1 (Hz)'
        dstr = 'D' + str(n) + '_1 (ppm)'
        data[dataset[k]]['delfstar'][n] = np.array(dataframe[dataset[k]][fstr] +
                    1j*dataframe[dataset[k]][dstr]*5*n/2)
        # reference to appropriate time
        tref_idx = qcm.find_nearest_idx(tref[dataset[k]], data[dataset[k]]['t'])
        data[dataset[k]]['delfstar'][n] = (data[dataset[k]]['delfstar'][n] -
            data[dataset[k]]['delfstar'][n][tref_idx])
             
    t = data[dataset[k]]['t']

    for n in [3, 5]:
        df = data[dataset[k]]['delfstar'][n].real
        dg = data[dataset[k]]['delfstar'][n].imag
        # a little smoothing doesn't hurt:
        df = savgol_filter(df, 99, 1)
        dg = savgol_filter(dg, 99, 1)
        # write these smoothed values back into the data dictionary so 
        # they get used by the solution functions in Figure 4
        data[dataset[k]]['delfstar_smooth'][n] = df + 1j*dg
        ax[k][0].plot(t, df/n, '-', label=r'n='+str(n))
        ax[k][1].plot(t, dg/n, '-', label=r'n='+str(n))

    for m in [0, 1]:
        ax[k][m].set_xlabel('t (hrs.)')
        ax[k][m].set_title(titles[k][m])
        ax[k][m].legend()
    ax[k][0].set_ylabel(r'$\Delta f/n$ (Hz)')
    ax[k][1].set_ylabel(r'$\Delta \Gamma /n$ (Hz)')
    ax[k][0].set_ylim(top=10)
    ax[k][1].set_ylim(bottom=-10)
    # add axis D axis
    qcm.add_D_axis(ax[k][1])
    
    fig.tight_layout()
    fig.savefig('Fig3.pdf')
    
#%% Figure 4
#calculate properties from the equilibrated data
sample = {}; parms = {'make_legend':False}
plt.close('all')

# calculate solutions for some range of time points
t_for_calc = np.linspace(2, 21.8, 30)
dataset = 'good'
idx_for_calc = qcm.find_nearest_idx(t_for_calc, data['good']['t'])
sample['xdata'] = data[dataset]['t'][idx_for_calc]

delfstar={}
for i in np.arange(len(idx_for_calc)):
    delfstar[i]={}
    for n in [3, 5]:
        delfstar[i][n] = data[dataset]['delfstar_smooth'][n][idx_for_calc[i]]
        
sample['delfstar'] = delfstar
sample['nhplot'] = [3, 5]
sample['xlabel'] = 't (hrs.)'
water = {'drho':np.inf, 'grho3':1e8, 'phi':90}
sample['overlayer'] = water
sample['prop_guess']={'drho':8e-5, 'grho3':3e8, 'phi':60}

# obtain 355 solution
sample['nhcalc'] = ['355']
power=qcm.solve_from_delfstar(sample, parms)
power['propfig']['drho_ax'].set_title('(A)')
power['propfig']['grho3_ax'].set_title('(B)')
power['propfig']['phi_ax'].set_title('(C)')

# format the |G*|rho axis 
power['propfig']['grho3_ax'].ticklabel_format(axis='y', style='sci', 
     scilimits=(0,0), useMathText=True)

# save the plot
power['propfig']['figure'].tight_layout()
power['propfig']['figure'].savefig('Fig4.pdf')


# extract results for comparison plot (Figure 6)
drho_G = power['results']['355']['film']['drho']
grho3_G = power['results']['355']['film']['grho3']
phi_G = power['results']['355']['film']['phi']
t_G = power['xdata']

#%% Figure 5
#Read Calculated values from Dfind into a dataframe
plt.close('all')
rho = 1000 #  density in kg/m^3 assumed by Dfind
mass_raw = pd.read_csv('Viscoelastic_mass.csv',
                   sep='\t', header=1, names=['t (s)', 'M (ng/cm^2)'], 
                   dtype=np.float64, na_values=['Failed'])

gprime_raw = pd.read_csv('Viscoelastic_elasticmodulus.csv',
         sep='\t', header=1, usecols=[1], names=['gprime (kPa)'], dtype=np.float64,
         na_values=['Failed'])

viscosity_raw = pd.read_csv('Viscoelastic_viscosity.csv',
         sep='\t', header=1, usecols=[1], names=['eta (micro Pa-s'], dtype=np.float64,
         na_values=['Failed'])

dfind_raw = mass_raw.join([gprime_raw, viscosity_raw])

# eliminate all the values where a solution was not found, and which 
# now contain nan values for the properties
dfind_raw = dfind_raw[~np.isnan(dfind_raw['gprime (kPa)'])]

# now convert to arrays with time in hours, other properties in SI units
t_D = np.array(dfind_raw['t (s)'])/3600
gprime_D = 1000*np.array(dfind_raw['gprime (kPa)'])
gdprime_D = 2*np.pi*3*5e6*np.array(dfind_raw['eta (micro Pa-s'])/1e6

drho_D = np.array(dfind_raw['M (ng/cm^2)'])/1e8
grho3_D = 1100*(gprime_D**2+gdprime_D**2)**(0.5)
phi_D = (180/np.pi)*np.arctan(gdprime_D/gprime_D)

# determine data to plot
t_for_plot = np.linspace(3, 21.8, 30)
idx_D = qcm.find_nearest_idx(t_for_plot, t_D)

# now make the plots
QCMD_props = qcm.make_prop_axes('Dfind_props', 't (hrs.)')

# replot data from Figure 4
QCMD_props['drho_ax'].plot(t_G, drho_G*1000, '-+', label=r'$\Gamma$')
QCMD_props['grho3_ax'].plot(t_G, grho3_G/1000,'-+', label=r'$\Gamma$')
QCMD_props['phi_ax'].plot(t_G, phi_G, '-+', label=r'$\Gamma$')

QCMD_props['drho_ax'].set_title('(A)')
QCMD_props['grho3_ax'].set_title('(B)')
QCMD_props['phi_ax'].set_title('(C)')

# plot the Dfind data
QCMD_props['drho_ax'].plot(t_D[idx_D], drho_D[idx_D]*1000, '-+', label=r'$D$')
QCMD_props['grho3_ax'].plot(t_D[idx_D], grho3_D[idx_D]/1000,'-+', label=r'$D$')
QCMD_props['phi_ax'].plot(t_D[idx_D], phi_D[idx_D], '-+', label=r'$D$')

# format the rho|G*| axis 
QCMD_props['grho3_ax'].ticklabel_format(axis='y', style='sci', 
     scilimits=(0,0), useMathText=True)

# add the legends
QCMD_props['drho_ax'].legend(loc='best')
QCMD_props['grho3_ax'].legend(loc='best')
QCMD_props['phi_ax'].legend(loc='best')

QCMD_props['figure'].tight_layout()

# save the plot
QCMD_props['figure'].savefig('Fig5.pdf')

#%% # Figure 6 - modified van Gurp Palmen plot of polyelectrolyte data
plt.close('all')

fig, ax = plt.subplots(1, 1, figsize=(3,3))
ax.set_xlabel(r'$|G_3^*|\rho$ (Pa$\cdot$g/cm$^3$)')
ax.set_ylabel(r'$\phi_3$ (deg.)')

# this is the data from Kazi's vgp plot
phi_vgp = np.array([86.6, 84.81, 76.51, 72.01, 68.06, 54.6, 52.83, 
                37.94, 31.33, 23.75, 18, 13.3, 12])
pG_vgp = 1000*np.array([2.5e5, 3.02e5, 4.367e5, 6.044e5, 1.17e6, 2.889e6, 4.14e6, 
               2.36e7, 3.35e7, 7.128e7, 1.442e8, 2.535e8, 3.375e8])
    
ax.semilogx(pG_vgp/1000, phi_vgp, '+', label=r'PEC ($\Gamma$)')
ax.plot(6.7e5, 74, 'o', label='Collagen (D)')
ax.plot([grhomin/1000, grhomax/1000], [90, 0], '-')
plt.yticks(ticks=[0, 15, 30, 45, 60, 75, 90])
ax.legend()

fig.tight_layout()
fig.savefig('Fig6.pdf')

#%% Figure 1
# regime map showing where different information can be obtained
plt.close('all')
npts1 = 200

def max_gamma(n,g3):
    # returns max. dissipation at any thickness, with corresponding thickness
    lam_rho = qcm.calc_lamrho(n, g3, phi(g3))
    soln = optimize.minimize_scalar(lambda x: -df_bulk(n, 1e8, 90).imag -
                                    calc_df(n, x, g3, phi(g3)).imag,
                                    bounds=[0, 0.3*lam_rho], method='bounded')
    return -soln['fun'], soln['x']

def min_gamma(n,g3):
    # returns min. dissipation after the initial maximum, with corresponding thickness
    lam_rho = qcm.calc_lamrho(n, g3, phi(g3))
    soln = optimize.minimize_scalar(lambda x: df_bulk(n, 1e8, 90).imag +
                    calc_df(n, x, g3, phi(g3)).imag,
                    bounds=[0.25*lam_rho, 0.75*lam_rho], method='bounded')
    return soln['fun'], soln['x']


# g3_max is the largest value of g3 for which gam5 is less than 20,000 Hz
# this portion of the code generates horizontal portion of the red line
# separating the viscoelastic and overdamped regions
g3_max = optimize.least_squares(lambda x: max_gamma(5, x)[0]-20000, 1e10,
                                  bounds = [2e9, 1e12])['x']

g31 = np.logspace(np.log10(g3_max), np.log10(0.9e12), npts1)
d_max1 = np.zeros(npts1)
for i in np.arange(npts1):
    lam_rho = qcm.calc_lamrho(5, g31[i], phi(g31[i]))
    max_gam_drho = max_gamma(5, g31[i])[1]
    d_max1[i] = optimize.least_squares(lambda x: calc_df(5, x, g31[i], 
                    phi(g31[i])).imag-20000, lam_rho/8, 
                    bounds=(0, max_gam_drho))['x']
    
   
# now we calculate line for minimum thickness
# this gets a bit tricky because we need to make sure we are getting the
# lowest value of d that meets are criterion
# this portion of the code generates the blue line separating the Sauerbrey 
# and viscoelastic regimes
npts3 = 200
g32 = np.logspace(np.log10(1e8), np.log10(0.8e12), npts3)
d_min = np.zeros(npts3)


for i in np.arange(npts3):
    lam_rho = qcm.calc_lamrho(3, g32[i], phi(g32[i]))
    # peak in this ratio is close to lamda/3 for fifth harmonic
    upper_lim = qcm.calc_lamrho(5, g32[i], phi(g32[i]))/3
    # now figure out where the first minimum is in df_diff
    first_min = optimize.minimize_scalar(lambda x: calc_df_diff(3, 5, x, g32[i], 
        phi(g32[i])), bounds=[0, upper_lim], method='bounded')
    if first_min['fun'] < -10:  # the value we are looking for is before this minimum
        d_min[i] = (optimize.least_squares(lambda x: abs(calc_df_diff(3, 5, x, g32[i],
             phi(g32[i])))-10, first_min['x']/2, bounds=[0, first_min['x']])['x'])
    else: # first minimum not shallow enough, we need the postive value that 
          # occurs for thicker films
        guess=np.average([first_min['x'], upper_lim])
        d_min[i] = optimize.least_squares(lambda x: abs(calc_df_diff(3, 5, x, g32[i],
             phi(g32[i])))-10, guess, bounds=[first_min['x'], upper_lim])['x']

# now we generate the black line separating the viscoelastic and bulk regions
g33 = np.logspace(8, np.log10(5.5e9), 20) 
delta_rho = np.zeros(20)
for i in np.arange(20):
    delta_rho[i] = qcm.calc_deltarho(3, g33[i], 
             phi(g33[i]))
    

# make the figure and plot all the lines
fig, ax = plt.subplots(1, 1, figsize=(3,3))
ax.set_xlabel(r'$|G_3^*|\rho$ (Pa$\cdot$g/cm$^3$)')
ax2=ax.twiny()
ax.set_ylabel(r'$d\rho$ ($\mu$m$\cdot$g/cm$^3$)')
ax.loglog(g31/1000, 1000*d_max1, 'r-')
ax.loglog(g32/1000, 1000*d_min, 'b-')
ax.loglog(g33/1000, 2000*delta_rho, 'k-')
ax.set_ylim(0.01, 10)
ax.set_xlim(1e5, 1e9)
ax2.set_xlim(phi(1e8), phi(1e12))
ax2.set_xlabel(r'$\phi_3$ (deg.)')
plt.xticks(ticks=[0, 15, 30, 45, 60, 75, 90])
ax2.plot([50.78,50.78], [0.92, 10], '-r')

# add text delineating different regions
plt.text(55, 0.03, 'Sauerbrey')
plt.text(80, 3, 'bulk')
plt.text(49,5, 'overdamped')
plt.text(62, 0.5,'viscoelastic') 
fig.tight_layout()
fig.savefig('Fig1.pdf')

