#import these functions/packages
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy import stats
from scipy import linalg
import datetime as datetime
import matplotlib.dates as mdates
from k0_fcns import k0ext_from_Vext_pHcal
from k0_fcns import pHext_from_Vext_k0ext
from k0_fcns import phcalc_jp
from k0_fcns import k0ext_from_Vext_and_k2_pHcal
from pathlib import Path
import os
## ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ ##

os.system("osascript -e 'mount volume \"smb://atlas.shore.mbari.org/ProjectLibrary\"'")  ## make sure computer is connected to poject library where data is stored  -- MAC specifc

#file names that contain the raw data
files = ['DSD372_5mMborate_ChD_07092024', 'DSD371_5mMborate_ChC_07092024', 'DSD371_3mMborate_ChC_08022024', 'DSD372_3mMborate_ChD_08022024', 'DSD371_1mMBorate_ChC_08222024', 'DSD372_1mMBorate_ChD_08222024']

# create data frame to hold linear regression value, temperature data, sensor #, and solution concentrations - will be populated later in the script
cols = ['Sensor', 'Borate Concentration (mM)', 'Temperature (C)', 'Slope', 'Intercept', 'r-value', 'p-value']
linregression_data = pd.DataFrame(columns= cols)

## ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ ##

## loop through the data files, process and plot the data
for filename in files: 

    # check which sensor the data is from and gather the correct k2 and f(p) coeffs for calculating pH 
    if 'DSD371' in filename:
        k2 = -0.001053          ## k2 from PTC from PTC run using 0.01N HCl
        pcoefs = np.array([1.50439E-25, -7.40391E-22, 1.52834E-18, -0.0000000000000017309, 0.00000000000118432, -0.000000000509298, 0.000000138888, -0.0000232671])  # pressure coeffecients from PTC run using 0.01N HCl aka f(p), use this one
        #pcoefs = np.array([-1.02955000E-22,	4.18146000E-19,	-7.06696000E-16,	0.000000000000647041,	-0.000000000349579,	0.000000113683,	-0.0000214628])  # pressure coeffecients from PTC run using 0.01N HCl aka f(p), old - onlyt 7th order poly? relooked at data and recomended 8th order? (nov 2024)
        sensor_name = 'DSD371'
        print('Data is from DSD', str(filename[3:6]))

    elif 'DSD372' in filename:
        k2 = -0.00109363        ## k2 from PTC from PTC run using 0.01N HCl
        pcoefs = np.array([1.03051E-17,	-0.0000000000000331615,	0.0000000000394921,	-0.0000000195721 , 0.00000173246 ])  # pressure coeffecients from PTC run using 0.01N HCl aka f(p), new 
        #pcoefs = np.array([9.89491E-18, -0.0000000000000321385, 0.0000000000391067, -0.0000000203264, 0.00000219769])  # pressure coeffecients from PTC run using 0.01N HCl aka f(p), old, uses 3rd order instead of 0 order (recomended) for some processing
        sensor_name = 'DSD372'
        print('Data is from DSD', str(filename[3:6]))

    else:
        print('Unknown sensor used')


    ## Check which solution was used for borate solutions for importing the correct spec pH value

    if '5mM' in filename:
        spec_pH = 8.6666
        density = 1.02502

        concentration = 5

        save_path = '/Volumes/ProjectLibrary/901805_Coastal_Biogeochemical_Sensing/PTC/Processed Data/5mM Borate/'
        print ('The spec pH is from the 5 mM borate solution')

    elif  '3mM' in filename:
        spec_pH = 8.6645
        density = 1.02476

        concentration = 3 

        save_path = '/Volumes/ProjectLibrary/901805_Coastal_Biogeochemical_Sensing/PTC/Processed Data/3mM Borate/'
        print ('The spec pH is from the 3 mM borate solution')

    elif '1mM' in filename:
        spec_pH = 8.6469
        density = 1.02476

        concentration = 1

        save_path = '/Volumes/ProjectLibrary/901805_Coastal_Biogeochemical_Sensing/PTC/Processed Data/1mM Borate/'
        print ('The spec pH is from the 1 mM borate solution')

    else:
        print('Spec pH not known for this solution')


    ## ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ ## 

    ## import the data, drop the header information, and make sure data is in correct readable format
    path = '/Volumes/ProjectLibrary/901805_Coastal_Biogeochemical_Sensing/PTC/Raw_Data_Copy/'  # where the data is stored on project library (also should be on Atlas/Carbon)

    ## import the data file
    data= pd.read_csv(path + filename +'.csv', header = 0, delimiter= ',', skipinitialspace= True)
    data  = data [~data ['Labview Date/Time'].str.startswith('H')] # drop rows with header information (first column has 'H')
    data['Labview Date/Time'] = pd.to_datetime(data['Labview Date/Time'])  # convert date/time information into python readable format 
    data['Thermistor Temp'] = pd.to_numeric(data['Thermistor Temp'])  
    data['Bath Set Point'] = pd.to_numeric(data['Bath Set Point'])
    data['StepNum'] = data['StepNum'].astype(int)

    ## determine the difference of temperature and presure between each samples
    data['Pressure_diff'] = data['Pressure (dBar)'].diff()
    data['Temp_offset (bath v thermistor)'] = abs (data['Thermistor Temp']- data['Bath Set Point'])

    # drop data from breakin cycles 
    data  = data[(data['StepNum'] >14)]

    ## ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ ##

    ## Determine k0 of the solution/sensor to calculate pH
    ambient_data = data[(data['Pressure (dBar)'] < 21) & (data['Thermistor Temp'] > 19.9) & (data['Thermistor Temp'] < 20.1)]  ## filter for only VRS measurments at 20 C and (close to) ambient pressure
    ambient_data['K0_T_corrected'], ambient_data['K0_0C'] = k0ext_from_Vext_and_k2_pHcal(ambient_data['VRS VOLTS'], spec_pH, k2 , ambient_data['Thermistor Temp'],35) ## calculate k0 using above data, pH measured via spectometer before the experiment, and salinity of solution (35)

    mean_k0 = np.mean(ambient_data['K0_0C'])    ## calculate the mean k0 

    ## calulate the free and total pH using VRS, pressure, temperature, sailinty, k0 estimated using above subdata set, etc
    data['pH_free'], data['pH_tot'] = phcalc_jp(data['VRS VOLTS'], data['Pressure (dBar)'], data['Thermistor Temp'], 35, mean_k0, k2, pcoefs)


    # Initialize the pH_diff column
    data['pH_diff'] = 0.0

    ## ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ ##

    ## calulate the pH difference for each step in inc/dec data ####

    # Function to check if pressure is decreasing, increasing, or staying staionary 
    def pressure_trend(step_data):
        first_pressure = step_data.iloc[0]
        last_pressure = step_data.iloc[-1]
        pressure_change = last_pressure - first_pressure
        
        if pressure_change < -200 :
            return 'dec'
        
        elif pressure_change > 200:
            return 'inc'
        else: 
            return 'staionary'


    # Create a new column 'Inc_Dec_Pressure' by applying the function for each StepNum
    data['Inc_Dec_Pressure'] = data.groupby('StepNum')['Pressure (dBar)'].transform(lambda x: pressure_trend(x))

    # Make subdata sets only containing increasing or decreasing pressure data
    dec_data = data[(data['Inc_Dec_Pressure'] == 'dec')]
    inc_data = data[(data['Inc_Dec_Pressure'] == 'inc')]

    # make a list of unique step numbers associated with the increasing or decreasing data sets
    dec_stepnums = dec_data['StepNum'].unique()
    inc_stepnums = inc_data['StepNum'].unique()


    ## initialize the data column with delta pH for each presure cycle
    data['pH_diff'] = 0

    # get the delta pH referencing the first measurement in that step (low pressure) 
    for step in inc_stepnums: 
        first_instance = data[data['StepNum'] == step].index[0]
        last_instance = data[data['StepNum'] ==step].index[-1]

        row = first_instance
        
        while row <= last_instance:
            pH_diff = (data.loc[row, 'pH_tot'] - data.loc[first_instance, 'pH_tot'])
            data.loc[row, 'pH_diff'] = pH_diff
            row +=1

    # get the delta pH referenceing the last measurement in that step (low pressure)
    for step in dec_stepnums:
        first_instance = data[data['StepNum'] == step].index[0]
        last_instance = data[data['StepNum'] ==step].index[-1]

        data.loc[last_instance, 'pH_diff'] = 0
        
        row = first_instance         

        while row < last_instance:
            pH_diff = (data.loc[row, 'pH_tot'] - data.loc[last_instance, 'pH_tot'])
            data.loc[row, 'pH_diff'] = pH_diff
            row +=1
    ## ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ ##

    # calculate -RTln(10) (pH0 - pHp)
    data['rtln10_delpH'] = -data['pH_diff'] * (data['Thermistor Temp'] +273.15) * np.log(10) * 83.15

    data = data[(data['Inc_Dec_Pressure'] != 'staionary')]  # drop data from when the sensor is at a constant presusre (high and low pressure holds)

    ## ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ ##

    ## plot the data

    ## Pressure vs VRS
    plt.figure()
    plt.rcParams['figure.dpi'] = 300
    plt.scatter(data['Pressure (dBar)'], data['VRS VOLTS'], c = data['Thermistor Temp'], cmap = 'jet')
    cbar = plt.colorbar()
    cbar.set_label('Temperature (˚C)', fontsize = 14)
    plt.ylabel ('VRS (V)', fontsize = 16)
    plt.xlabel ('Pressure (dBar)', fontsize = 16)
    plt.savefig(save_path + 'Overall Pressure vs VRS_'+ sensor_name +'.jpeg')


    ## Pressure vs pH
    plt.figure()
    plt.rcParams['figure.dpi'] = 300
    plt.scatter(data['Pressure (dBar)'], data['pH_tot'], c = data['Thermistor Temp'], cmap = 'jet')
    cbar = plt.colorbar()
    cbar.set_label('Temperature (˚C)', fontsize = 14)
    plt.ylabel ('pH', fontsize = 16)
    plt.xlabel ('Pressure (dBar)', fontsize = 16)
    plt.savefig(save_path + 'Overall Pressure vs pH_'+ sensor_name +'.jpeg')


    ## pressure vs delta pH
    plt.figure()
    plt.rcParams['figure.dpi'] = 300
    plt.scatter(data['Pressure (dBar)'], data['pH_diff'], c = data['Thermistor Temp'], cmap = 'jet')
    cbar = plt.colorbar()
    cbar.set_label('Temperature (˚C)', fontsize = 14)
    plt.ylabel ('Delta pH', fontsize = 16)
    plt.xlabel ('Pressure (dBar)', fontsize = 16)
    plt.savefig(save_path + 'Overall Pressure vs delta pH_'+ sensor_name +'.jpeg')



    ## pressure (bar) vs -RTln(10)(pH0 - pHp)
    plt.figure()
    plt.rcParams['figure.dpi'] = 300
    plt.scatter(data['Pressure (dBar)']/10, data['rtln10_delpH'], c = data['Thermistor Temp'], cmap = 'jet')
    cbar = plt.colorbar()
    cbar.set_label('Temperature (˚C)', fontsize = 14)
    plt.ylabel ('-RTln(10)delta pH', fontsize = 16)
    plt.xlabel ('Pressure (Bar)', fontsize = 16)
    plt.savefig(save_path + 'Overall Pressure vs -RTlnpH_'+ sensor_name +'.jpeg')


    ## ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ ##

    ## plot cycles from individual temperatures and calculate slope
    temps = data['Bath Set Point'].unique()

    for temp in temps:
        subdata = data[(data['Bath Set Point'] == temp)]

        slope, intercept, r, p, std= stats.linregress((subdata['Pressure (dBar)']/10), subdata['rtln10_delpH'])

        txt ='Temperature = ' +str(temp) + '\nSlope = ' + str("{:.4f}".format(slope)) + '\nIntercept = ' + str("{:.2f}".format(intercept)) + '\nr = ' + str("{:.4f}".format(r))

        ## append data to data frame with linear regression data, this will be used to make delV vs [borate] and temp plots
        new_row = {"Sensor": sensor_name, "Borate Concentration (mM)": concentration, "Temperature (C)": temp, "Slope": slope, "Intercept": intercept, "r-value": r, "p-value": p}
        linregression_data = pd.concat([linregression_data, pd.DataFrame([new_row])], ignore_index = True)


        ## pressure vs delta pH
        fig1 = plt.figure()
        plt.rcParams['figure.dpi'] = 300
        plt.plot(subdata['Pressure (dBar)'],subdata['pH_diff'], color ='black')
        plt.ylabel ('Delta pH', fontsize = 16)
        plt.xlabel ('Pressure (dBar)', fontsize = 16)
        plt.savefig(save_path + str(temp) + '_Pressure vs delta pH_'+ sensor_name +'.jpeg')


        ## pressure (bar) vs RTln(10)(pH0 - pHp)
        fig2 = plt.figure()
        plt.rcParams['figure.dpi'] = 300
        plt.plot(subdata['Pressure (dBar)']/10, subdata['rtln10_delpH'], color ='black')
        plt.ylabel ('-RTln(10)delta pH', fontsize = 16)
        plt.xlabel ('Pressure (Bar)', fontsize = 16)
        fig2.text(0.65, 0.15, txt)
        plt.savefig(save_path + str(temp)+ '_Pressure vs -RTln10delpH_'+ sensor_name +'.jpeg')


## ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ ##

## plot delv vs temperature and borate concentrations 

concs = [1,3,5]

dfs = ['DSD371', 'DSD372']

for df in dfs: 
    subset = linregression_data[(linregression_data['Sensor'] == df)]

    # plot overall delV vs borate concentraion as a scatter plot with temperature as a third variable
    plt.figure()
    plt.scatter(subset['Slope'], subset['Borate Concentration (mM)'], c = subset['Temperature (C)'], cmap = 'jet')
    cbar = plt.colorbar()
    cbar.set_label('Temperature (˚C)', fontsize = 14)
    plt.ylabel('[Borate] (mM)')
    plt.xlabel(r'delV $_\mathrm{borate}$ [cm$^{3}$ mol$^{-1}$]')
    plt.savefig(save_path + 'all_temps_' + str(df)+ '_delV_vs_borate_conc' + '.jpeg')

    for conc in concs:
        subdata = subset[(subset['Borate Concentration (mM)'] == conc )]
        s, i, r, p, std = stats.linregress(subdata['Temperature (C)'], subdata['Slope'])

        info = str(df) + '\nBorate Concentraion: ' + str(conc) + 'mM' + '\ndelV =  ' + str("{:.4f}".format(s)) + 'T(˚C) ' + str("{:.4f}".format(i)) + ' ;  r = ' + str("{:.4f}".format(r))

        # plot temperature vs delV data for each borate concentraion run (for each DSD)
        fig3 = plt.figure()
        plt.rcParams['figure.dpi'] = 300
        plt.plot(subdata['Temperature (C)'], subdata['Temperature (C)']* s + i, '-', color = 'black')
        plt.plot(subdata['Temperature (C)'], subdata['Slope'],'o')
        fig3.text(0.15, 0.15, info)
        plt.ylabel(r'delV $_\mathrm{borate}$ [cm$^{3}$ mol$^{-1}$]')
        plt.xlabel('Temperature (˚C)')
        plt.savefig(save_path + str(conc) + '_' + str(df) + '_Temp_vs_delV' + '.jpeg')

    for temp in temps:
        subtempdata = subset[(subset['Temperature (C)'] == temp)]
        #s, i, r, p, std = stats.linregress(subtempdata['Slope'], subtempdata['Borate Concentration (mM)'])

        info = str(df) + '\nTemperature: ' + str(temp)

        # plot delV vs borate concentraion for each temperature cycle (for each DSD)
        fig4 = plt.figure()
        plt.rcParams['figure.dpi'] = 300 
        plt.plot(subtempdata['Slope'], subtempdata['Borate Concentration (mM)'], 'o')
        plt.ylabel('[Borate] (mM)')
        #fig4.text(0.55, 0.15, info)
        plt.xlabel(r'delV $_\mathrm{borate}$ [cm$^{3}$ mol$^{-1}$]')
        plt.title(info)
        plt.savefig(save_path + str(temp) + '_' + str(df)+ '_delV_vs_borate_conc' + '.jpeg')