import sys
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
from scipy.optimize import curve_fit

## Usage
# You can call this script from a command line like so:
# `python OD_calibration.py filename.txt`
# Or make use of it in another python script by pasting this file into the same folder and importing:
# `import OD_calibration as odc
#  odc.OD_calibration(filename)`

def model(OD, p0, p1, p2):
    return p0 + p1*OD**p2

def OD_calibration(file):
    # Read file into pandas dataframe
    df = pd.read_csv(file, engine='python', skiprows=5, skipfooter=5, header=[0,1])

    fig, ax = plt.subplots(2,1,height_ratios=[3,1],sharex=True)
    plt.subplots_adjust(hspace=0)

    flasks = ['A','B','C','D']
    for idx, f in enumerate(flasks):
        ODs = df[f'{f}']['OD'].to_numpy()
        readings = df[f'{f}']['Reading'].to_numpy()

        # Remove missing readings
        ODs = ODs[~np.isnan(ODs)]
        readings = readings[~np.isnan(readings)]

        # Scipy's `curve_fit` can't handle ODs = 0.00, so replace with 0.00001:
        ODs[ODs == 0.00] = 0.00001

        # Must still have at least three readings - if not, plot dummy data for the legend and skip this flask
        if (np.size(ODs) < 3 or np.size(readings) < 3):
            ax[0].scatter(ODs, readings, label=f'{f}: not enough data')
            ax[0].plot([], [], 'o-')
            ax[1].plot([], [], 'o-')
            continue

        # Fit model
        popt, pcov = curve_fit(model, ODs, readings)

        # Get R2 values
        ss_res = np.dot((readings - model(ODs, *popt)), (readings - model(ODs, *popt)))
        mean_readings = np.mean(readings)
        ss_tot = np.dot((readings - mean_readings), (readings - mean_readings))
        r_2 = 1 - (ss_res / ss_tot)

        # Plot
        ax[0].scatter(
            ODs,
            readings,
            label=f'{f}: p0 = {popt[0]:.2f}, p1 = {popt[1]:.2f}, p2 = {popt[2]:.4f}, R² = {r_2:.4f}',
        )
        xdata = np.linspace(0, max(ODs))
        ax[0].plot(xdata, model(xdata, *popt))
        ax[1].plot(ODs, readings - model(ODs, *popt), 'o-')

        fig.suptitle(f"OD Calibration - {file}")
        ax[0].set_ylabel('reading')
        ax[1].set_ylabel('residuals')
        ax[1].set_xlabel('OD')
        ax[0].legend()

if __name__ == '__main__':
    if (len(sys.argv) < 2):
        print("Adv. Cal. analysis: no file given")
    else:
        OD_calibration(sys.argv[1])

        plt.show()

