### Zeta pH Curve Analysis - Adaline Farrow 2024

# import stuff
import matplotlib.pyplot as plt
import numpy as np
import os 
import pandas as pd
from scipy.optimize import *

def readzeta(folder):
    os.chdir(folder)
    zeta = []
    pH = []
    
    for i in os.listdir():
        data = pd.read_excel(i)
        
        zeta.append(np.array(data)[7,6])
        
        pH.append(float(i.split("pH",1)[1][:5].replace('_', '')))
        
    return zeta, pH

def readcond(folder):
    os.chdir(folder)
    cond = []
    pH = []
    
    for i in os.listdir():
        data = pd.read_excel(i)
        
        cond.append(np.array(data)[7,9])
        
        pH.append(float(i.split("pH",1)[1][:5].replace('_', '')))
        
    return cond, pH

def resid(x, y, popt):
    residuals = y - func(x, *popt)
    ss_res = np.sum(residuals**2)
    ss_tot = np.sum((y-np.mean(y))**2)
    r_squared = 1 - (ss_res / ss_tot)
    
    return residuals, r_squared

def zetaplot(pH, zeta):
    zetaavg = []

    for i in zeta:
        zetaavg.append(np.average(i))

    zetaavg = np.array(zetaavg)
    
    xln = np.linspace(min(pH)-0.5, max(pH) + 0.5, 100)
    
    plt.plot(pH, zetaavg, 'ro')
    plt.xlabel('pH', size = 13)
    plt.ylabel('Zeta Potential (mV)', size = 13)
    plt.legend(loc='best', fontsize = 12)
    
    return zetaavg

def condplot(cond, pH):    
    xln = np.linspace(min(pH)-0.5, max(pH) + 0.5, 100)
    
    plt.plot(pH, cond, 'ro')
    plt.xlabel('pH', size = 13)
    plt.ylabel('Conductance (uS)', size = 13)
    plt.legend(loc='best', fontsize = 12)

def func(x, a, b, c):
    return a*(x+c)**2 + b

def plotquad(x, y):
    xln = np.linspace(min(x)-0.1, max(x)+0.1, 1000)
    
    popt, pcov = curve_fit(func, x, y)
    
    residuals, r_square = resid(x, y, popt)
    
    plt.plot(xln, func(xln, *popt), 'k--', label=f'''{np.round(popt[0], 2)}*(x {np.round(popt[2],2)})$^2$ {np.round(popt[1],2)}, 
    R$^2$ = {np.round(r_square,4)}''')
    
    plt.legend(loc='best', fontsize = 12)
    
def save(folder, filename):
    os.chdir(folder)
    plt.savefig(filename+'.png')

"""
Func: readzeta
Input: (folder directory:str) - directory in which zeta files are stored - ONLY READS pH IF 'pHXX.XX' IN FILENAME
Output: None
Returns: zeta values:list, pH values:list
"""
zeta, pH = readzeta('C:/Users/ajfar/Desktop/McGill/Research/Harrington_Singh_2024/pH Zeta Data/solcol/')

###########################################################################################################################
"""
Func: zetaplot
Input: (pH values:list, zeta values:list) - OBTAINED FROM READZETA
Output: Plot of zeta values vs. pH values
Returns: None
"""
zetaplot(pH, zeta)

###########################################################################################################################
"""
Func: plotquad
Input: (pH values:list, zeta values:list) - OBTAINED FROM READZETA
Output: Plot of quadratic fit on inputted data
Returns: None
"""
plotquad(pH, zeta)

###########################################################################################################################
"""
Func: save
Input: (Output Folder Directory:str, Desired Filename:str) - WILL ADD .png AUTOMATICALLY
Output: High-res png to prespecified folder
Returns: None
"""
save('C:/Users/ajfar/Desktop/McGill/Research/Harrington_Singh_2024/pH Zeta Data/plots/', 'lowpH_plot')

"""
Func: readcond
Input: (folder) - directory in which zeta files are stored - ONLY READS pH IF 'pHXX.XX' IN FILENAME
Output: None
Returns: conductance values, pH values - as lists
"""
cond, pH = readcond('C:/Users/ajfar/Desktop/McGill/Research/Harrington_Singh_2024/pH Zeta Data/ghcl/')

"""
Func: condplot
Input: (conductance values, zeta values) - OBTAINED FROM READCOND
Output: Plot of conductance values vs. pH values
Returns: None
"""
condplot(cond, pH)