## Required imports
# basic imports
import os
import sys
import random
import re
import time
import datetime
from dateutil.relativedelta import relativedelta

# ds imports
import pandas as pd
import numpy as np
from math import sqrt
import scipy
import scipy.stats as stats
from scipy.spatial import cKDTree
from scipy.optimize import curve_fit
from statsmodels.nonparametric.smoothers_lowess import lowess

# visualization imports
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from matplotlib import colors
from matplotlib.dates import date2num, num2date
from matplotlib.cm import ScalarMappable
from matplotlib.colors import Normalize
from matplotlib.lines import Line2D
from matplotlib.patches import Polygon
from matplotlib.collections import PatchCollection
import seaborn as sns

# ml imports
from supersmoother import SuperSmoother
from sklearn import preprocessing
from sklearn.preprocessing import StandardScaler, MinMaxScaler
from sklearn.decomposition import PCA
from sklearn.cluster import KMeans
from sklearn.gaussian_process import GaussianProcessRegressor
from sklearn.ensemble import RandomForestRegressor
from sklearn.pipeline import make_pipeline
from sklearn.gaussian_process.kernels import Matern, WhiteKernel
from sklearn.preprocessing import PolynomialFeatures
from sklearn.linear_model import LinearRegression, Ridge, Lasso, RidgeCV, LassoCV
from sklearn.model_selection import GridSearchCV, LeaveOneOut
from sklearn.metrics import mean_squared_error, r2_score
import tensorflow as tf
from tensorflow.keras.models import Model, Sequential
from tensorflow.keras.layers import Input, LSTM, RepeatVector, TimeDistributed, Dense

# geospatial imports
import geopandas as gpd
import contextily as cx
import folium
import shapely
from shapely.geometry import LineString, MultiLineString, Point
from shapely.ops import nearest_points
from shapely.strtree import STRtree
from pyproj import Proj, Transformer
from ipyleaflet import (Map, basemaps, WidgetControl, GeoJSON, 
                        LayersControl, Icon, Marker,FullScreenControl,
                        CircleMarker, Popup, AwesomeIcon)

# other setup
import warnings
warnings.filterwarnings("ignore")

plt.rcParams["font.family"] = "Times New Roman"


##################################################################################################
##################################################################################################


## class PylenmDataFactory
class PylenmDataFactory(object):
    """Class object that initilaizes Pylenm given data.
    """
    
    def __init__(self, data: pd.DataFrame):
            """Initlizes pylenm with a Pandas DataFrame

            Args:
                data (pd.DataFrame): Data to be imported.
            """
            self.setData(data)
            self.__jointData = [None, 0]

    # DATA VALIDATION     
    def __isValid_Data(self, data):
        """Validates data for Pylenm usage.

        Returns:
            Tuple[bool, str]: Returns a tuple of length 2. 
                First value is `True` if data is valid, `False` otherwise.
                Second value is a string describing the validation result.
        """

        if(str(type(data)).lower().find('dataframe') == -1):
            return (False, 'Make sure the data is a pandas DataFrame.\n')
        if(not self.__hasColumns_Data(data)):
            return (False, 'Make sure that ALL of the columns specified in the REQUIREMENTS are present.\n')
        else:
            return (True, None)

    def __isValid_Construction_Data(self, data):
        if(str(type(data)).lower().find('dataframe') == -1):
            return (False, 'Make sure the data is a pandas DataFrame.\n')
        if(not self.__hasColumns_Construction_Data(data)):
            return (False, 'Make sure that ALL of the columns specified in the REQUIREMENTS are present.\n')
        else:
            return (True, None)

    # COLUMN VALIDATION 
    def __hasColumns_Data(self, data):
        find = ['COLLECTION_DATE','STATION_ID','ANALYTE_NAME','RESULT','RESULT_UNITS']
        cols = list(data.columns)
        cols = [x.upper() for x in cols]
        hasCols =  all(item in cols for item in find)
        return hasCols

    def __hasColumns_Construction_Data(self, data):
        find = ['STATION_ID', 'AQUIFER', 'WELL_USE', 'LATITUDE', 'LONGITUDE', 'GROUND_ELEVATION', 'TOTAL_DEPTH']
        cols = list(data.columns)
        cols = [x.upper() for x in cols]
        hasCols =  all(item in cols for item in find)
        return hasCols

    # SETTING DATA
    def setData(self, data: pd.DataFrame, verbose: bool = True) -> None:
        """Saves the dataset into pylenm

        Args:
            data (pd.DataFrame): Dataset to be imported.
            verbose (bool, optional): Prints success message. Defaults to True.

        Returns:
            None
        """
        validation = self.__isValid_Data(data)
        if(validation[0]):
            # Make all columns all caps
            cols_upper = [x.upper() for x in list(data.columns)]
            data.columns = cols_upper
            self.data = data
            if(verbose):
                print('Successfully imported the data!\n')
            self.__set_units()
        else:
            print('ERROR: {}'.format(validation[1]))
            return self.__REQUIREMENTS_DATA()

    def setConstructionData(self, construction_data: pd.DataFrame, verbose=True):
        """Imports the addtitional well information as a separate DataFrame.

        Args:
            construction_data (pd.DataFrame): Data with additonal details.
            verbose (bool, optional): Prints success message. Defaults to True.

        Returns:
            None
        """
        validation = self.__isValid_Construction_Data(construction_data)
        if(validation[0]):
            # Make all columns all caps
            cols_upper = [x.upper() for x in list(construction_data.columns)]
            construction_data.columns = cols_upper
            self.construction_data = construction_data.set_index(['STATION_ID'])
            if(verbose):
                print('Successfully imported the construction data!\n')
        else:
            print('ERROR: {}'.format(validation[1]))
            return self.__REQUIREMENTS_CONSTRUCTION_DATA()

    def jointData_is_set(self, lag):
        """Checks to see if getJointData function was already called and saved for given lag.

        Args:
            lag (int): number of days to look ahead and behind the specified date (+/-)

        Returns:
            bool: True if JointData was already calculated, False, otherwise.
        """
        if(str(type(self.__jointData[0])).lower().find('dataframe') == -1):
            return False
        else:
            if(self.__jointData[1]==lag):
                return True
            else:
                return False

    def __set_jointData(self, data, lag):
        self.__jointData[0] = data
        self.__jointData[1] = lag

    # GETTING DATA      
    def getData(self):
        """Returns the concentration data in pylenm 

        Returns:
            pd.DataFrame: concentration data that was passed into pylenm
        """
        return self.data

    def get_Construction_Data(self):
        """Returns the construction data in pylenm 

        Returns:
            pd.DataFrame: construction data that was passed into pylenm
        """
        return self.construction_data

    # MESSAGES FOR INVALID DATA          
    def __REQUIREMENTS_DATA(self):
        print('PYLENM DATA REQUIREMENTS:\nThe imported data needs to meet ALL of the following conditions to have a successful import:')
        print('   1) Data should be a pandas dataframe.')
        print("   2) Data must have these column names: \n      ['COLLECTION_DATE','STATION_ID','ANALYTE_NAME','RESULT','RESULT_UNITS']")

    def __REQUIREMENTS_CONSTRUCTION_DATA(self):
        print('PYLENM CONSTRUCTION REQUIREMENTS:\nThe imported construction data needs to meet ALL of the following conditions to have a successful import:')
        print('   1) Data should be a pandas dataframe.')
        print("   2) Data must have these column names: \n      ['station_id', 'aquifer', 'well_use', 'latitude', 'longitude', 'ground_elevation', 'total_depth']")

    # Helper function for plot_correlation
    # Sorts analytes in a specific order: 'TRITIUM', 'URANIUM-238','IODINE-129','SPECIFIC CONDUCTANCE', 'PH', 'DEPTH_TO_WATER'
    def __custom_analyte_sort(self, analytes):
        my_order = 'TURISPDABCEFGHJKLMNOQVWXYZ-_abcdefghijklmnopqrstuvwxyz135790 2468'
        return sorted(analytes, key=lambda word: [my_order.index(c) for c in word])

    def __plotUpperHalf(self, *args, **kwargs):
        corr_r = args[0].corr(args[1], 'pearson')
        corr_text = f"{corr_r:2.2f}"
        ax = plt.gca()
        ax.set_axis_off()
        marker_size = abs(corr_r) * 10000
        ax.scatter([.5], [.5], marker_size, [corr_r], alpha=0.6, cmap="coolwarm",
                    vmin=-1, vmax=1, transform=ax.transAxes)
        font_size = abs(corr_r) * 40 + 5
        ax.annotate(corr_text, [.5, .48,],  xycoords="axes fraction", # [.5, .48,]
                    ha='center', va='center', fontsize=font_size, fontweight='bold')


    def simplify_data(self, data=None, inplace=False, columns=None, save_csv=False, file_name= 'data_simplified', save_dir='data/'):
        """Removes all columns except 'COLLECTION_DATE', 'STATION_ID', 'ANALYTE_NAME', 'RESULT', and 'RESULT_UNITS'.
            
            If the user specifies additional columns in addition to the ones listed above, those columns will be kept.
            The function returns a dataframe and has an optional parameter to be able to save the dataframe to a csv file.

        Args:
            data (pd.DataFrame, optional): data to simplify. Defaults to None.
            inplace (bool, optional): save data to current working dataset. Defaults to False.
            columns (list, optional): list of any additional columns on top of  ['COLLECTION_DATE', 'STATION_ID', 'ANALYTE_NAME', 'RESULT', and 'RESULT_UNITS'] to be kept in the dataframe. Defaults to None.
            save_csv (bool, optional): flag to determine whether or not to save the dataframe to a csv file. Defaults to False.
            file_name (str, optional): name of the csv file you want to save. Defaults to 'data_simplified'.
            save_dir (str, optional): name of the directory you want to save the csv file to. Defaults to 'data/'.

        Returns:
            pd.DataFrame
        """
        if(str(type(data)).lower().find('dataframe') == -1):
            data = self.data
        else:
            data = data
        if(columns==None):
            sel_cols = ['COLLECTION_DATE','STATION_ID','ANALYTE_NAME','RESULT','RESULT_UNITS']
        else:
            hasColumns =  all(item in list(data.columns) for item in columns)
            if(hasColumns):            
                sel_cols = ['COLLECTION_DATE','STATION_ID','ANALYTE_NAME','RESULT','RESULT_UNITS'] + columns
            else:
                print('ERROR: specified column(s) do not exist in the data')
                return None

        data = data[sel_cols]
        data.COLLECTION_DATE = pd.to_datetime(data.COLLECTION_DATE)
        data = data.sort_values(by="COLLECTION_DATE")
        dup = data[data.duplicated(['COLLECTION_DATE', 'STATION_ID','ANALYTE_NAME', 'RESULT'])]
        data = data.drop(dup.index)
        data = data.reset_index().drop('index', axis=1)
        if(save_csv):
            if not os.path.exists(save_dir):
                os.makedirs(save_dir)
            data.to_csv(save_dir + file_name + '.csv')
            print('Successfully saved "' + file_name +'.csv" in ' + save_dir)
        if(inplace):
            self.setData(data, verbose=False)
        return data

    def get_MCL(self, analyte_name):
        """Returns the Maximum Concentration Limit value for the specified analyte. Example: 'TRITIUM' returns 1.3

        Args:
            analyte_name (str): name of the analyte to be processed

        Returns:
            float: MLC value
        """
        mcl_dictionary = {'TRITIUM': 1.3, 'URANIUM-238': -1.69,  'NITRATE-NITRITE AS NITROGEN': 1,
                            'TECHNETIUM-99': 2.95, 'IODINE-129': -3, 'STRONTIUM-90': -2.1
                            }
        return mcl_dictionary[analyte_name]

    def __set_units(self):
        analytes = list(np.unique(self.data[['ANALYTE_NAME']]))
        mask1 = ~self.data[['ANALYTE_NAME','RESULT_UNITS']].duplicated()
        res = self.data[['ANALYTE_NAME','RESULT_UNITS']][mask1]
        mask2 = ~self.data[['ANALYTE_NAME']].duplicated()
        res = res[mask2]
        unit_dictionary = pd.Series(res.RESULT_UNITS.values,index=res.ANALYTE_NAME).to_dict()
        self.unit_dictionary = unit_dictionary
        
    def get_unit(self, analyte_name):
        """Returns the unit of the analyte you specify. Example: 'DEPTH_TO_WATER' may return 'ft'

        Args:
            analyte_name (str): ame of the analyte to be processed

        Returns:
            str: unit of analyte
        """
        return self.unit_dictionary[analyte_name]

    def filter_by_column(self, data=None, col=None, equals=[]):
        """Filters construction data based on one column. You only specify ONE column to filter by, but can selected MANY values for the entry.

        Args:
            data (pd.DataFrame, optional): dataframe to filter. Defaults to None.
            col (str, optional): column to filter. Example: col='STATION_ID'. Defaults to None.
            equals (list, optional): values to filter col by. Examples: equals=['FAI001A', 'FAI001B']. Defaults to [].

        Returns:
            pd.DataFrame: returns filtered dataframe
        """
        if(data is None):
            return 'ERROR: DataFrame was not provided to this function.'
        else:
            if(str(type(data)).lower().find('dataframe') == -1):
                return 'ERROR: Data provided is not a pandas DataFrame.'
            else:
                data = data
        # DATA VALIDATION
        if(col==None):
            return 'ERROR: Specify a column name to filter by.'
        data_cols = list(data.columns)
        if((col in data_cols)==False): # Make sure column name exists 
            return 'Error: Column name "{}" does not exist'.format(col)
        if(equals==[]):
            return 'ERROR: Specify a value that "{}" should equal to'.format(col)
        data_val = list(data[col])
        for value in equals:
            if((value in data_val)==False):
                return 'ERROR: No value equal to "{}" in "{}".'.format(value, col)

        # QUERY
        final_data = pd.DataFrame()
        for value in equals:
            current_data = data[data[col]==value]
            final_data = pd.concat([final_data, current_data])
        return final_data

    def filter_wells(self, units):
        """Returns a list of the well names filtered by the unit(s) specified.

        Args:
            units (list): Letter of the well to be filtered (e.g. [‘A’] or [‘A’, ‘D’])

        Returns:
            list: well names filtered by the unit(s) specified
        """
        data = self.data
        if(units==None):
            units= ['A', 'B', 'C', 'D']
        def getUnits():
            wells = list(np.unique(data.STATION_ID))
            wells = pd.DataFrame(wells, columns=['STATION_ID'])
            for index, row in wells.iterrows():
                mo = re.match('.+([0-9])[^0-9]*$', row.STATION_ID)
                last_index = mo.start(1)
                wells.at[index, 'unit'] = row.STATION_ID[last_index+1:]
                u = wells.unit.iloc[index]
                if(len(u)==0): # if has no letter, use D
                    wells.at[index, 'unit'] = 'D'
                if(len(u)>1): # if has more than 1 letter, remove the extra letter
                    if(u.find('R')>0):
                        wells.at[index, 'unit'] = u[:-1]
                    else:
                        wells.at[index, 'unit'] = u[1:]
                u = wells.unit.iloc[index]
                if(u=='A' or u=='B' or u=='C' or u=='D'):
                    pass
                else:
                    wells.at[index, 'unit'] = 'D'
            return wells
        df = getUnits()
        res = df.loc[df.unit.isin(units)]
        return list(res.STATION_ID)

    def remove_outliers(self, data, z_threshold=4):
        """Removes outliers from a dataframe based on the z_scores and returns the new dataframe.

        Args:
            data (pd.DataFrame): data for the outliers to removed from
            z_threshold (int, optional): z_score threshold to eliminate. Defaults to 4.

        Returns:
            pd.DataFrame: data with outliers removed
        """
        z = np.abs(stats.zscore(data))
        row_loc = np.unique(np.where(z > z_threshold)[0])
        data = data.drop(data.index[row_loc])
        return data

    def get_analyte_details(self, analyte_name, filter=False, col=None, equals=[], save_to_file = False, save_dir='analyte_details'):
        """Returns a csv file saved to save_dir with details pertaining to the specified analyte. Details include the well names, the date ranges and the number of unique samples.

        Args:
            analyte_name (str): name of the analyte to be processed
            filter (bool, optional): whether to filter the data. Defaults to False.
            col (str, optional): column to filter. Example: col='STATION_ID'. Defaults to None.
            equals (list, optional): values to filter col by. Examples: equals=['FAI001A', 'FAI001B']. Defaults to [].
            save_to_file (bool, optional): whether to save data to file. Defaults to False.
            save_dir (str, optional): name of the directory you want to save the csv file to. Defaults to 'analyte_details'.

        Returns:
            pd.DataFrame: Table with well information
        """
        data = self.data
        data = data[data.ANALYTE_NAME == analyte_name].reset_index().drop('index', axis=1)
        data = data[~data.RESULT.isna()]
        data = data.drop(['ANALYTE_NAME', 'RESULT', 'RESULT_UNITS'], axis=1)
        data.COLLECTION_DATE = pd.to_datetime(data.COLLECTION_DATE)
        if(filter):
            filter_res = self.filter_by_column(data=self.construction_data, col=col, equals=equals)
            if('ERROR:' in str(filter_res)):
                return filter_res
            query_wells = list(data.STATION_ID.unique())
            filter_wells = list(filter_res.index.unique())
            intersect_wells = list(set(query_wells) & set(filter_wells))
            if(len(intersect_wells)<=0):
                return 'ERROR: No results for this query with the specifed filter parameters.'
            data = data[data['STATION_ID'].isin(intersect_wells)]        

        info = []
        wells = np.unique(data.STATION_ID.values)
        for well in wells:
            current = data[data.STATION_ID == well]
            startDate = current.COLLECTION_DATE.min().date()
            endDate = current.COLLECTION_DATE.max().date()
            numSamples = current.duplicated().value_counts()[0]
            info.append({'Well Name': well, 'Start Date': startDate, 'End Date': endDate,
                            'Date Range (days)': endDate-startDate ,
                            'Unique samples': numSamples})
            details = pd.DataFrame(info)
            details.index = details['Well Name']
            details = details.drop('Well Name', axis=1)
            details = details.sort_values(by=['Start Date', 'End Date'])
            details['Date Range (days)'] = (details['Date Range (days)']/ np.timedelta64(1, 'D')).astype(int)
        if(save_to_file):
            if not os.path.exists(save_dir):
                os.makedirs(save_dir)
            details.to_csv(save_dir + '/' + analyte_name + '_details.csv')
        return details

    def get_data_summary(self, analytes=None, sort_by='date', ascending=False, filter=False, col=None, equals=[]):
        """Returns a dataframe with a summary of the data for certain analytes. Summary includes the date ranges and the number of unique samples and other statistics for the analyte results.

        Args:
            analytes (list, optional): list of analyte names to be processed. If left empty, a list of all the analytes in the data will be used. Defaults to None.
            sort_by (str, optional): {‘date’, ‘samples’, ‘wells’} sorts the data by either the dates by entering: ‘date’, the samples by entering: ‘samples’, or by unique well locations by entering ‘wells’. Defaults to 'date'.
            ascending (bool, optional): flag to sort in ascending order.. Defaults to False.
            filter (bool, optional): flag to indicate filtering. Defaults to False.
            col (str, optional): column to filter. Example: col='STATION_ID'. Defaults to None.
            equals (list, optional): values to filter col by. Examples: equals=['FAI001A', 'FAI001B']. Defaults to [].

        Returns:
            pd.DataFrame: Table with well information
        """
        data = self.data
        if(analytes == None):
            analytes = data.ANALYTE_NAME.unique()
        data = data.loc[data.ANALYTE_NAME.isin(analytes)].drop(['RESULT_UNITS'], axis=1)
        data = data[~data.duplicated()] # remove duplicates
        data.COLLECTION_DATE = pd.to_datetime(data.COLLECTION_DATE)
        data = data[~data.RESULT.isna()]
        if(filter):
            filter_res = self.filter_by_column(data=self.construction_data, col=col, equals=equals)
            if('ERROR:' in str(filter_res)):
                return filter_res
            query_wells = list(data.STATION_ID.unique())
            filter_wells = list(filter_res.index.unique())
            intersect_wells = list(set(query_wells) & set(filter_wells))
            if(len(intersect_wells)<=0):
                return 'ERROR: No results for this query with the specifed filter parameters.'
            data = data[data['STATION_ID'].isin(intersect_wells)]

        info = []
        for analyte_name in analytes:
            query = data[data.ANALYTE_NAME == analyte_name]
            startDate = min(query.COLLECTION_DATE)
            endDate = max(query.COLLECTION_DATE)
            numSamples = query.shape[0]
            wellCount = len(query.STATION_ID.unique())
            stats = query.RESULT.describe().drop('count', axis=0)
            stats = pd.DataFrame(stats).T
            stats_col = [x for x in stats.columns]

            result = {'Analyte Name': analyte_name, 'Start Date': startDate, 'End Date': endDate,
                        'Date Range (days)':endDate-startDate, '# unique wells': wellCount,'# samples': numSamples,
                        'Unit': self.get_unit(analyte_name) }
            for num in range(len(stats_col)):
                result[stats_col[num]] = stats.iloc[0][num] 

            info.append(result)

            details = pd.DataFrame(info)
            details.index = details['Analyte Name']
            details = details.drop('Analyte Name', axis=1)
            if(sort_by.lower() == 'date'):
                details = details.sort_values(by=['Start Date', 'End Date', 'Date Range (days)'], ascending=ascending)
            elif(sort_by.lower() == 'samples'):
                details = details.sort_values(by=['# samples'], ascending=ascending)
            elif(sort_by.lower() == 'wells'):
                details = details.sort_values(by=['# unique wells'], ascending=ascending)

        return details

    def get_well_analytes(self, well_name=None, filter=False, col=None, equals=[]):
        """Displays the analyte names available at given well locations.

        Args:
            well_name (str, optional): name of the well. If left empty, all wells are returned.. Defaults to None.
            filter (bool, optional): flag to indicate filtering. Defaults to False.
            col (str, optional): column to filter. Example: col='STATION_ID'. Defaults to None.
            equals (list, optional): values to filter col by. Examples: equals=['FAI001A', 'FAI001B']. Defaults to [].

        Returns:
            None
        """
        data = self.data
        bb = "\033[1m"
        be = "\033[0m"
        if(filter):
            filter_res = self.filter_by_column(data=self.construction_data, col=col, equals=equals)
            if('ERROR:' in str(filter_res)):
                return filter_res
            query_wells = list(data.STATION_ID.unique())
            filter_wells = list(filter_res.index.unique())
            intersect_wells = list(set(query_wells) & set(filter_wells))
            if(len(intersect_wells)<=0):
                return 'ERROR: No results for this query with the specifed filter parameters.'
            data = data[data['STATION_ID'].isin(intersect_wells)]
        
        if(well_name==None):
            wells = list(data.STATION_ID.unique())
        else:
            wells = [well_name]
        for well in wells:
            print("{}{}{}".format(bb,str(well), be))
            analytes = sorted(list(data[data.STATION_ID==well].ANALYTE_NAME.unique()))
            print(str(analytes) +'\n')


    def query_data(self, well_name, analyte_name):
        """Filters data by passing the data and specifying the well_name and analyte_name

        Args:
            well_name (str): name of the well to be processed
            analyte_name (str): name of the analyte to be processed

        Returns:
            pd.DataFrame: filtered data based on query conditons
        """
        data = self.data
        query = data[data.STATION_ID == well_name]
        query = query[query.ANALYTE_NAME == analyte_name]
        if(query.shape[0] == 0):
            return 0
        else:
            return query
    
    def plot_data(self, well_name, analyte_name, log_transform=True, alpha=0,
              plot_inline=True, year_interval=2, x_label='Years', y_label='', save_dir='plot_data', filter=False, col=None, equals=[]):
        """Plot concentrations over time of a specified well and analyte with a smoothed curve on interpolated data points.

        Args:

            well_name (str): name of the well to be processed
            analyte_name (str): name of the analyte to be processed
            log_transform (bool, optional): choose whether or not the data should be transformed to log base 10 values. Defaults to True.
            alpha (int, optional): alue between 0 and 10 for line smoothing. Defaults to 0.
            plot_inline (bool, optional): choose whether or not to show plot inline. Defaults to True.
            year_interval (int, optional): plot by how many years to appear in the axis e.g.(1 = every year, 5 = every 5 years, ...). Defaults to 2.
            x_label (str, optional): x axis label. Defaults to 'Years'.
            y_label (str, optional): y axis label. Defaults to ''.
            save_dir (str, optional): name of the directory you want to save the plot to. Defaults to 'plot_data'.
            filter (bool, optional): flag to indicate filtering. Defaults to False.
            col (str, optional): column to filter. Example: col='STATION_ID'. Defaults to None.
            equals (list, optional): values to filter col by. Examples: equals=['FAI001A', 'FAI001B']. Defaults to [].

        Returns:
            None
        """
    
        # Gets appropriate data (well_name and analyte_name)
        query = self.query_data(well_name, analyte_name)
        query = self.simplify_data(data=query)

        if(type(query)==int and query == 0):
            return 'No results found for {} and {}'.format(well_name, analyte_name)
        else:
            if(filter):
                filter_res = self.filter_by_column(data=self.construction_data, col=col, equals=equals)
                if('ERROR:' in str(filter_res)):
                    return filter_res
                query_wells = list(query.STATION_ID.unique())
                filter_wells = list(filter_res.index.unique())
                intersect_wells = list(set(query_wells) & set(filter_wells))
                if(len(intersect_wells)<=0):
                    return 'ERROR: No results for this query with the specifed filter parameters.'
                query = query[query['STATION_ID'].isin(intersect_wells)]
            x_data = query.COLLECTION_DATE
            x_data = pd.to_datetime(x_data)
            y_data = query.RESULT
            if(log_transform):
                y_data = np.log10(y_data)

            x_RR = x_data.astype('int64').to_numpy()

            nu = x_data.shape[0]

            model = SuperSmoother(alpha=alpha)
            model.fit(x_RR, y_data)
            y_pred = model.predict(x_RR)
            r = model.cv_residuals()
            out = abs(r) > 2.2*np.std(r)
            out_x = x_data[out]
            out_y = y_data[out]

            plt.figure(figsize=(8,8))
            ax = plt.axes()
            years = mdates.YearLocator(year_interval)  # every year
            months = mdates.MonthLocator()  # every month
            yearsFmt = mdates.DateFormatter('%Y')

            for label in ax.get_xticklabels():
                label.set_rotation(30)
                label.set_horizontalalignment('center')

            ax = plt.gca()
            ax.xaxis.set_major_locator(years)
            ax.xaxis.set_major_formatter(yearsFmt)
            ax.autoscale_view()

            unit = query.RESULT_UNITS.values[0]

            ax.set_title(str(well_name) + ' - ' + analyte_name, fontweight='bold')
            ttl = ax.title
            ttl.set_position([.5, 1.05])
            if(y_label==''):    
                if(log_transform):
                    ax.set_ylabel('log-Concentration (' + unit + ')')
                else:
                    ax.set_ylabel('Concentration (' + unit + ')')
            else:
                ax.set_ylabel(y_label)
            ax.set_xlabel(x_label)
            small_fontSize = 15
            large_fontSize = 20
            plt.rc('axes', titlesize=large_fontSize)
            plt.rc('axes', labelsize=large_fontSize)
            plt.rc('legend', fontsize=small_fontSize)
            plt.rc('xtick', labelsize=small_fontSize)
            plt.rc('ytick', labelsize=small_fontSize) 
            ax.plot(x_data, y_data, ls='', marker='o', ms=5, color='black', alpha=1)
            ax.plot(x_data, y_pred, ls='-', marker='', ms=5, lw=2, color='blue', alpha=0.5, label="Super Smoother")
            ax.plot(out_x , out_y, ls='', marker='o', ms=5, color='red', alpha=1, label="Outliers")
            ax.legend(bbox_to_anchor=(1.04, 1), loc='upper left', borderaxespad=0.)
            props = dict(boxstyle='round', facecolor='grey', alpha=0.15)       
            ax.text(1.05, 0.85, 'Samples: {}'.format(nu), transform=ax.transAxes, 
                    fontsize=small_fontSize,
                    fontweight='bold',
                    verticalalignment='top', 
                    bbox=props)
            if not os.path.exists(save_dir):
                os.makedirs(save_dir)
            # plt.savefig(save_dir + '/' + str(well_name) + '-' + analyte_name +'.png', bbox_inches="tight")
            if(plot_inline):
                plt.show()
            plt.clf()
            plt.cla()
            plt.close()


    def plot_all_data(self, log_transform=True, alpha=0, year_interval=2, plot_inline=True, save_dir='plot_data'):
        """Plot concentrations over time for every well and analyte with a smoothed curve on interpolated data points.

        Args:
            log_transform (bool, optional): choose whether or not the data should be transformed to log base 10 values. Defaults to True.
            alpha (int, optional): alue between 0 and 10 for line smoothing. Defaults to 0.
            plot_inline (bool, optional): choose whether or not to show plot inline. Defaults to True.
            year_interval (int, optional): plot by how many years to appear in the axis e.g.(1 = every year, 5 = every 5 years, ...). Defaults to 2.
            save_dir (str, optional): name of the directory you want to save the plot to. Defaults to 'plot_data'.
        """
        analytes = ['TRITIUM','URANIUM-238','IODINE-129','SPECIFIC CONDUCTANCE', 'PH', 'DEPTH_TO_WATER']
        wells = np.array(self.data.STATION_ID.values)
        wells = np.unique(wells)
        success = 0
        errors = 0
        for well in wells:
            for analyte in analytes:
                plot = self.plot_data(well, analyte, 
                                    log_transform=log_transform, 
                                    alpha=alpha, 
                                    year_interval=year_interval,
                                    plot_inline=plot_inline,
                                    save_dir=save_dir)
                if 'ERROR:' in str(plot):
                    errors = errors + 1
                else:
                    success = success + 1
        print("Success: ", success)
        print("Errors: ", errors)

    def plot_correlation_heatmap(self, well_name, show_symmetry=True, color=True, save_dir='plot_correlation_heatmap'):
        """ Plots a heatmap of the correlations of the important analytes over time for a specified well.

        Args:
            well_name (str): name of the well to be processed
            show_symmetry (bool, optional): choose whether or not the heatmap should show the same information twice over the diagonal. Defaults to True.
            color (bool, optional): choose whether or not the plot should be in color or in greyscale. Defaults to True.
            save_dir (str, optional): name of the directory you want to save the plot to. Defaults to 'plot_correlation_heatmap'.

        Returns:
            None
        """
        data = self.data
        query = data[data.STATION_ID == well_name]
        a = list(np.unique(query.ANALYTE_NAME.values))
        b = ['TRITIUM','IODINE-129','SPECIFIC CONDUCTANCE', 'PH','URANIUM-238', 'DEPTH_TO_WATER']
        analytes = self.__custom_analyte_sort(list(set(a) and set(b)))
        query = query.loc[query.ANALYTE_NAME.isin(analytes)]
        analytes = self.__custom_analyte_sort(np.unique(query.ANALYTE_NAME.values))
        x = query[['COLLECTION_DATE', 'ANALYTE_NAME']]
        unique = ~x.duplicated()
        query = query[unique]
        piv = query.reset_index().pivot(index='COLLECTION_DATE',columns='ANALYTE_NAME', values='RESULT')
        piv = piv[analytes]
        totalSamples = piv.shape[0]
        piv = piv.dropna()
        samples = piv.shape[0]
        if(samples < 5):
            return 'ERROR: {} does not have enough samples to plot.'.format(well_name)
        else:
            scaler = StandardScaler()
            pivScaled = scaler.fit_transform(piv)
            pivScaled = pd.DataFrame(pivScaled, columns=piv.columns)
            pivScaled.index = piv.index
            corr_matrix = pivScaled.corr()
            if(show_symmetry):
                mask = None
            else:
                mask = np.triu(corr_matrix)
            if(color):
                cmap = 'RdBu'
            else:
                cmap = 'binary'
            fig, ax = plt.subplots(figsize=(8,6))
            ax.set_title(well_name + '_correlation', fontweight='bold')
            ttl = ax.title
            ttl.set_position([.5, 1.05])
            props = dict(boxstyle='round', facecolor='grey', alpha=0.15)
            ax.text(1.3, 1.05, 'Start date:  {}\nEnd date:    {}\n\nSamples:     {} of {}'.format(piv.index[0], piv.index[-1], samples, totalSamples), transform=ax.transAxes, fontsize=15, fontweight='bold', verticalalignment='bottom', bbox=props)
            ax = sns.heatmap(corr_matrix,
                                    ax=ax,
                                    mask=mask,
                                    vmin=-1, vmax=1,
                                    xticklabels=corr_matrix.columns,
                                    yticklabels=corr_matrix.columns,
                                    cmap=cmap,
                                    annot=True,
                                    linewidths=1,
                                    cbar_kws={'orientation': 'vertical'})
            if not os.path.exists(save_dir):
                os.makedirs(save_dir)
            fig.savefig(save_dir + '/' + well_name + '_correlation.png', bbox_inches="tight")

    def plot_all_correlation_heatmap(self, show_symmetry=True, color=True, save_dir='plot_correlation_heatmap'):
        """Plots a heatmap of the correlations of the important analytes over time for each well in the dataset.

        Args:
            show_symmetry (bool, optional): choose whether or not the heatmap should show the same information twice over the diagonal. Defaults to True.
            color (bool, optional): choose whether or not the plot should be in color or in greyscale. Defaults to True.
            save_dir (str, optional): name of the directory you want to save the plot to. Defaults to 'plot_correlation_heatmap'.
        """
        data = self.data
        wells = np.array(data.STATION_ID.values)
        wells = np.unique(wells)
        for well in wells:
            self.plot_correlation_heatmap(well_name=well,
                                            show_symmetry=show_symmetry,
                                            color=color,
                                            save_dir=save_dir)

    def interpolate_well_data(self, well_name, analytes, frequency='2W'):
        """Resamples the data based on the frequency specified and interpolates the values of the analytes.

        Args:
            well_name (str): name of the well to be processed.
            analytes (list): list of analyte names to use
            frequency (str, optional): {‘D’, ‘W’, ‘M’, ‘Y’} frequency to interpolate. See https://pandas.pydata.org/pandas-docs/stable/user_guide/timeseries.html for valid frequency inputs. (e.g. ‘W’ = every week, ‘D ’= every day, ‘2W’ = every 2 weeks). Defaults to '2W'.

        Returns:
            pd.DataFrame
        """
        data = self.data
        inter_series = {}
        query = data[data.STATION_ID == well_name]
        for analyte in analytes:
            series = query[query.ANALYTE_NAME == analyte]
            series = (series[['COLLECTION_DATE', 'RESULT']])
            series.COLLECTION_DATE = pd.to_datetime(series.COLLECTION_DATE)
            series.index = series.COLLECTION_DATE
            original_dates = series.index
            series = series.drop('COLLECTION_DATE', axis=1)
            series = series.rename({'RESULT': analyte}, axis=1)
            upsampled = series.resample(frequency).mean()
            interpolated = upsampled.interpolate(method='linear', order=2)
            inter_series[analyte] = interpolated
        join = inter_series[analytes[0]]
        join = join.drop(analytes[0], axis=1)
        for analyte in analytes:
            join = join.join(inter_series[analyte])
        join = join.dropna()
        return join

    def plot_corr_by_well(self, well_name, analytes, plot_figure=True, remove_outliers=True, z_threshold=4,
                          interpolate=False, frequency='2W', save_dir='plot_correlation',
                          log_transform=False, fontsize=20, return_data=False, remove=[], no_log=None):
        """Plots the correlations with the physical plots as well as the correlations of the important analytes over time for a specified well.

        Args:
            well_name (str): name of the well to be processed
            analytes (list): list of analyte names to use
            remove_outliers (bool, optional): choose whether or to remove the outliers. Defaults to True.
            z_threshold (int, optional): z_score threshold to eliminate outliers. Defaults to 4.
            interpolate (bool, optional): choose whether or to interpolate the data. Defaults to False.
            frequency (str, optional): {‘D’, ‘W’, ‘M’, ‘Y’} frequency to interpolate. Note: See https://pandas.pydata.org/pandas-docs/stable/user_guide/timeseries.html for valid frequency inputs. (e.g. ‘W’ = every week, ‘D ’= every day, ‘2W’ = every 2 weeks). Defaults to '2W'.
            save_dir (str, optional): name of the directory you want to save the plot to. Defaults to 'plot_correlation'.
            log_transform (bool, optional): flag for log base 10 transformation. Defaults to False.
            fontsize (int, optional): font size. Defaults to 20.
            returnData (bool, optional): flag to return data used to perfrom correlation analysis. Defaults to False.
            remove (list, optional): wells to remove. Defaults to [].
            no_log (list, optional): list of column names to not apply log transformation to. Defaults to None.

        Returns:
            None
        """
        data = self.data
        query = data[data.STATION_ID == well_name]
        a = list(np.unique(query.ANALYTE_NAME.values))    # get all analytes from dataset
        for value in analytes:
            if((value in a)==False):
                return 'ERROR: No analyte named "{}" in data.'.format(value)
        analytes = sorted(analytes)
        query = query.loc[query.ANALYTE_NAME.isin(analytes)]
        x = query[['COLLECTION_DATE', 'ANALYTE_NAME']]
        unique = ~x.duplicated()
        query = query[unique]
        piv = query.reset_index().pivot(index='COLLECTION_DATE',columns='ANALYTE_NAME', values='RESULT')
        piv = piv[analytes]
        piv.index = pd.to_datetime(piv.index)
        totalSamples = piv.shape[0]
        piv = piv.dropna()
        if(interpolate):
            piv = self.interpolate_well_data(well_name, analytes, frequency=frequency)
            file_extension = '_interpolated_' + frequency
            title = well_name + '_correlation - interpolated every ' + frequency
        else:
            file_extension = '_correlation'
            title = well_name + '_correlation'
        samples = piv.shape[0]
        if(samples < 5):
            if(interpolate):
                return 'ERROR: {} does not have enough samples to plot.\n Try a different interpolation frequency'.format(well_name)
            return 'ERROR: {} does not have enough samples to plot.'.format(well_name)
        else:
            # scaler = StandardScaler()
            # pivScaled = scaler.fit_transform(piv)
            # pivScaled = pd.DataFrame(pivScaled, columns=piv.columns)
            # pivScaled.index = piv.index
            # piv = pivScaled
            
            if(log_transform):
                piv[piv <= 0] = 0.00000001
                temp = piv.copy()
                piv = np.log10(piv)
                if(no_log !=None):
                    for col in no_log:
                        piv[col] = temp[col]

            # Remove outliers
            if(remove_outliers):
                piv = self.remove_outliers(piv, z_threshold=z_threshold)
            samples = piv.shape[0]

            idx = piv.index.date
            dates = [dates.strftime('%Y-%m-%d') for dates in idx]
            remaining = [i for i in dates if i not in remove]
            piv = piv.loc[remaining]

            if plot_figure:
            
                sns.set_style("white", {"axes.facecolor": "0.95"})
                g = sns.PairGrid(piv, aspect=1.2, diag_sharey=False, despine=False)
                g.fig.suptitle(title, fontweight='bold', y=1.08, fontsize=25)
                g.map_lower(sns.regplot, lowess=True, ci=False, line_kws={'color': 'red', 'lw': 3},
                                                                scatter_kws={'color': 'black', 's': 20})
                g.map_diag(sns.distplot, kde_kws={'color': 'black', 'lw': 3}, hist_kws={'histtype': 'bar', 'lw': 2, 'edgecolor': 'k', 'facecolor':'grey'})
                g.map_upper(self.__plotUpperHalf)
                for ax in g.axes.flat:
                    ax.tick_params("y", labelrotation=0, labelsize=fontsize)
                    ax.set_xticklabels(ax.get_xticklabels(), rotation=45, fontsize=fontsize)
                    ax.set_xlabel(ax.get_xlabel(), fontsize=fontsize, fontweight='bold') #HERE
                    ax.set_ylabel(ax.get_ylabel(), fontsize=fontsize,fontweight='bold')
                    
                g.fig.subplots_adjust(wspace=0.3, hspace=0.3)
                ax = plt.gca()

                props = dict(boxstyle='round', facecolor='grey', alpha=0.15)
                ax.text(1.3, 6.2, 'Start date:  {}\nEnd date:    {}\n\nOriginal samples:     {}\nSamples used:     {}'.format(piv.index[0].date(), piv.index[-1].date(), totalSamples, samples), transform=ax.transAxes, fontsize=20, fontweight='bold', verticalalignment='bottom', bbox=props)
                # Add titles to the diagonal axes/subplots
                for ax, col in zip(np.diag(g.axes), piv.columns):
                    ax.set_title(col, y=0.82, fontsize=15)
                if not os.path.exists(save_dir):
                    os.makedirs(save_dir)
                g.fig.savefig(save_dir + '/' + well_name + file_extension + '.png', bbox_inches="tight")

            if return_data:
                return piv
            

    def plot_all_corr_by_well(self, analytes, remove_outliers=True, z_threshold=4, interpolate=False, frequency='2W', save_dir='plot_correlation', log_transform=False, fontsize=20):
        """Plots the correlations with the physical plots as well as the important analytes over time for each well in the dataset.

        Args:
            analytes (list): list of analyte names to use
            remove_outliers (bool, optional): choose whether or to remove the outliers. Defaults to True.
            z_threshold (int, optional): z_score threshold to eliminate outliers. Defaults to 4.
            interpolate (bool, optional): choose whether or to interpolate the data. Defaults to False.
            frequency (str, optional): {‘D’, ‘W’, ‘M’, ‘Y’} frequency to interpolate. Note: See https://pandas.pydata.org/pandas-docs/stable/user_guide/timeseries.html for valid frequency inputs. (e.g. ‘W’ = every week, ‘D ’= every day, ‘2W’ = every 2 weeks). Defaults to '2W'.
            save_dir (str, optional): name of the directory you want to save the plot to. Defaults to 'plot_correlation'.
            log_transform (bool, optional): flag for log base 10 transformation. Defaults to False.
            fontsize (int, optional): font size. Defaults to 20.
        """
        data = self.data
        wells = np.array(data.STATION_ID.values)
        wells = np.unique(wells)
        for well in wells:
            self.plot_corr_by_well(well_name=well, analytes=analytes,remove_outliers=remove_outliers, z_threshold=z_threshold, interpolate=interpolate, frequency=frequency, save_dir=save_dir, log_transform=log_transform, fontsize=fontsize)
        
    def plot_corr_by_date_range(self, date, analytes, lag=0, min_samples=10, save_dir='plot_corr_by_date', log_transform=False, fontsize=20, returnData=False, no_log=None):
        """Plots the correlations with the physical plots as well as the correlations of the important analytes for ALL the wells on a specified date or range of dates if a lag greater than 0 is specifed.

        Args:
            date (str): date to be analyzed
            analytes (_type_): list of analyte names to use
            lag (int, optional): number of days to look ahead and behind the specified date (+/-). Defaults to 0.
            min_samples (int, optional): minimum number of samples the result should contain in order to execute.. Defaults to 10.
            save_dir (str, optional): name of the directory you want to save the plot to. Defaults to 'plot_corr_by_date'.
            log_transform (bool, optional): flag for log base 10 transformation. Defaults to False.
            fontsize (int, optional): font size. Defaults to 20.
            returnData (bool, optional): flag to return data used to perfrom correlation analysis. Defaults to False.
            no_log (list, optional): list of column names to not apply log transformation to. Defaults to None.
        """
        if(lag==0):
            data = self.data
            data = self.simplify_data(data=data)
            query = data[data.COLLECTION_DATE == date]
            a = list(np.unique(query.ANALYTE_NAME.values))# get all analytes from dataset
            for value in analytes:
                if((value in a)==False):
                    return 'ERROR: No analyte named "{}" in data.'.format(value)
            analytes = sorted(analytes)
            query = query.loc[query.ANALYTE_NAME.isin(analytes)]
            if(query.shape[0] == 0):
                return 'ERROR: {} has no data for all of the analytes.'.format(date)
            samples = query[['COLLECTION_DATE', 'STATION_ID', 'ANALYTE_NAME']].duplicated().value_counts()[0]
            if(samples < min_samples):
                return 'ERROR: {} does not have at least {} samples.'.format(date, min_samples)
            else:
                piv = query.reset_index().pivot_table(index = 'STATION_ID', columns='ANALYTE_NAME', values='RESULT',aggfunc=np.mean)
                # return piv
        else:
            # If the data has already been calculated with the lag specified, retrieve it
            if(self.jointData_is_set(lag=lag)==True): 
                data = self.__jointData[0]
            # Otherwise, calculate it
            else:
                data = self.getJointData(analytes, lag=lag)
                self.__set_jointData(data=data, lag=lag)
            # get new range based on the lag and create the pivor table to be able to do the correlation
            dateStart, dateEnd = self.__getLagDate(date, lagDays=lag)
            dateRange_key = str(dateStart.date()) + " - " + str(dateEnd.date())
            piv = pd.DataFrame(data.loc[dateRange_key]).unstack().T
            piv.index = piv.index.droplevel()
            piv = pd.DataFrame(piv).dropna(axis=0, how='all')
            num_NaNs = int(piv.isnull().sum().sum())
            samples = (piv.shape[0]*piv.shape[1])-num_NaNs
            for col in piv.columns:
                piv[col] = piv[col].astype('float64', errors = 'raise')
            if(lag>0):
                date = dateRange_key
            # return piv
        title = date + '_correlation'
        # scaler = StandardScaler()
        # pivScaled = scaler.fit_transform(piv)
        # pivScaled = pd.DataFrame(pivScaled, columns=piv.columns)
        # pivScaled.index = piv.index
        # piv = pivScaled

        if(log_transform):
            piv[piv <= 0] = 0.00000001
            temp = piv.copy()
            piv = np.log10(piv)
            if(no_log !=None):
                for col in no_log:
                    piv[col] = temp[col]

        sns.set_style("white", {"axes.facecolor": "0.95"})
        g = sns.PairGrid(piv, aspect=1.2, diag_sharey=False, despine=False)
        g.fig.suptitle(title, fontweight='bold', y=1.08, fontsize=25)
        g.map_lower(sns.regplot, lowess=True, ci=False, line_kws={'color': 'red', 'lw': 3},
                                                        scatter_kws={'color': 'black', 's': 20})
        g.map_diag(sns.distplot, kde_kws={'color': 'black', 'lw': 3}, hist_kws={'histtype': 'bar', 'lw': 2, 'edgecolor': 'k', 'facecolor':'grey'})
        g.map_upper(self.__plotUpperHalf)
        for ax in g.axes.flat:
                ax.tick_params("y", labelrotation=0, labelsize=fontsize)
                ax.set_xticklabels(ax.get_xticklabels(), rotation=45, fontsize=fontsize)
                ax.set_xlabel(ax.get_xlabel(), fontsize=fontsize, fontweight='bold') #HERE
                ax.set_ylabel(ax.get_ylabel(), fontsize=fontsize,fontweight='bold')
        g.fig.subplots_adjust(wspace=0.3, hspace=0.3)
        ax = plt.gca()

        props = dict(boxstyle='round', facecolor='grey', alpha=0.15)
        ax.text(1.3, 3, 'Date:  {}\n\nWells:     {}\nSamples used:     {}'.format(date, piv.shape[0] ,samples), transform=ax.transAxes, fontsize=20, fontweight='bold', verticalalignment='bottom', bbox=props)
        # Add titles to the diagonal axes/subplots
        for ax, col in zip(np.diag(g.axes), piv.columns):
            ax.set_title(col, y=0.82, fontsize=15)
        if not os.path.exists(save_dir):
            os.makedirs(save_dir)
        g.fig.savefig(save_dir + '/' + date + '.png', bbox_inches="tight")
        if(returnData):
            return piv


    def plot_corr_by_year(self, queried_data, year, analytes, quarter=None, remove_outliers=True, z_threshold=4, min_samples=10,
                          save_dir='plot_corr_by_year', file_suffix="", log_transform=False, 
                          fontsize=20, return_data=False, no_log=None):
        """
        Plots correlation heatmaps and scatter matrices for specified analytes in a given year.
        
        Parameters:
            queried_data (pd.DataFrame): DataFrame containing the data to be analyzed.
            year (int): Year of data to analyze.
            analytes (list): List of analyte names to include in the plot.
            quarter (int): Quarter of the year to filter data. Default is None (no filtering).
            remove_outliers (bool): Whether to remove outliers using z-score threshold. Default is True.
            z_threshold (int): Z-score threshold for outlier removal. Default is 4.
            min_samples (int): Minimum number of samples required to generate plots. Default is 10.
            save_dir (str): Directory to save plots. Default is 'plot_corr_by_year'.
            file_suffix (str): Suffix to add to saved plot filenames.
            log_transform (bool): Whether to apply log10 transformation to data. Default is False.
            fontsize (int): Font size for plot labels and titles. Default is 20.
            return_data (bool): If True, returns the processed pivot DataFrame. Default is False.
            no_log (list): List of columns to exclude from log transformation. Default is None.

        Returns:
            pd.DataFrame (optional): Pivot table used for correlation analysis if return_data is True.
        """
        # Prepare and filter data
        # query = self.simplify_data(self.data)
        # query['COLLECTION_DATE'] = pd.to_datetime(query['COLLECTION_DATE'])
        # query = query[query['COLLECTION_DATE'].dt.year == year]

        query = queried_data.copy()

        available_analytes = np.unique(query['ANALYTE_NAME'].values)
        missing_analytes = [a for a in analytes if a not in available_analytes]
        if missing_analytes:
            return f"ERROR: Missing analytes in data: {', '.join(missing_analytes)}"
        
        query = query[query['ANALYTE_NAME'].isin(analytes)]
        if query.empty:
            return f"ERROR: {year} has no data for the selected analytes."

        samples = query[['COLLECTION_DATE', 'STATION_ID', 'ANALYTE_NAME']].duplicated().value_counts().get(0, 0)
        if samples < min_samples:
            return f"ERROR: {year} does not have at least {min_samples} samples." 
        
        # Pivot data for correlation analysis
        piv = query.pivot_table(index='STATION_ID', columns='ANALYTE_NAME', values='RESULT', aggfunc=np.mean)

        # Remove outliers if requested
        if remove_outliers:
            piv = self.remove_outliers(piv, z_threshold=z_threshold)
        samples = piv.count().sum()

        # Log transformation if requested
        if log_transform:
            piv[piv <= 0] = 1e-8  # Replace non-positive values before log
            temp_piv = piv.copy()
            piv = np.log10(piv)
            if no_log:
                for col in no_log:
                    piv[col] = temp_piv[col]

        # Set 
        if quarter is not None:
            time_string = f"{year}_Q{quarter}"
        else:
            time_string = f"{year}"

        # Plotting using Seaborn PairGrid
        sns.set_style("white", {"axes.facecolor": "0.95"})
        g = sns.PairGrid(piv, aspect=1.2, diag_sharey=False, despine=False)
        title = f"{time_string}_correlation{file_suffix}"
        g.fig.suptitle(title, fontweight='bold', y=1.08, fontsize=25)

        g.map_lower(sns.regplot, lowess=True, ci=False, 
                    line_kws={'color': 'red', 'lw': 3},
                    scatter_kws={'color': 'black', 's': 20})

        g.map_diag(sns.histplot, kde=True, stat='density', kde_kws={'cut':3}, alpha=0.4, color='grey', edgecolor='black', linewidth=1)
        # g.map_diag(sns.distplot, kde_kws={'color': 'black', 'lw': 3}, hist_kws={'histtype': 'bar', 'lw': 2, 'edgecolor': 'k', 'facecolor':'grey'})
        g.map_upper(self.__plotUpperHalf)

        # Customize axes labels and titles
        for ax in g.axes.flat:
            ax.tick_params(labelsize=fontsize)
            ax.set_xticklabels(ax.get_xticklabels(), rotation=45, fontsize=fontsize)
            ax.set_xlabel(ax.get_xlabel(), fontsize=fontsize, fontweight='bold')
            ax.set_ylabel(ax.get_ylabel(), fontsize=fontsize, fontweight='bold')

        # Add per-axis titles on the diagonal plots
        for ax_diag, col in zip(np.diag(g.axes), piv.columns):
            ax_diag.set_title(col, y=0.82, fontsize=15)

        # Add annotation with metadata
        ax_main = plt.gca()
        props = dict(boxstyle='round', facecolor='grey', alpha=0.15)
        ax_main.text(1.3, 3, f'Date: {time_string.replace("_"," ")}\n\nSamples Used: {samples}', 
                    transform=ax_main.transAxes, fontsize=20, fontweight='bold', 
                    verticalalignment='bottom', bbox=props)

        g.fig.subplots_adjust(wspace=0.3, hspace=0.3)

        # Save figure
        os.makedirs(save_dir, exist_ok=True)
        file_name = f"{time_string}{file_suffix}.png"
        save_path = os.path.join(save_dir, file_name)
        g.fig.savefig(save_path, bbox_inches="tight")

        if return_data:
            return piv

            
    def plot_MCL(self, well_name, analyte_name, year_interval=5, save_dir='plot_MCL'):
        """Plots the linear regression line of data given the analyte_name and well_name. The plot includes the prediction where the line of best fit intersects with the Maximum Concentration Limit (MCL).

        Args:
            well_name (str): ame of the well to be processed
            analyte_name (str): name of the analyte to be processed
            year_interval (int, optional): lot by how many years to appear in the axis e.g.(1 = every year, 5 = every 5 years, ...). Defaults to 5.
            save_dir (str, optional): name of the directory you want to save the plot to. Defaults to 'plot_MCL'.
        """
        data = self.data
        # finds the intersection point of 2 lines given the slopes and y-intercepts
        def line_intersect(m1, b1, m2, b2):
            if m1 == m2:
                print ('The lines are parallel')
                return None
            x = (b2 - b1) / (m1 - m2)
            y = m1 * x + b1
            return x,y

        # Gets appropriate data (well_name and analyte_name)
        query = self.query_data(well_name, analyte_name)
        query = query[query.RESULT > 0]  # drop non-positive values

        if(type(query)==int and query == 0):
            return 'No results found for {} and {}'.format(well_name, analyte_name)
        else:   

            test = query.groupby(['COLLECTION_DATE'])[['RESULT']].mean()
            test.index = pd.to_datetime(test.index)

            x = date2num(test.index)
            y = np.log10(test.RESULT)
            ylabel = 'log-Concentration (' + self.get_unit(analyte_name) + ')'
            y = y.rename(ylabel)

            p, cov = np.polyfit(x, y, 1, cov=True)  # parameters and covariance from of the fit of 1-D polynom.

            m_unc = np.sqrt(cov[0][0])
            b_unc = np.sqrt(cov[1][1])

            f = np.poly1d(p)

            try:
                MCL = self.get_MCL(analyte_name)
                m1, b1 = f # line of best fit
                m2, b2 = 0, MCL # MCL constant

                intersection = line_intersect(m1, b1, m2, b2)

                ## Get confidence interval intersection points with MCL
                data = list(zip(x,y))
                n = len(data)
                list_slopes = []
                list_intercepts = []
                random.seed(50)
                for _ in range(80):
                    sampled_data = [ random.choice(data) for _ in range(n) ]
                    x_s, y_s = zip(*sampled_data)
                    x_s = np.array(x_s)
                    y_s = np.array(y_s)

                    m_s, b_s, r, p, err = scipy.stats.linregress(x_s,y_s)
                    ymodel = m_s*x_s + b_s
                    list_slopes.append(m_s)
                    list_intercepts.append(b_s)

                max_index = list_slopes.index(max(list_slopes))
                min_index = list_slopes.index(min(list_slopes))
                intersection_left = line_intersect(list_slopes[min_index], list_intercepts[min_index], m2, b2)
                intersection_right = line_intersect(list_slopes[max_index], list_intercepts[max_index], m2, b2)
                MCL_date_lower = min(intersection_left[0], intersection_right[0])
                MCL_date_upper = max(intersection_left[0], intersection_right[0])

                fig, ax = plt.subplots(figsize=(10, 6))

                ax.set_title(well_name + ' - ' + analyte_name, fontweight='bold')
                ttl = ax.title
                ttl.set_position([.5, 1.05])
                years = mdates.YearLocator(year_interval)  # every year
                months = mdates.MonthLocator()  # every month
                yearsFmt = mdates.DateFormatter('%Y') 

                ax.xaxis.set_major_locator(years)
                ax = plt.gca()
                ax.xaxis.set_major_locator(years)
                ax.xaxis.set_major_formatter(yearsFmt)
                ax.autoscale_view()
                ax.grid(True, alpha=0.4)
                small_fontSize = 15
                large_fontSize = 20
                plt.rc('axes', titlesize=large_fontSize)
                plt.rc('axes', labelsize=large_fontSize)
                plt.rc('legend', fontsize=small_fontSize)
                plt.rc('xtick', labelsize=small_fontSize)
                plt.rc('ytick', labelsize=small_fontSize)

                ax.set_xlabel('Years')
                ax.set_ylabel('$log_{10}$-Concentration (' + self.get_unit(analyte_name) + ')')

                if(intersection[0] < min(x)):
                    ax.set_ylim([min(y)-1, max(y)+1])
                    ax.set_xlim([MCL_date_lower-1000, max(x)+1000])
                elif(intersection[0] < max(x) and intersection[0] > min(x)):
                    ax.set_ylim([min(y)-1, max(y)+1])
                    ax.set_xlim(min(x)-1000, max(x)+1000)
                else:
                    ax.set_ylim([min(y)-1, max(y)+1])
                    ax.set_xlim([min(x)-1000, MCL_date_upper+1000])

                ax = sns.regplot(x=x, y=y, logx=False, truncate=False, seed=42, n_boot=1000, ci=95) # Line of best fit
                ax.plot(x, y, ls='', marker='o', ms=5, color='black', alpha=1) # Data
                ax.axhline(y=MCL, color='r', linestyle='--') # MCL
                ax.plot(intersection[0], intersection[1], color='blue', marker='o', ms=10)
                ax.plot(intersection_left[0], intersection_left[1], color='green', marker='o', ms=5)
                ax.plot(intersection_right[0], intersection_right[1], color='green', marker='o', ms=5)

                predict = num2date(intersection[0]).date()
                l_predict = num2date(MCL_date_lower).date()
                u_predict = num2date(MCL_date_upper).date()
                ax.annotate(predict, (intersection[0], intersection[1]), xytext=(intersection[0], intersection[1]+1), 
                            bbox=dict(boxstyle="round", alpha=0.1),ha='center', arrowprops=dict(arrowstyle="->", color='blue'), fontsize=small_fontSize, fontweight='bold')
                props = dict(boxstyle='round', facecolor='grey', alpha=0.15)
                ax.text(1.1, 0.5, 'Lower confidence:  {}\n            Prediction:  {}\nUpper confidence:  {}'.format(l_predict, predict, u_predict), transform=ax.transAxes, fontsize=small_fontSize, fontweight='bold', verticalalignment='bottom', bbox=props)

                if not os.path.exists(save_dir):
                    os.makedirs(save_dir)
                plt.savefig(save_dir + '/' + well_name + '-' + analyte_name +'.png', bbox_inches="tight")

            except Exception as e:
                print(f"{well_name}: {e}")
                return None

    def plot_PCA_by_date(self, date, analytes, lag=0, n_clusters=4, return_clusters=False, min_samples=3, show_labels=True, save_dir='plot_PCA_by_date', filter=False, col=None, equals=[]):
        """Gernates a PCA biplot (PCA score plot + loading plot) of the data given a date in the dataset. The data is also clustered into n_clusters.

        Args:
            date (str): date to be analyzed
            analytes (str): list of analyte names to use
            lag (int, optional): number of days to look ahead and behind the specified date (+/-). Defaults to 0.
            n_clusters (int, optional): number of clusters to split the data into.. Defaults to 4.
            return_clusters (bool, optional): Flag to return the cluster data to be used for spatial plotting.. Defaults to False.
            min_samples (int, optional): minimum number of samples the result should contain in order to execute.. Defaults to 3.
            show_labels (bool, optional): choose whether or not to show the name of the wells.. Defaults to True.
            save_dir (str, optional): name of the directory you want to save the plot to. Defaults to 'plot_PCA_by_date'.
            filter (bool, optional): flag to indicate filtering. Defaults to False.
            col (str, optional): column to filter. Example: col='STATION_ID'. Defaults to None.
            equals (list, optional): values to filter col by. Examples: equals=['FAI001A', 'FAI001B']. Defaults to [].
        """
        if(lag==0):
            data = self.data
            data = self.simplify_data(data=data)
            query = data[data.COLLECTION_DATE == date]
            if(filter):
                filter_res = self.filter_by_column(data=self.construction_data, col=col, equals=equals)
                if('ERROR:' in str(filter_res)):
                    return filter_res
                query_wells = list(query.STATION_ID.unique())
                filter_wells = list(filter_res.index.unique())
                intersect_wells = list(set(query_wells) & set(filter_wells))
                if(len(intersect_wells)<=0):
                    return 'ERROR: No results for this query with the specifed filter parameters.'
                query = query[query['STATION_ID'].isin(intersect_wells)]
            a = list(np.unique(query.ANALYTE_NAME.values))# get all analytes from dataset
            for value in analytes:
                if((value in a)==False):
                    return 'ERROR: No analyte named "{}" in data.'.format(value)
            analytes = sorted(analytes)
            query = query.loc[query.ANALYTE_NAME.isin(analytes)]

            if(query.shape[0] == 0):
                return 'ERROR: {} has no data for the 6 analytes.'.format(date)
            samples = query[['COLLECTION_DATE', 'STATION_ID', 'ANALYTE_NAME']].duplicated().value_counts()[0]
            if(samples < min_samples):
                return 'ERROR: {} does not have at least {} samples.'.format(date, min_samples)
            # if(len(np.unique(query.ANALYTE_NAME.values)) < 6):
            #     return 'ERROR: {} has less than the 6 analytes we want to analyze.'.format(date)
            else:
                # analytes = self.__custom_analyte_sort(np.unique(query.ANALYTE_NAME.values))
                analytes = sorted(analytes)
                piv = query.reset_index().pivot_table(index = 'STATION_ID', columns='ANALYTE_NAME', values='RESULT',aggfunc=np.mean)
                
                # return piv
        else:
            # If the data has already been calculated with the lag specified, retrieve it
            if(self.jointData_is_set(lag=lag)==True): 
                data = self.__jointData[0]
            # Otherwise, calculate it
            else:
                data = self.getJointData(analytes, lag=lag)
                self.__set_jointData(data=data, lag=lag)
            # get new range based on the lag and create the pivor table to be able to do the correlation
            dateStart, dateEnd = self.__getLagDate(date, lagDays=lag)
            dateRange_key = str(dateStart.date()) + " - " + str(dateEnd.date())
            piv = pd.DataFrame(data.loc[dateRange_key]).unstack().T
            piv.index = piv.index.droplevel()
            piv = pd.DataFrame(piv).dropna(axis=0, how='all')
            num_NaNs = int(piv.isnull().sum().sum())
            samples = (piv.shape[0]*piv.shape[1])-num_NaNs
            for col in piv.columns:
                piv[col] = piv[col].astype('float64', errors = 'raise')
            if(lag>0):
                date = dateRange_key
            # return piv

            
        main_data = piv.dropna()
        
        scaler = StandardScaler()
        X = scaler.fit_transform(main_data)
        pca = PCA(n_components=2)
        x_new = pca.fit_transform(X)
        
        pca_points = pd.DataFrame(x_new, columns=["x1", "x2"])
        k_Means = KMeans(n_clusters=n_clusters, random_state=42)
        model = k_Means.fit(pca_points[['x1', 'x2']])
        predict = model.predict(pca_points[['x1', 'x2']])
        # attach predicted cluster to original points
        pca_points['predicted'] = model.labels_
        # Create a dataframe for cluster_centers (centroids)
        centroids = pd.DataFrame(model.cluster_centers_, columns=["x1", "x2"])
        colors = ['red', 'blue', 'orange', 'purple', 'green', 'beige', 'pink', 'black', 'cadetblue', 'lightgreen']
        pca_points['color'] = pca_points['predicted'].map(lambda p: colors[p])

        fig, ax = plt.subplots(figsize=(10,10))
        ax = plt.axes()

        small_fontSize = 15
        large_fontSize = 20
        plt.rc('axes', titlesize=large_fontSize)
        plt.rc('axes', labelsize=large_fontSize)
        plt.rc('legend', fontsize=small_fontSize)
        plt.rc('xtick', labelsize=small_fontSize)
        plt.rc('ytick', labelsize=small_fontSize)

        def myplot(score,coeff,labels=None,c='r', centroids=None):
            xs = score.iloc[:,0]
            ys = score.iloc[:,1]
            n = coeff.shape[0]
            scalex = 1.0/(xs.max() - xs.min())
            scaley = 1.0/(ys.max() - ys.min())
            scatt_X = xs * scalex
            scatt_Y = ys * scaley
            scatter = plt.scatter(scatt_X, scatt_Y, alpha=0.8, label='Wells', c=c)
            centers = plt.scatter(centroids.iloc[:,0]* scalex, centroids.iloc[:,1]* scaley,
                                    c = colors[0:n_clusters],
                                    marker='X', s=550)

            for i in range(n):
                arrow = plt.arrow(0, 0, coeff[i,0], coeff[i,1], color = 'r', alpha = 0.9, head_width=0.05, head_length=0.05, label='Loadings')
                if labels is None:
                    plt.text(coeff[i,0]* 1.15, coeff[i,1] * 1.15, "Var"+str(i+1), color = 'g', ha = 'center', va = 'center')
                else:
                    plt.text(coeff[i,0]* 1.15, coeff[i,1] * 1.15, labels[i], color = 'g', ha = 'center', va = 'bottom')

            if(show_labels):
                for x_pos, y_pos, label in zip(scatt_X, scatt_Y, main_data.index):
                    ax.annotate(label, # The label for this point
                    xy=(x_pos, y_pos), # Position of the corresponding point
                    xytext=(7, 0),     # Offset text by 7 points to the right
                    textcoords='offset points', # tell it to use offset points
                    ha='left',         # Horizontally aligned to the left
                    va='center',       # Vertical alignment is centered
                    color='black', alpha=0.8)
            plt.legend( [scatter, centers, arrow], ['Wells', 'Well centroids','Loadings'])

        samples = x_new.shape[0]*piv.shape[1]
        props = dict(boxstyle='round', facecolor='grey', alpha=0.15)
        ax.text(1.1, 0.5, 'Date:  {}\n\nSamples:          {}\nWells:               {}'.format(date, samples, x_new.shape[0]), 
                    transform=ax.transAxes, fontsize=20, fontweight='bold', verticalalignment='bottom', bbox=props)

        plt.xlim(-1,1)
        plt.ylim(-1,1)
        plt.xlabel("PC{}".format(1))
        plt.ylabel("PC{}".format(2))
        ax.set_title('PCA Biplot - ' + date, fontweight='bold')
        plt.grid(alpha=0.5)

        #Call the function. Use only the 2 PCs.
        myplot(pca_points,np.transpose(pca.components_[0:2, :]), labels=piv.columns, c=pca_points['color'], centroids=centroids)
        plt.show()
        if not os.path.exists(save_dir):
            os.makedirs(save_dir)
        fig.savefig(save_dir + '/' + 'PCA Biplot - '+ date +'.png', bbox_inches="tight")
        
        if(return_clusters):
            stations = list(main_data.index)
            color_wells = list(pca_points.color)
            def merge(list1, list2): 
                merged_list = [(list1[i], list2[i]) for i in range(0, len(list1))] 
                return merged_list
            color_df = pd.DataFrame(merge(stations, color_wells), columns=['STATION_ID', 'color'])
            if(self.get_Construction_Data==None):
                print('You need to set the GPS data first using the getConstructionData function.')
                return None
            else:
                gps_color = pd.merge(self.get_Construction_Data(), color_df, on=['STATION_ID'])
                return gps_color


    def plot_PCA_by_year(self, year, analytes, n_clusters=4, return_clusters=False, min_samples=10, show_labels=True, save_dir='plot_PCA_by_year', filter=False, col=None, equals=[]):
        """Gernates a PCA biplot (PCA score plot + loading plot) of the data given a year in the dataset. The data is also clustered into n_clusters.

        Args:
            year (int): year to be analyzed
            analytes (str): list of analyte names to use
            n_clusters (int, optional): number of clusters to split the data into.. Defaults to 4.
            return_clusters (bool, optional): Flag to return the cluster data to be used for spatial plotting.. Defaults to False.
            min_samples (int, optional): minimum number of samples the result should contain in order to execute.. Defaults to 3.
            show_labels (bool, optional): choose whether or not to show the name of the wells.. Defaults to True.
            save_dir (str, optional): name of the directory you want to save the plot to. Defaults to 'plot_PCA_by_date'.
            filter (bool, optional): flag to indicate filtering. Defaults to False.
            col (str, optional): column to filter. Example: col='STATION_ID'. Defaults to None.
            equals (list, optional): values to filter col by. Examples: equals=['FAI001A', 'FAI001B']. Defaults to [].
        """
        data = self.data
        query = self.simplify_data(data=data)
        query.COLLECTION_DATE = pd.to_datetime(query.COLLECTION_DATE)
        query = query[query.COLLECTION_DATE.dt.year == year]
        if(filter):
            filter_res = self.filter_by_column(data=self.construction_data, col=col, equals=equals)
            if('ERROR:' in str(filter_res)):
                return filter_res
            query_wells = list(query.STATION_ID.unique())
            filter_wells = list(filter_res.index.unique())
            intersect_wells = list(set(query_wells) & set(filter_wells))
            if(len(intersect_wells)<=0):
                return 'ERROR: No results for this query with the specifed filter parameters.'
            query = query[query['STATION_ID'].isin(intersect_wells)]
        a = list(np.unique(query.ANALYTE_NAME.values))# get all analytes from dataset
        for value in analytes:
            if((value in a)==False):
                return 'ERROR: No analyte named "{}" in data.'.format(value)
        analytes = sorted(analytes)
        query = query.loc[query.ANALYTE_NAME.isin(analytes)]

        if(query.shape[0] == 0):
            return 'ERROR: {} has no data for the 6 analytes.'.format(year)
        samples = query[['COLLECTION_DATE', 'STATION_ID', 'ANALYTE_NAME']].duplicated().value_counts()[0]
        if(samples < min_samples):
            return 'ERROR: {} does not have at least {} samples.'.format(year, min_samples)
        # if(len(np.unique(query.ANALYTE_NAME.values)) < 6):
        #     return 'ERROR: {} has less than the 6 analytes we want to analyze.'.format(year)
        else:
            # analytes = self.__custom_analyte_sort(np.unique(query.ANALYTE_NAME.values))
            analytes = sorted(analytes)
            piv = query.reset_index().pivot_table(index = 'STATION_ID', columns='ANALYTE_NAME', values='RESULT',aggfunc=np.mean)
            
            main_data = piv.dropna()
            # # FILTERING CODE
            # if(filter):
            #     res_wells = self.filter_wells(filter_well_by)
            #     main_data = main_data.loc[main_data.index.isin(res_wells)]
            
            scaler = StandardScaler()
            X = scaler.fit_transform(main_data)
            pca = PCA(n_components=2)
            x_new = pca.fit_transform(X)

            pca_points = pd.DataFrame(x_new, columns=["x1", "x2"])
            k_Means = KMeans(n_clusters=n_clusters, random_state=42)
            model = k_Means.fit(pca_points[['x1', 'x2']])
            predict = model.predict(pca_points[['x1', 'x2']])
            # attach predicted cluster to original points
            pca_points['predicted'] = model.labels_
            # Create a dataframe for cluster_centers (centroids)
            centroids = pd.DataFrame(model.cluster_centers_, columns=["x1", "x2"])
            colors = ['red', 'blue', 'orange', 'purple', 'green', 'beige', 'pink', 'black', 'cadetblue', 'lightgreen']
            pca_points['color'] = pca_points['predicted'].map(lambda p: colors[p])
            
            fig, ax = plt.subplots(figsize=(15,15))
            ax = plt.axes()

            small_fontSize = 15
            large_fontSize = 20
            plt.rc('axes', titlesize=large_fontSize)
            plt.rc('axes', labelsize=large_fontSize)
            plt.rc('legend', fontsize=small_fontSize)
            plt.rc('xtick', labelsize=small_fontSize)
            plt.rc('ytick', labelsize=small_fontSize) 

            def myplot(score,coeff,labels=None,c='r', centroids=None):
                xs = score[:,0]
                ys = score[:,1]
                n = coeff.shape[0]
                scalex = 1.0/(xs.max() - xs.min())
                scaley = 1.0/(ys.max() - ys.min())
                scatt_X = xs * scalex
                scatt_Y = ys * scaley
                scatter = plt.scatter(scatt_X, scatt_Y, alpha=0.8, label='Wells', c=c)
                centers = plt.scatter(centroids.iloc[:,0]* scalex, centroids.iloc[:,1]* scaley,
                                        c = colors[0:n_clusters],
                                        marker='X', s=550)
                for i in range(n):
                    arrow = plt.arrow(0, 0, coeff[i,0], coeff[i,1], color = 'r', alpha = 0.9, head_width=0.05, head_length=0.05)
                    if labels is None:
                        plt.text(coeff[i,0]* 1.15, coeff[i,1] * 1.15, "Var"+str(i+1), color = 'g', ha = 'center', va = 'center')
                    else:
                        plt.text(coeff[i,0]* 1.15, coeff[i,1] * 1.15, labels[i], color = 'g', ha = 'center', va = 'bottom')

                if(show_labels):
                    for x_pos, y_pos, label in zip(scatt_X, scatt_Y, main_data.index):
                        ax.annotate(label, # The label for this point
                        xy=(x_pos, y_pos), # Position of the corresponding point
                        xytext=(7, 0),     # Offset text by 7 points to the right
                        textcoords='offset points', # tell it to use offset points
                        ha='left',         # Horizontally aligned to the left
                        va='center', color='black', alpha=0.8)       # Vertical alignment is centered
                plt.legend( [scatter, centers, arrow], ['Wells', 'Well centroids','Loadings'])

            samples = x_new.shape[0]*piv.shape[1]    
            props = dict(boxstyle='round', facecolor='grey', alpha=0.15)
            ax.text(1.1, 0.5, 'Date:  {}\n\nSamples:          {}\nWells:               {}'.format(year,samples, x_new.shape[0]), 
                        transform=ax.transAxes, fontsize=20, fontweight='bold', verticalalignment='bottom', bbox=props)

            plt.xlim(-1,1)
            plt.ylim(-1,1)
            plt.xlabel("PC{}".format(1))
            plt.ylabel("PC{}".format(2))
            ax.set_title('PCA Biplot - ' + str(year), fontweight='bold')
            plt.grid(alpha=0.5)

            #Call the function. Use only the 2 PCs.
            myplot(x_new[:,0:2],np.transpose(pca.components_[0:2, :]), labels=piv.columns, c=pca_points['color'], centroids=centroids)

            plt.show()
            if not os.path.exists(save_dir):
                os.makedirs(save_dir)
            fig.savefig(save_dir + '/' + 'PCA Biplot - '+ str(year) +'.png', bbox_inches="tight")
            
            if(return_clusters):
                stations = list(main_data.index)
                color_wells = list(pca_points.color)
                def merge(list1, list2): 
                    merged_list = [(list1[i], list2[i]) for i in range(0, len(list1))] 
                    return merged_list
                color_df = pd.DataFrame(merge(stations, color_wells), columns=['STATION_ID', 'color'])
                if(self.get_Construction_Data==None):
                    print('You need to set the GPS data first using the setConstructionData function.')
                    return None
                else:
                    gps_color = pd.merge(self.get_Construction_Data(), color_df, on=['STATION_ID'])
                    return gps_color

    def plot_PCA_by_well(self, well_name, analytes, interpolate=False, frequency='2W', min_samples=10, show_labels=True, save_dir='plot_PCA_by_well'):
        """Gernates a PCA biplot (PCA score plot + loading plot) of the data given a well_name in the dataset. Only uses the 6 important analytes.

        Args:
            well_name (str): name of the well to be processed
            analytes (str): list of analyte names to use
            interpolate (bool, optional): choose to interpolate the data. Defaults to False.
            frequency (str, optional): {‘D’, ‘W’, ‘M’, ‘Y’} frequency to interpolate. See https://pandas.pydata.org/pandas-docs/stable/user_guide/timeseries.html for valid frequency inputs. (e.g. ‘W’ = every week, ‘D ’= every day, ‘2W’ = every 2 weeks). Defaults to '2W'.
            min_samples (int, optional): minimum number of samples the result should contain in order to execute.. Defaults to 3.
            show_labels (bool, optional): choose whether or not to show the name of the wells.. Defaults to True.
            save_dir (str, optional): name of the directory you want to save the plot to. Defaults to 'plot_PCA_by_date'.
        """
        data = self.data
        query = data[data.STATION_ID == well_name]
        a = list(np.unique(query.ANALYTE_NAME.values))# get all analytes from dataset
        for value in analytes:
            if((value in a)==False):
                return 'ERROR: No analyte named "{}" in data.'.format(value)
        analytes = sorted(analytes)
        query = query.loc[query.ANALYTE_NAME.isin(analytes)]
        x = query[['COLLECTION_DATE', 'ANALYTE_NAME']]
        unique = ~x.duplicated()
        query = query[unique]
        piv = query.reset_index().pivot(index='COLLECTION_DATE',columns='ANALYTE_NAME', values='RESULT')
        piv = piv[analytes]
        piv.index = pd.to_datetime(piv.index)
        totalSamples = piv.stack().shape[0]
        piv = piv.dropna()
        if(interpolate):
            piv = self.interpolate_well_data(well_name, analytes, frequency=frequency)
            title = 'PCA Biplot - ' + well_name + ' - interpolated every ' + frequency
        else:
            title = 'PCA Biplot - ' + well_name

        if(query.shape[0] == 0):
            return 'ERROR: {} has no data for the 6 analytes.'.format(well_name)
        samples = query[['COLLECTION_DATE', 'STATION_ID', 'ANALYTE_NAME']].duplicated().value_counts()[0]
        if(samples < min_samples):
            return 'ERROR: {} does not have at least {} samples.'.format(well_name, min_samples)
        # if(len(np.unique(query.ANALYTE_NAME.values)) < 6):
        #     return 'ERROR: {} has less than the 6 analytes we want to analyze.'.format(well_name)
        else:
            scaler = StandardScaler()
            X = scaler.fit_transform(piv.dropna())
            pca = PCA(n_components=2)
            x_new = pca.fit_transform(X)

            fig, ax = plt.subplots(figsize=(15,15))
            ax = plt.axes()

            small_fontSize = 15
            large_fontSize = 20
            plt.rc('axes', titlesize=large_fontSize)
            plt.rc('axes', labelsize=large_fontSize)
            plt.rc('legend', fontsize=small_fontSize)
            plt.rc('xtick', labelsize=small_fontSize)
            plt.rc('ytick', labelsize=small_fontSize) 

            def myplot(score,coeff,labels=None):
                xs = score[:,0]
                ys = score[:,1]
                n = coeff.shape[0]
                scalex = 1.0/(xs.max() - xs.min())
                scaley = 1.0/(ys.max() - ys.min())
                scatt_X = xs * scalex
                scatt_Y = ys * scaley
                scatter = plt.scatter(scatt_X, scatt_Y, alpha=0.8, label='Date samples')

                for i in range(n):
                    arrow = plt.arrow(0, 0, coeff[i,0], coeff[i,1], color = 'r', alpha = 0.9, head_width=0.05, head_length=0.05, label='Loadings')
                    if labels is None:
                        plt.text(coeff[i,0]* 1.15, coeff[i,1] * 1.15, "Var"+str(i+1), color = 'g', ha = 'center', va = 'center')
                    else:
                        plt.text(coeff[i,0]* 1.15, coeff[i,1] * 1.15, labels[i], color = 'g', ha = 'center', va = 'bottom')

                if(show_labels):
                    for x_pos, y_pos, label in zip(scatt_X, scatt_Y, piv.dropna().index.date):
                        ax.annotate(label, # The label for this point
                        xy=(x_pos, y_pos), # Position of the corresponding point
                        xytext=(7, 0),     # Offset text by 7 points to the right
                        textcoords='offset points', # tell it to use offset points
                        ha='left',         # Horizontally aligned to the left
                        va='center',       # Vertical alignment is centered
                        color='black', alpha=0.8)
                plt.legend( [scatter, arrow], ['Date samples', 'Loadings'])

            samples = x_new.shape[0]*piv.shape[1]
            props = dict(boxstyle='round', facecolor='grey', alpha=0.15)      
            ax.text(1.1, 0.5, 'Start date:  {}\nEnd date:    {}\n\nOriginal samples:     {}\nSamples used:     {}\nDate samples:               {}'
                        .format(piv.index[0].date(), piv.index[-1].date(), totalSamples, samples, x_new.shape[0]), 
                        transform=ax.transAxes, fontsize=20, fontweight='bold', verticalalignment='bottom', bbox=props)

            plt.xlim(-1,1)
            plt.ylim(-1,1)
            plt.xlabel("PC{}".format(1))
            plt.ylabel("PC{}".format(2))
            ax.set_title(title, fontweight='bold')
            plt.grid(alpha=0.5)

            #Call the function. Use only the 2 PCs.
            myplot(x_new[:,0:2],np.transpose(pca.components_[0:2, :]), labels=piv.columns)
            plt.show()
            if not os.path.exists(save_dir):
                os.makedirs(save_dir)
            fig.savefig(save_dir + '/' + title +'.png', bbox_inches="tight")
            
    def plot_coordinates_to_map(self, gps_data, center=[33.271459, -81.675873], zoom=14) -> folium.Map:
        """Plots the well locations on an interactive map given coordinates.

        Args:
            gps_data (pd.DataFrame): Data frame with the following column names: station_id, latitude, longitude, color. If the color column is not passed, the default color will be blue.
            center (list, optional): latitude and longitude coordinates to center the map view. Defaults to [33.271459, -81.675873].
            zoom (int, optional): value to determine the initial scale of the map. Defaults to 14.

        Returns:
            folium.Map
        """

        # get the center of the map
        center_data = [gps_data['LATITUDE'].mean(), gps_data['LONGITUDE'].mean()]

        m = folium.Map(location=center_data, zoom_start=zoom)

        # add station markers to the map
        for _, row in gps_data.iterrows():
            
            color = row['color'] if 'color' in gps_data.columns else 'blue'
            
            folium.Marker(
                location=[row['LATITUDE'], row['LONGITUDE']],
                popup=row['STATION_ID'],
                icon=folium.Icon(color=color)
            ).add_to(m)


        return m

    def interpolate_wells_by_analyte(self, analyte, frequency='2W', rm_outliers=True, z_threshold=3):
        """Resamples analyte data based on the frequency specified and interpolates the values in between. NaN values are replaced with the average value per well.

        Args:
            analyte (_type_): analyte name for interpolation of all present wells.
            frequency (str, optional): {‘D’, ‘W’, ‘M’, ‘Y’} frequency to interpolate. See https://pandas.pydata.org/pandas-docs/stable/user_guide/timeseries.html for valid frequency inputs. (e.g. ‘W’ = every week, ‘D ’= every day, ‘2W’ = every 2 weeks). Defaults to '2W'.
            rm_outliers (bool, optional): flag to remove outliers in the data. Defaults to True.
            z_threshold (int, optional): z_score threshold to eliminate outliers. Defaults to 3.

        Returns:
            pd.DataFrame: interpolated data
        """
        data = self.data
        df_t, dates = self.__transform_time_series( 
                                                    analytes=[analyte], 
                                                    resample=frequency, 
                                                    rm_outliers=True, 
                                                    z_threshold=z_threshold)
        res_interp = self.__get_individual_analyte_df(data=df_t, dates=dates, analyte=analyte)
        res_interp = res_interp.dropna(axis=1, how='all')
        return res_interp

    # IN THE WORKS
    def __transform_time_series(self, analytes=[], resample='2W', rm_outliers=False, z_threshold=4):
        data = self.data
        def transform_time_series_by_analyte(data, analyte_name):
            wells_analyte = np.unique(data[data.ANALYTE_NAME == analyte_name].STATION_ID)
            condensed = data[data.ANALYTE_NAME == analyte_name].groupby(['STATION_ID','COLLECTION_DATE']).agg({'RESULT': lambda x: x.mean()})
            analyte_df_resample = pd.DataFrame(index=wells_analyte, columns=t)
            analyte_df_resample.sort_index(inplace=True)
            for well in wells_analyte:
                for date in condensed.loc[well].index:
                    analyte_df_resample.at[well, date] = condensed.loc[well,date].RESULT
            analyte_df_resample = analyte_df_resample.astype('float').T
            analyte_df_resample = analyte_df_resample.interpolate(method='linear')
            return analyte_df_resample

        all_dates = np.unique(data.COLLECTION_DATE)
        # Create array of equally spaced dates
        start = pd.Timestamp(all_dates.min())
        end = pd.Timestamp(all_dates.max())
        delta = end - start
        t = np.linspace(start.value, end.value, delta.days)
        t = pd.to_datetime(t)
        t = pd.Series(t)
        t = t.apply(lambda x: x.replace(minute=0, hour=0, second=0, microsecond=0, nanosecond=0))

        cutoff_dates = []
        # Save each analyte data
        analyte_data = []
        for analyte in analytes:
            ana_data = transform_time_series_by_analyte(data, analyte)
            if(rm_outliers):
                col_num = ana_data.shape[1]
                for col in range(col_num):
                    
                    # original_column = ana_data.iloc[:,col]
                    filtered_column = self.remove_outliers(ana_data.iloc[:,col].dropna(), z_threshold=z_threshold)
                    
                    # assign back only the filtered values (keeping original NaNs)
                    ana_data.loc[filtered_column.index, ana_data.columns[col]] = filtered_column

                # ana_data = ana_data.dropna()
                ana_data = ana_data.interpolate(method='linear')
            ana_data.index = pd.to_datetime(ana_data.index)
            # Resample
            ana_data_resample = ana_data.resample(resample).mean()
            # Save data
            analyte_data.append(ana_data_resample)
            # Determine cuttoff point for number of NaNs in dataset
            passes_limit = []
            for date in ana_data_resample.index:
                limit = 0.7 * ana_data_resample.shape[1]
                curr = ana_data_resample.isna().loc[date,:].value_counts()
                if('False' in str(curr)):
                    curr_total = ana_data_resample.isna().loc[date,:].value_counts()[0]
                    if curr_total > limit:
                        passes_limit.append(date)
            passes_limit = pd.to_datetime(passes_limit)
            cutoff_dates.append(passes_limit.min())
        start_index = pd.Series(cutoff_dates).max()

        # Get list of shared wells amongst all the listed analytes
        combined_well_list = []
        for x in range(len(analytes)):
            combined_well_list = combined_well_list + list(analyte_data[x].columns)
        combined_count = pd.Series(combined_well_list).value_counts()
        shared_wells = list(combined_count[list(pd.Series(combined_well_list).value_counts()==len(analytes))].index)

        # Vectorize data
        vectorized_df = pd.DataFrame(columns=analytes, index = shared_wells)

        for analyte, num in zip(analytes, range(len(analytes))):
            for well in shared_wells:
                analyte_data_full = analyte_data[num][well].fillna(analyte_data[num][well].mean())
                vectorized_df.at[well, analyte] = analyte_data_full[start_index:].values

        dates = ana_data_resample[start_index:].index
        return vectorized_df, dates

    def __get_individual_analyte_df(self, data, dates, analyte):
        sample = data[analyte]
        sample_analyte = pd.DataFrame(sample, index=dates, columns=sample.index)
        for well in sample.index:
            sample_analyte[well] = sample[well]
        return sample_analyte

    def __cluster_data_OLD(self, data, n_clusters=4, log_transform=False, filter=False, filter_well_by=['D'], return_clusters=False):
        if(filter):
            res_wells = self.filter_wells(filter_well_by)
            data = data.T
            data = data.loc[data.index.isin(res_wells)]
            data = data.T
        if(log_transform):
            data = np.log10(data)
            data = data.dropna(axis=1)
        temp = data.T
        k_Means = KMeans(n_clusters=n_clusters, random_state=42)
        km = k_Means.fit(temp)
        predict = km.predict(temp)
        temp['predicted'] = km.labels_
        colors = ['red', 'blue', 'orange', 'purple', 'green', 'beige', 'pink', 'black', 'cadetblue', 'lightgreen']
        temp['color'] = temp['predicted'].map(lambda p: colors[p])

        fig, ax = plt.subplots(figsize=(20,10))
        ax = plt.axes()

        color = temp['color']
        for x in range(temp.shape[0]):
            curr = data.iloc[:,x]
            ax.plot(curr, label=curr.name, color=color[x])
            ax.legend()
        
        if(return_clusters):
            color_df = pd.DataFrame(temp['color'])
            color_df['STATION_ID'] = color_df.index
            if(self.get_Construction_Data==None):
                print('You need to set the GPS data first using the setConstructionData function.')
                return None
            else:
                gps_color = pd.merge(self.get_Construction_Data(), color_df, on=['STATION_ID'])
                return gps_color

    def cluster_data(self, data, analyte_name=["ANALYTE_NAME"], n_clusters=4, filter=False, col=None, equals=[], year_interval=5, y_label = 'Concentration', return_clusters=False ):
        """Clusters time series concentration data using kmeans algorithm and plots it.

        Args:
            data (pd.DataFrame): data to be used in clustering.
            analyte_name (list, optional): analytes to use to cluster. Defaults to ["ANALYTE_NAME"].
            n_clusters (int, optional): number of clusters for kmeans. Defaults to 4.
            filter (bool, optional): flag to indicate filtering. Defaults to False.
            col (str, optional): column to filter. Example: col='STATION_ID'. Defaults to None.
            equals (list, optional): values to filter col by. Examples: equals=['FAI001A', 'FAI001B']. Defaults to [].
            year_interval (int, optional): plot x_label interval in years. Defaults to 5.
            y_label (str, optional): y axis label. Defaults to 'Concentration'.
            return_clusters (bool, optional): flag to return cluster assignemnt. Defaults to False.
        """
        data = data.copy()
        if(filter):
            filter_res = self.filter_by_column(data=self.get_Construction_Data(), col=col, equals=equals)
            if('ERROR:' in str(filter_res)):
                return filter_res
            query_wells = list(data.columns)
            filter_wells = list(filter_res.index.unique())
            intersect_wells = list(set(query_wells) & set(filter_wells))
            print(intersect_wells)
            if(len(intersect_wells)<=0):
                return 'ERROR: No results for this query with the specifed filter parameters.'
            data = data[intersect_wells]
        data.index = date2num(data.index)
        temp = data.T
        k_Means = KMeans(n_clusters=n_clusters, random_state=43)
        km = k_Means.fit(temp)
        predict = km.predict(temp)
        temp['predicted'] = km.labels_
        colors = ['red', 'blue', 'orange', 'purple', 'green', 'pink', 'black', 'cadetblue', 'lightgreen','beige']
        temp['color'] = temp['predicted'].map(lambda p: colors[p])
        
        fig, ax = plt.subplots(figsize=(10,10), dpi=100)
        ax.minorticks_off()
        ax = plt.axes()

        color = temp['color']
        for x in range(temp.shape[0]):
            curr = data.iloc[:,x]
            ax.plot(curr, label=curr.name, color=color[x])
            
        years = mdates.YearLocator(year_interval)  # every year
        months = mdates.MonthLocator()  # every month
        yearsFmt = mdates.DateFormatter('%Y') 
        ax = plt.gca()
        ax.xaxis.set_major_locator(years)
        ax.xaxis.set_major_locator(years)
        ax.xaxis.set_major_formatter(yearsFmt)
        ax.autoscale_view()
        ax.set_xlabel("Years", fontsize=20)
        plt.xticks(fontsize=20)
        ax.set_ylabel(y_label, fontsize=20)
        plt.yticks(fontsize=20)
        ax.set_title("{}: {} clusters".format(analyte_name, n_clusters), fontsize=20)
        if(return_clusters):
            color_df = pd.DataFrame(temp['color'])
            color_df['STATION_ID'] = color_df.index
            if(self.get_Construction_Data==None):
                print('You need to set the GPS data first using the setConstructionData function.')
                return None
            else:
                gps_color = pd.merge(self.get_Construction_Data(), color_df, on=['STATION_ID'])
                return gps_color
        # if(return_data):
        #     return color, temp

    def plot_all_time_series_simple(self, analyte_name=None, start_date=None, end_date=None, title='Dataset: Time ranges', x_label='Well', y_label='Year',
                                min_days=10, x_min_lim=-5, x_max_lim = 170, y_min_date='1988-01-01', y_max_date='2020-01-01', return_data=False, filter=False, col=None, equals=[]):
        """Plots the start and end date of analyte readings for differnt locations/sensors/wells.

        Args:
            analyte_name (str, optional): analyte to examine. Defaults to None.
            start_date (str, optional): start date of horizontal time to show alignment. Defaults to None.
            end_date (str, optional): end date of horizontal time to show alignment.. Defaults to None.
            title (str, optional): plot title. Defaults to 'Dataset: Time ranges'.
            x_label (str, optional): x axis label. Defaults to 'Well'.
            y_label (str, optional): y axis label. Defaults to 'Year'.
            min_days (int, optional): minimum number of days required to plot the time series . Defaults to 10.
            x_min_lim (int, optional): x axis starting point. Defaults to -5.
            x_max_lim (int, optional): x axis ending point. Defaults to 170.
            y_min_date (str, optional): y axis starting date. Defaults to '1988-01-01'.
            y_max_date (str, optional): y axis ending date. Defaults to '2020-01-01'.
            return_data (bool, optional): flag to return data. Defaults to False.
            filter (bool, optional): flag to indicate filtering. Defaults to False.
            col (str, optional): column to filter. Example: col='STATION_ID'. Defaults to None.
            equals (list, optional): values to filter col by. Examples: equals=['FAI001A', 'FAI001B']. Defaults to [].
        """
        data = self.simplify_data()
        if(filter):
            filter_res = self.filter_by_column(data=self.construction_data, col=col, equals=equals)
            if('ERROR:' in str(filter_res)):
                return filter_res
            query_wells = list(data.STATION_ID.unique())
            filter_wells = list(filter_res.index.unique())
            intersect_wells = list(set(query_wells) & set(filter_wells))
            if(len(intersect_wells)<=0):
                return 'ERROR: No results for this query with the specifed filter parameters.'
            data = data[data['STATION_ID'].isin(intersect_wells)]

        if(analyte_name!=None):
            data = data[data.ANALYTE_NAME == analyte_name]
        wells = data.STATION_ID.unique()
        wells_dateRange=pd.DataFrame(columns=['STATION_ID','START_DATE','END_DATE'])
        for i in range(len(wells)):
            wellName=wells[i]
            wellNamedData=data[data['STATION_ID']==wells[i]]
            minDate=min(wellNamedData['COLLECTION_DATE'])
            maxDate=max(wellNamedData['COLLECTION_DATE'])
            wells_dateRange.loc[wells_dateRange.shape[0]]=[wellName,minDate,maxDate]

        wells_dateRange["RANGE"] = wells_dateRange.END_DATE - wells_dateRange.START_DATE
        # wells_dateRange.RANGE = wells_dateRange.RANGE.astype('timedelta64[D]').astype('int')
        wells_dateRange = wells_dateRange[wells_dateRange.RANGE.dt.days>min_days]
        wells_dateRange.sort_values(by=["RANGE","END_DATE","START_DATE"], ascending = (False, False, True), inplace=True)
        wells_dateRange.reset_index(inplace=True)
        wells_dateRange.drop('index', axis=1, inplace=True)
        wells = np.array(wells_dateRange.STATION_ID)

        fig, ax = plt.subplots(1, 1, sharex=False,figsize=(20,6),dpi=300)

        ax.set_xticks(range(len(wells)))
        ax.set_xticklabels(wells, rotation='vertical', fontsize=6)

        ax.plot(wells_dateRange['START_DATE'], c='blue', marker='o',lw=0, label='Start date')
        ax.plot(wells_dateRange['END_DATE'], c='red', marker='o',lw=0, label='End date')

        ax.hlines([max(wells_dateRange['END_DATE'])], x_min_lim, x_max_lim, colors='purple', label='Selected end date')
        if(start_date==None):
            ax.hlines([min(wells_dateRange['START_DATE'])], x_min_lim, x_max_lim, colors='green', label='Selected start date')
        else:
            ax.hlines([pd.to_datetime(start_date)], x_min_lim, x_max_lim, colors='green', label='Selected start date')

        x_label = x_label + ' (count: ' + str(wells_dateRange.shape[0])+ ')'
        ax.set_xlabel(x_label, fontsize=20)
        ax.set_ylabel(y_label, fontsize=20)   
        ax.set_xlim([x_min_lim, x_max_lim])
        ax.set_ylim([pd.to_datetime(y_min_date), pd.to_datetime(y_max_date)]) 
        ax.plot([], [], ' ', label="Time series with at least {} days".format(min_days))
        ax.legend()
        if(analyte_name!=None):
            title = title + ' (' + analyte_name + ')'
        fig.suptitle(title, fontsize=20)
        for i in range(wells_dateRange.shape[0]):
            ax.vlines(i,wells_dateRange.loc[i,'START_DATE'],wells_dateRange.loc[i,'END_DATE'],colors='k')
        if(return_data):
            return wells_dateRange

    def plot_all_time_series(self, analyte_name=None, title='Dataset: Time ranges', x_label='Well', y_label='Year', x_label_size=8, marker_size=30,
                            min_days=10, x_min_lim=None, x_max_lim=None, y_min_date=None, y_max_date=None, sort_by_distance=True, source_coordinate=[436642.70,3681927.09], log_transform=False, cmap=mpl.cm.rainbow, 
                            drop_cols=[], return_data=False, filter=False, col=None, equals=[], cbar_min=None, cbar_max=None, reverse_y_axis=False, fontsize = 20, figsize=(20,6), dpi=300, y_2nd_label=None):
        """Plots the start and end date of analyte readings for differnt locations/sensors/wells with colored concentration reading.

        Args:
            analyte_name (str, optional): analyte to examine. Defaults to None.
            title (str, optional): plot title. Defaults to 'Dataset: Time ranges'.
            x_label (str, optional): x axis label. Defaults to 'Well'.
            y_label (str, optional): y axis label. Defaults to 'Year'.
            x_label_size (int, optional): x axis label font size. Defaults to 8.
            marker_size (int, optional): point size for time series. Defaults to 30.
            min_days (int, optional): minimum number of days required to plot the time series . Defaults to 10.
            x_min_lim (int, optional): x axis starting point. Defaults to None.
            x_max_lim (int, optional): x axis ending point. Defaults to None.
            y_min_date (str, optional): y axis starting date. Defaults to None.
            y_max_date (str, optional): y axis ending date. Defaults to None.
            sort_by_distance (bool, optional): flag to sort by distance from source center. Defaults to True.
            source_coordinate (list, optional): Easting, Northing coordinate of source center. Defaults to [436642.70,3681927.09].
            log_transform (bool, optional): flag to toggle log base 10 transformation. Defaults to False.
            cmap (cmap, optional): color map for plotting. Defaults to mpl.cm.rainbow.
            drop_cols (list, optional): columns, usually wells, to exclude. Defaults to [].
            return_data (bool, optional): flag to return data. Defaults to False.
            filter (bool, optional): flag to indicate filtering. Defaults to False.
            col (str, optional): column to filter. Example: col='STATION_ID'. Defaults to None.
            equals (list, optional): values to filter col by. Examples: equals=['FAI001A', 'FAI001B']. Defaults to [].
            cbar_min (float, optional): color bar lower boundary. Defaults to None.
            cbar_max (float, optional): color bar upper boundary. Defaults to None.
            reverse_y_axis (bool, optional): flag that reverses y axis. Defaults to False.
            fontsize (int, optional): plot font size. Defaults to 20.
            figsize (tuple, optional): matplotlib style figure size. Defaults to (20,6).
            dpi (int, optional): DPI of figure. Defaults to 300.
            y_2nd_label (str, optional): color bar label manual override. Defaults to None.
        """
        dt = self.getCleanData([analyte_name])
        dt = dt[analyte_name]
        try:
            if(filter):
                filter_res = self.filter_by_column(data=self.get_Construction_Data(), col=col, equals=equals)
                if('ERROR:' in str(filter_res)):
                    return filter_res
                query_wells = list(dt.columns.unique())
                filter_wells = list(filter_res.index.unique())
                intersect_wells = list(set(query_wells) & set(filter_wells) & set(dt.columns))
                if(len(intersect_wells)<=0):
                    return 'ERROR: No results for this query with the specifed filter parameters.'
                dt = dt[intersect_wells]
        
            well_info = self.get_Construction_Data()
            shared_wells = list(set(well_info.index) & set(dt.columns))
            dt = dt[shared_wells]
            well_info = well_info.T[shared_wells]
            dt = dt.reindex(sorted(dt.columns), axis=1)
            well_info = well_info.reindex(sorted(well_info.columns), axis=1)
            well_info = well_info.T
            transformer = Transformer.from_crs("epsg:4326", "epsg:26917") # Latitude/Longitude to UTM
            UTM_x, UTM_y = transformer.transform(well_info.LATITUDE, well_info.LONGITUDE)
            X = np.vstack((UTM_x,UTM_y)).T
            well_info = pd.DataFrame(X, index=list(well_info.index),columns=['Easting', 'Northing'])
            well_info = self.add_dist_to_source(well_info, source_coordinate=source_coordinate)
            if(sort_by_distance):
                well_info.sort_values(by=['dist_to_source'], ascending = True, inplace=True)
            dt = dt[well_info.index]
        except:
            pass

        dt = dt.interpolate()
        dt = dt.drop(drop_cols, axis=1) # DROP BAD ONES 
        
        if(log_transform):
            dt[dt <= 0] = 0.00000001
            dt = np.log10(dt)
        wells = dt.columns
        if(cbar_min==None):
            cbar_min = dt.min().min()
        if(cbar_max==None):
            cbar_max = dt.max().max() 
        norm = mpl.colors.Normalize(vmin=cbar_min, vmax=cbar_max)

        fig, ax = plt.subplots(1, 2, sharex=False, figsize=figsize, dpi=dpi, gridspec_kw={'width_ratios': [40, 1]})
        ax[0].set_xticks(range(len(wells)))
        ax[0].set_xticklabels(wells, rotation='vertical', fontsize=x_label_size)

        for col in wells:
            curr_start = dt[col].first_valid_index()
            curr_end =  dt[col].last_valid_index()
            length = len(list(dt[col].loc[curr_start:curr_end].index))
            color_vals = list(dt[col].loc[curr_start:curr_end])
            color_vals = [cmap(norm((x))) for x in color_vals]
            x = col
            ys = list(dt[col].loc[curr_start:curr_end].index)
            ax[0].scatter([str(x)]*length, ys, c=color_vals, marker='o',lw=0,s=marker_size, alpha=0.75)

        x_label = x_label + ' (count: ' + str(dt.shape[1])+ ')'
        ax[0].set_xlabel(x_label, fontsize=fontsize)
        ax[0].set_ylabel(y_label, fontsize=fontsize)
        if(x_min_lim==None):
            x_min_lim = -5
        if(x_max_lim==None):
            x_max_lim = len(wells)+5
        ax[0].set_xlim([x_min_lim, x_max_lim])
        if(y_min_date==None):
            y_min_date = dt.index.min() + relativedelta(years=-1)
        if(y_max_date==None):
            y_max_date = dt.index.max() + relativedelta(years=1)
        ax[0].set_ylim([pd.to_datetime(y_min_date), pd.to_datetime(y_max_date)]) 
        ax[0].plot([], [], ' ', label="Time series with at least {} days".format(min_days))
        ax[0].set_facecolor((0, 0, 0,0.1))
                
        # COLORBAR
        label_cb = "Concentration ({})".format(self.get_unit(analyte_name))
        if(log_transform):
            label_cb = "Log " + label_cb
        cbar = fig.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap=cmap),
                cax=ax[1], orientation='vertical')
        if(y_2nd_label!=None):
            label_cb = y_2nd_label
        cbar.set_label(label=label_cb,size=fontsize)

        ax[0].tick_params(axis='y', labelsize=fontsize)
        cbar.ax.tick_params(labelsize=fontsize)
        plt.tight_layout()
        if(reverse_y_axis):
            ax[0].invert_yaxis()
        if(analyte_name!=None):
            title = title + ' (' + analyte_name + ')'
        fig.suptitle(title, fontsize=fontsize, y=1.05)
        if(return_data):
            return dt



    # Helper function to return start and end date for a date and a lag (+/- days)
    def __getLagDate(self, date, lagDays=7):
        date = pd.to_datetime(date)
        dateStart = date - pd.DateOffset(days=lagDays)
        dateEnd = date + pd.DateOffset(days=lagDays)
        return dateStart, dateEnd


    def getCleanData(self, analytes):
        """Creates a table filling the data from the concentration dataset for a given analyte list where the columns are multi-indexed as follows [analytes, well names] and the index is all of the dates in the dataset. Many NaN should be expected.

        Args:
            analytes (list): list of analyte names to use

        Returns:
            pd.DataFrame
        """
        curr = self.data[['STATION_ID', 'COLLECTION_DATE', 'ANALYTE_NAME', 'RESULT']]
        main = pd.DataFrame()
        for ana in analytes:
            main = pd.concat([main, curr[curr.ANALYTE_NAME==ana]])
        piv = main.pivot_table(index=['COLLECTION_DATE'],columns=['ANALYTE_NAME', 'STATION_ID'], values='RESULT', aggfunc=np.mean)
        piv.index = pd.to_datetime(piv.index)
        piv.sort_index(inplace=True)
        return piv

    def getCommonDates(self, analytes, lag=[3,7,10]):
        """Creates a table which counts the number of wells within a range specified by a list of lag days.

        Args:
            analytes (list): list of analyte names to use
            lag (list, optional): list of days to look ahead and behind the specified date (+/-). Defaults to [3,7,10].

        Returns:
            pd.DataFrame
        """
        piv = self.getCleanData(analytes)
        dates = piv.index
        names=['Dates', 'Lag']
        tuples = [dates, lag]
        finalData = pd.DataFrame(index=pd.MultiIndex.from_product(tuples, names=names), columns=['Date Ranges', 'Number of wells'])
        for date in dates:
            for i in lag:
                dateStart, dateEnd = self.__getLagDate(date, lagDays=i)
                mask = (piv.index > dateStart) & (piv.index <= dateEnd)
                result = piv[mask].dropna(axis=1, how='all')
                numWells = len(list(result.columns.get_level_values(1).unique()))
                dateRange = str(dateStart.date()) + " - " + str(dateEnd.date())
                finalData.loc[date, i]['Date Ranges'] = dateRange
                finalData.loc[date, i]['Number of wells'] = numWells
        return finalData

    def getJointData(self, analytes, lag=3):
        """Creates a table filling the data from the concentration dataset for a given analyte list where the columns are multi-indexed as follows [analytes, well names] and the index is the date ranges secified by the lag.

        Args:
            analytes (list): list of analyte names to use
            lag (int, optional): number of days to look ahead and behind the specified date (+/-). Defaults to 3.

        Returns:
            pd.DataFrame
        """
        if(self.jointData_is_set(lag=lag)==True):
            finalData = self.__jointData[0]
            return finalData
        piv = self.getCleanData(analytes)
        dates = piv.index
        dateRanges = []
        for date in dates:
            dateStart, dateEnd = self.__getLagDate(date, lagDays=lag)
            dateRange = str(dateStart.date()) + " - " + str(dateEnd.date())
            dateRanges.append(dateRange)
        finalData = pd.DataFrame(columns=piv.columns, index=dateRanges)
        numLoops = len(dates)
        everySomePercent = []
        print("Generating data with a lag of {}.".format(lag).upper())
        print("Progress:")
        for x in list(np.arange(1, 100, 1)):
            everySomePercent.append(round((x/100)*numLoops))
        for date, iteration in zip(dates, range(numLoops)):
            if(iteration in everySomePercent):
                print(str(round(iteration/numLoops*100)) + "%", end=', ')
            dateStart, dateEnd = self.__getLagDate(date, lagDays=lag)
            dateRange = str(dateStart.date()) + " - " + str(dateEnd.date())
            mask = (piv.index > dateStart) & (piv.index <= dateEnd)
            result = piv[mask].dropna(axis=1, how='all')
            resultCollapse = pd.concat([result[col].dropna().reset_index(drop=True) for col in result], axis=1)
            # HANDLE MULTIPLE VALUES
            if(resultCollapse.shape[0]>1):
                resultCollapse = pd.DataFrame(resultCollapse.mean()).T
            resultCollapse = resultCollapse.rename(index={0: dateRange})
            for ana_well in resultCollapse.columns:
                finalData.loc[dateRange, ana_well] =  resultCollapse.loc[dateRange, ana_well]
            # Save data to the pylenm global variable
            self.__set_jointData(data=finalData, lag=lag)
        for col in finalData.columns:
            finalData[col] = finalData[col].astype('float64')
        print("Completed")
        return finalData

    def mse(self, y_true, y_pred):
        """Error Metric: Mean Squared Error 

        Args:
            y_true (numpy.array): true values
            y_pred (numpy.array): predicted values

        Returns:
            float: mean squared error
        """
        return mean_squared_error(y_true, y_pred)

    def get_Best_GP(self, X, y, smooth=True, seed = 42):
        """Returns the best Gaussian Process model for a given X and y.

        Args:
            X (numpy.array): array of dimension (number of wells, 2) where each element is a pair of UTM coordinates.
            y (numpy.array): array of size (number of wells) where each value corresponds to a concentration value at a well.
            smooth (bool, optional): flag to toggle WhiteKernel on and off. Defaults to True.
            seed (int, optional): random state setting. Defaults to 42.

        Returns:
            GaussianProcessRegressor: best GP model
        """
        gp = GaussianProcessRegressor(normalize_y=True, random_state=seed)
        # Kernel models
        if(smooth):
            k1 = Matern(length_scale=400, nu=1.5, length_scale_bounds=(100.0, 5000.0))+WhiteKernel()
            k2 = Matern(length_scale=800, nu=1.5, length_scale_bounds=(100.0, 5000.0))+WhiteKernel()
            k3 = Matern(length_scale=1200, nu=1.5, length_scale_bounds=(100.0, 5000.0))+WhiteKernel()
            k4 = Matern(length_scale=400, nu=1.5, length_scale_bounds=(100.0, 5000.0))+WhiteKernel()
            k5 = Matern(length_scale=800, nu=1.5, length_scale_bounds=(100.0, 5000.0))+WhiteKernel()
            k6 = Matern(length_scale=1200, nu=2.5, length_scale_bounds=(100.0, 5000.0))+WhiteKernel()
            k7 = Matern(length_scale=400, nu=2.5, length_scale_bounds=(100.0, 5000.0))+WhiteKernel()
            k8 = Matern(length_scale=800, nu=2.5, length_scale_bounds=(100.0, 5000.0))+WhiteKernel()
            k9 = Matern(length_scale=1200, nu=2.5, length_scale_bounds=(100.0, 5000.0))+WhiteKernel()
            k10 = Matern(length_scale=400, nu=np.inf, length_scale_bounds=(100.0, 5000.0))+WhiteKernel()
            k11 = Matern(length_scale=800, nu=np.inf, length_scale_bounds=(100.0, 5000.0))+WhiteKernel()
            k12 = Matern(length_scale=1200, nu=np.inf, length_scale_bounds=(100.0, 5000.0))+WhiteKernel()
        else:
            k1 = Matern(length_scale=400, nu=1.5, length_scale_bounds=(100.0, 5000.0))
            k2 = Matern(length_scale=800, nu=1.5, length_scale_bounds=(100.0, 5000.0))
            k3 = Matern(length_scale=1200, nu=1.5, length_scale_bounds=(100.0, 5000.0))
            k4 = Matern(length_scale=400, nu=1.5, length_scale_bounds=(100.0, 5000.0))
            k5 = Matern(length_scale=800, nu=1.5, length_scale_bounds=(100.0, 5000.0))
            k6 = Matern(length_scale=1200, nu=2.5, length_scale_bounds=(100.0, 5000.0))
            k7 = Matern(length_scale=400, nu=2.5, length_scale_bounds=(100.0, 5000.0))
            k8 = Matern(length_scale=800, nu=2.5, length_scale_bounds=(100.0, 5000.0))
            k9 = Matern(length_scale=1200, nu=2.5, length_scale_bounds=(100.0, 5000.0))
            k10 = Matern(length_scale=400, nu=np.inf, length_scale_bounds=(100.0, 5000.0))
            k11 = Matern(length_scale=800, nu=np.inf, length_scale_bounds=(100.0, 5000.0))
            k12 = Matern(length_scale=1200, nu=np.inf, length_scale_bounds=(100.0, 5000.0))
        parameters = {'kernel': [k1,k2,k3,k4,k5,k6,k7,k8,k9,k10,k11,k12]}
        model = GridSearchCV(gp, parameters, scoring='r2')
        model.fit(X, y)
        return model

    def fit_gp(self, X, y, xx, ft=['Easting','Northing'], model=None, smooth=True):
        """Fits Gaussian Process for X and y and returns both the GP model and the predicted values

        Args:
            X (numpy.array): array of dimension (number of wells, 2) where each element is a pair of UTM coordinates.
            y (numpy.array): array of size (number of wells) where each value corresponds to a concentration value at a well.
            xx (numpy.array): prediction locations
            ft (list, optional): feature names to train on. Defaults to ['Easting','Northing'].
            model (GaussianProcessRegressor, optional): model to fit. Defaults to None.
            smooth (bool, optional): flag to toggle WhiteKernel on and off. Defaults to True.

        Returns:
            GaussianProcessRegressor, numpy.array: GP model, prediction of xx
        """
        if(model==None):
            gp = self.get_Best_GP(X[ft], y, smooth) # selects best kernel params to fit
        else:
            gp = model
        gp.fit(X[ft], y)
        y_pred = gp.predict(xx[ft])
        return gp, y_pred

    def interpolate_topo(self, X, y, xx, ft=['Elevation'], gp_kernel=None, smooth=True, regression='linear', seed = 42):
        """Spatially interpolate the water table as a function of topographic metrics using Gaussian Process. Uses regression to generate trendline adds the values to the GP map.

        Args:
            X (numpy.array): training values. Must include "Easting" and "Northing" columns.
            y (numpy.array): array of size (number of wells) where each value corresponds to a concentration value at a well.
            xx (numpy.array): prediction locations
            ft (list, optional): eature names to train on. Defaults to ['Elevation'].
            model (GaussianProcessRegressor, optional): model to fit. Defaults to None.
            smooth (bool, optional): flag to toggle WhiteKernel on and off. Defaults to True.
            regression (str, optional): choice between 'linear' for linear regression, 'rf' for random forest regression, 'ridge' for ridge regression, or 'lasso' for lasso regression.. Defaults to 'linear'.
            seed (int, optional): random state setting. Defaults to 42.

        Returns:
            numpy.array: predicton of locations xx
        """
        alpha_Values = [1e-5, 5e-5, 0.0001, 0.0005, 0.005, 0.05, 0.1, 0.3, 1, 3, 5, 10, 15, 30, 80]
        if(regression.lower()=='linear'):
            reg = LinearRegression()
        if(regression.lower()=='rf'):
            reg = RandomForestRegressor(n_estimators=20, random_state=seed)
        if(regression.lower()=='ridge'):
            # reg = make_pipeline(PolynomialFeatures(3), Ridge())
            reg = RidgeCV(alphas=alpha_Values)
        if(regression.lower()=='lasso'):
            # reg = make_pipeline(PolynomialFeatures(3), Lasso())
            reg = LassoCV(alphas=alpha_Values)
        
        
        reg.fit(X[ft], y)
        y_est = reg.predict(X[ft])
        residuals = y - y_est
        if(gp_kernel==None):
            model = self.get_Best_GP(X[['Easting','Northing']], residuals, smooth=smooth, seed=seed)
        else:
            model = GaussianProcessRegressor(kernel=gp_kernel, optimizer=None, normalize_y=True)
            model.fit(X[['Easting','Northing']], residuals)
        reg_trend = reg.predict(xx[ft])

        
        r_map = model.predict(xx[['Easting','Northing']])
        y_map = reg_trend + r_map
        
        return y_map, r_map, residuals, reg_trend, model

    # Helper fucntion for get_Best_Wells
    def __get_Best_Well(self, X, y, xx, ref, selected, leftover, ft=['Elevation'], regression='linear', verbose=True, smooth=True, model=None):
        num_selected=len(selected)
        errors = []
        if(model==None):
            if(len(selected)<5):
                model, pred = self.fit_gp(X, y, xx)
            else:
                model = None
        else:
            model=model
        if(verbose):  
            print("# of wells to choose from: ", len(leftover))
        if(num_selected==0):
            if(verbose): 
                print("Selecting first well")
            for ix in leftover:
                y_pred, r_map, residuals, lr_trend = self.interpolate_topo(X=X.iloc[ix:ix+1,:], y=y[ix:ix+1], xx=xx, ft=ft, regression=regression, model=model, smooth=smooth)
                y_err = self.mse(ref, y_pred)
                errors.append((ix, y_err))
        
        if(num_selected > 0):
            for ix in leftover:
                joined = selected + [ix]
                y_pred, r_map, residuals, lr_trend = self.interpolate_topo(X=X.iloc[joined,:], y=y[joined], xx=xx, ft=ft, regression=regression, model=model, smooth=smooth)
                y_err = self.mse(ref, y_pred)
                errors.append((ix, y_err))
            
        err_ix = [x[0] for x in errors]
        err_vals = [x[1] for x in errors]
        min_val = min(err_vals)
        min_ix = err_ix[err_vals.index(min(err_vals))]
        if(verbose):
            print("Selected well: {} with a MSE error of {}\n".format(min_ix, min_val))
        return min_ix, min_val

    def get_Best_Wells(self, X, y, xx, ref, initial, max_wells, ft=['Elevation'], regression='linear', verbose=True, smooth=True, model=None):
        """Greedy optimization function to select a subset of wells as to minimizes the MSE from a reference map

        Args:
            X (numpy.array): array of dimension (number of wells, 2) where each element is a pair of UTM coordinates.
            y (numpy.array): array of size (number of wells) where each value corresponds to a concentration value at a well.
            xx (numpy.array): prediction locations
            ref (numpy.array): reference field to optimize for (aka best/true map)
            initial (list): indices of wells as the starting wells for optimization
            max_wells (int): number of wells to optimize for
            ft (list, optional): feature names to train on. Defaults to ['Elevation'].
            regression (str, optional): choice between 'linear' for linear regression, 'rf' for random forest regression, 'ridge' for ridge regression, or 'lasso' for lasso regression.. Defaults to 'linear'.
            verbose (bool, optional): v. Defaults to True.
            smooth (bool, optional): flag to toggle WhiteKernel on and off. Defaults to True.
            model (GaussianProcessRegressor, optional): model to fit. Defaults to None.

        Returns:
            list: index of best wells in order from best to worst
        """
        tot_err = []
        selected = initial
        leftover = list(range(0, X.shape[0])) # all indexes from 0 to number of well
        
        # Remove the initial set of wells from pool of well indices to choose from
        for i in initial:
            leftover.remove(i)

        for i in range(max_wells-len(selected)):
            if(i==0): # select first well will min error
                well_ix, err = self.__get_Best_Well(X=X,y=y, xx=xx, ref=ref, selected=selected, leftover=leftover, ft=ft, regression=regression, verbose=verbose, smooth=smooth, model=model)
                selected.append(well_ix)
                leftover.remove(well_ix)
                tot_err.append(err)
            else:
                well_ix, err = self.__get_Best_Well(X=X,y=y, xx=xx, ref=ref, selected=selected, leftover=leftover, ft=ft, regression=regression, verbose=verbose, smooth=smooth, model=model)
                selected.append(well_ix)
                leftover.remove(well_ix)
                tot_err.append(err)
        print(selected)
        return selected, tot_err


    def dist(self, p1, p2):
        """2D Euclidean distance function

        Args:
            p1 (tuple): first point
            p2 (tuple): second point

        Returns:
            float: Euclidean distance
        """
        return sqrt(((p1[0]-p2[0])**2)+((p1[1]-p2[1])**2))

    def add_dist_to_source(self, XX, source_coordinate=[436642.70,3681927.09], col_name='dist_to_source'):
        """adds column to data with the distance of a record to the source coordinate

        Args:
            XX (pd.DataFrame): data with coordinate information
            source_coordinate (list, optional): source coordinate. Defaults to [436642.70,3681927.09].
            col_name (str, optional): name to assign new column. Defaults to 'dist_to_source'.

        Returns:
            pd.DataFrame: returns original data with additional column with the distance.
        """
        x1,y1 = source_coordinate
        distances = []
        for i in range(XX.shape[0]):
            x2,y2 = XX.iloc[i][0], XX.iloc[i][1]
            distances.append(self.dist([x1,y1],[x2,y2]))
        XX[col_name] = distances
        return XX
    

##################################################################################################
##################################################################################################
######### new functions 2024-08 #########

    ## by K. Whiteaker 2024-08, kwhit@alum.mit.edu
    def remove_outliers_lowess(self, data, lowess_frac=0.1, std_thresh=2.2, return_difference=False):
        """Identifies outliers in time series data as deviations of std_thresh standard deviations from a LOWESS fit, then sets these outliers to np.nan
    
        Args:
            data (pd.Series, or pd.DataFrame with 1 column): data for the outliers to removed from. MUST BE indexed by datetime (ex. collection date). Can convert pylenm concentration dataframe (for one station and one analyte) into a valid input via dataframe.set_index('COLLECTION_DATE').RESULT
            lowess_frac (float, optional): fraction of total data points considered in local fits of LOWESS smoother. A smaller value captures more local behaviour, and may be required for large datasets. Defaults to 0.1
            std_thresh (int, optional): number of standard deviations in (observation - fit) outside of which is considered an outlier. Defaults to 2.2
            return_difference (bool, optional): if True, return a pd.Series containing the difference between data and lowess fit
    
        Returns:
            if return_difference:
                pd.Series: input data with outliers set to np.nan
                pd.Series: input data - lowess fit
            else:
                pd.Series: input data with outliers set to np.nan
        """

        # create a copy so the data isn't modified inplace
        working_data = data.copy()
        
        # lowess() reads datetime as nanoseconds and has issues with large numbers & low frac, so need to scale down x-axis during fitting
        scaleDown = 1e17
        x_data = pd.to_datetime(working_data.index)
        x_readable = x_data.astype('int64').to_numpy()/scaleDown
        data_lowess = lowess(working_data, x_readable, frac=lowess_frac, return_sorted=False)

        # identify outlier locations and mark them as nan
        difference = working_data - data_lowess  # this is a pd.Series
        thresh = std_thresh*np.std(difference)
        difference_ignoreNaN = np.ma.array(difference, mask=np.isnan(difference)) # Use a mask to mark & ignore the NaNs as per https://stackoverflow.com/questions/37749900/how-to-disregard-the-nan-data-point-in-numpy-array-and-generate-the-normalized-d
        outliers_iloc = np.where(np.abs(difference_ignoreNaN)>thresh)[0]
        outliers = working_data.iloc[outliers_iloc]
        working_data.iloc[outliers_iloc] = np.nan

        if return_difference:
            return working_data, difference
        else:
            return working_data



    ## by K. Whiteaker 2024-08, kwhit@alum.mit.edu
    # Helper function for plot_data_weekAvg() and plot_data_lowess()
    # copied from self.plot_data(), with a few additional plot items added
    def __plot_data_xOutliers_fit(self, station_name, analyte_name, concentration_data_xOutliers, concentration_data_xOutliers_fit, outliers, units,
                                  difference, fitColor, fitName, rm_outliers, std_thresh, lowess_frac, show_difference, x_label, y_label, 
                                  year_interval, log_transform, y_zoom, return_data, save, save_dir, plot_inline, save_as_pdf):
        # define figure shape and size
        plt.figure(figsize=(8,8))
        ax = plt.axes()
        years = mdates.YearLocator(year_interval)  # every year
        months = mdates.MonthLocator()  # every month
        yearsFmt = mdates.DateFormatter('%Y')
        for label in ax.get_xticklabels():
            label.set_rotation(30)
            label.set_horizontalalignment('center')
        ax = plt.gca()
        ax.xaxis.set_major_locator(years)
        ax.xaxis.set_major_formatter(yearsFmt)
        ax.autoscale_view()
        if y_label is None:
            y_label = analyte_name + f' [{units}]'
            if(log_transform):
                y_label = 'log_10 of ' + analyte_name + f' [{units}]'
        ax.set_ylabel(y_label)
        ax.set_xlabel(x_label)
        small_fontSize = 15
        large_fontSize = 20
        plt.rc('axes', titlesize=large_fontSize)
        plt.rc('axes', labelsize=large_fontSize)
        plt.rc('legend', fontsize=small_fontSize)
        plt.rc('xtick', labelsize=small_fontSize)
        plt.rc('ytick', labelsize=small_fontSize)
        props = dict(boxstyle='round', facecolor='grey', alpha=0.15)
        ax.text(1.05, 0.85, 'Samples: {}'.format(concentration_data_xOutliers.size), transform=ax.transAxes, 
                fontsize=small_fontSize,
                fontweight='bold',
                verticalalignment='top', 
                bbox=props)
        ax.set_title(str(station_name) + ' - ' + analyte_name, fontweight='bold')
        ttl = ax.title
        ttl.set_position([.5, 1.05])

        # plot items
        ax.plot(concentration_data_xOutliers.index, concentration_data_xOutliers, ls='', marker='o', ms=2, color='black', alpha=1)
        if rm_outliers:
            ax.plot(outliers.index, outliers, ls='', marker='x', ms=5, color='red', alpha=0.75, label="Outliers ({})".format(outliers.size))
        # need to use pd.Series.dropna() when plotting so that matplotlib doesn't insert any gaps
        ax.plot(concentration_data_xOutliers_fit.dropna().index, concentration_data_xOutliers_fit.dropna(), ls='-', marker='', ms=2, lw=1, color=fitColor, alpha=1, label=fitName)
        if y_zoom:
            plt.ylim(min(concentration_data_xOutliers), max(concentration_data_xOutliers))
        ax.legend(bbox_to_anchor=(1.04, 1), loc='upper left', borderaxespad=0.)
        if save_as_pdf:
            fileType = '.pdf'
        else:
            fileType = '.png'
        if save:
            if not os.path.exists(save_dir):
                os.makedirs(save_dir)
            plt.savefig(save_dir + '/' + str(station_name) + '_' + analyte_name + '_' + fitName.replace(" ", "") + fileType, bbox_inches="tight")
        if(plot_inline):
            plt.show()
        plt.clf()
        plt.cla()
        plt.close()
        
        # plot difference between fit and observation, with outliers marked
        if rm_outliers and show_difference:
            outliers_loc = outliers.index
            plt.scatter(difference.index, difference, s=0.5, label='Difference')
            plt.scatter(difference[outliers_loc].index, difference[outliers_loc], s=10, marker='x', label=f'Outliers outside {std_thresh} SD')
            plt.title("(obs - fit), fraction "+ str(lowess_frac) + ", station "+ str(station_name))
            plt.ylabel('Observation - LOWESS Fit [ft]')
            plt.legend()
            if save:
                plt.savefig(save_dir + '/' + str(station_name) + '_' + analyte_name + '_difference' + fileType, bbox_inches="tight")
            if(plot_inline):
                plt.show()
            plt.clf()
            plt.cla()
            plt.close()



    ## by K. Whiteaker 2024-08, kwhit@alum.mit.edu
    def plot_data_rollAvg(self, station_name, analyte_name, window='1W', rm_outliers=True, std_thresh=2.2, lowess_frac=0.1, show_difference=False, x_label=None, 
                         y_label=None, year_interval=2, log_transform=False, y_zoom=False, return_data=False, save=False, save_dir='plot_data_lowess',
                         plot_inline=True, save_as_pdf=False):
        """Plot time series data for a specified station and analyte, alongside a curve calculated via a rolling time average (default 1 week) centred around each point
    
        Args:
            station_name (str): name of the station to be processed
            analyte_name (str): name of the analyte to be processed
            window (str): the time length of the rolling window used for averaging, ex. '1W', '2D'. Defaults to '1W'. Months are not supported; valid inputs are listed here: https://pandas.pydata.org/pandas-docs/stable/user_guide/timeseries.html
            rm_outliers (bool, optional): if True, use lowess function to remove outliers. Defaults to True
            std_thresh (int, optional): number of standard deviations in (observation - LOWESS fit) outside of which is considered an outlier. Defaults to 2.2
            lowess_frac (float, optional): fraction of total data points considered in local fits of LOWESS smoother. A smaller value captures more local behaviour, and may be required for large datasets. Defaults to 0.1
            show_difference (bool, optional): if True, plots (observation - LOWESS fit) at each measurement time
            x_label (str, optional): x axis label. Defaults to None
            y_label (str, optional): y axis label. Defaults to None. If left as None, will be set to analyte_name + f' [{units}]'
            year_interval (int, optional): plot by how many years to appear in the axis e.g.(1 = every year, 5 = every 5 years, ...). Defaults to 2
            log_transform (bool, optional): choose whether or not the data should be transformed to log base 10 values. Defaults to False
            y_zoom (bool, optional): if True, plot y axes will zoom in to [minimum y value, maximum y value] after removing outliers. Defaults to False
            return_data (bool, optional): if True, return the data used to make the plots. Defaults to False
            save (bool, optional): if True, save plots to file in save_dir. Defaults to False
            save_dir (str, optional): name of the directory you want to save the plot to. Defaults to 'plot_data_lowess'
            plot_inline (bool, optional): choose whether or not to show plot inline. Defaults to True
            save_as_pdf (bool, optional): if True, saves figure as vectorized pdf instead of png. Defaults to False
    
        Returns:
            if return_data:
                if rm_outliers:
                    pd.Series: concentration data with outliers set to np.nan, indexed by date
                    pd.Series: rolling average of post-outlier-removal concentration data, indexed by date
                    pd.Series: outliers, indexed by date
                else:
                    pd.Series: concentration data, indexed by date
                    pd.Series: rolling average of concentration data, indexed by date
            else:
                None
        """
    
        # import data and filter for the chosen analyte
        query = self.query_data(station_name, analyte_name)
        query = self.simplify_data(data=query)
        if(type(query)==int and query == 0):
            return 'No results found for {} and {}'.format(station_name, analyte_name)
        
        # reshape data so it can be passed into lowess() and remove_outliers_lowess()
        concentration_data = query.set_index('COLLECTION_DATE').RESULT
        time_data = concentration_data.index
        units = query.RESULT_UNITS.values[0]  # assumes same units for every entry
        if(log_transform):
            concentration_data = concentration_data.apply(np.log10)

        # if rm_outliers toggled, call remove_outliers_lowess() to set outliers to np.nan via lowess fit
        if rm_outliers:
            concentration_data_xOutliers, difference = self.remove_outliers_lowess(concentration_data, lowess_frac=lowess_frac,
                                                                                   std_thresh=std_thresh, return_difference=True)
        else:
            concentration_data_xOutliers = concentration_data
            difference = pd.Series  # will not be used, just need something to pass into __plot_data_xOutliers_fit
        outliers_loc = concentration_data_xOutliers.compare(concentration_data).index  # dates at which outliers occurred
        outliers = concentration_data[outliers_loc]  # outlier values indexed by their date of occurrence
        num_outliers = outliers.size
        analyte_max_xOutliers, analyte_min_xOutliers = max(concentration_data_xOutliers), min(concentration_data_xOutliers)
                
        # now that outliers have been dropped, calculate rolling average of y values (result of average is assigned to the centre of the window)
        avg_window = pd.Timedelta(window)  # could also do pd.Timedelta(7, "d")
        rolling_ser = concentration_data_xOutliers
        roll = rolling_ser.rolling(window=avg_window, min_periods=1, center=True)
        concentration_data_xOutliers_avg = roll.mean()
        
        # plot rolling average
        self.__plot_data_xOutliers_fit(station_name=station_name, analyte_name=analyte_name, concentration_data_xOutliers=concentration_data_xOutliers,
                                       concentration_data_xOutliers_fit=concentration_data_xOutliers_avg, outliers=outliers, units=units, difference=difference, 
                                       fitColor='green', fitName=window+' Average', rm_outliers=rm_outliers, std_thresh=std_thresh,
                                       lowess_frac=lowess_frac, show_difference=show_difference, x_label=x_label, y_label=y_label,
                                       year_interval=year_interval, log_transform=log_transform, y_zoom=y_zoom, return_data=return_data,
                                       save=save, save_dir=save_dir, plot_inline=plot_inline, save_as_pdf=save_as_pdf)
        
        # return data
        if return_data:
            if rm_outliers:
                return concentration_data_xOutliers, concentration_data_xOutliers_avg, outliers
            else:
                return concentration_data_xOutliers, concentration_data_xOutliers_avg
    


    ## by K. Whiteaker 2024-08, kwhit@alum.mit.edu
    def plot_data_lowess(self, station_name, analyte_name, rm_outliers=True, std_thresh=2.2, lowess_frac=0.1, show_difference=False, x_label=None, 
                         y_label=None, year_interval=2, log_transform=False, y_zoom=False, return_data=False, save=False, save_dir='plot_data_lowess',
                         plot_inline=True, save_as_pdf=False):
        """Plot time series data for a specified station and analyte, alongside a smoothed curve on interpolated data points (LOWESS)
    
        Args:
            station_name (str): name of the station to be processed
            analyte_name (str): name of the analyte to be processed
            rm_outliers (bool, optional): if True, use lowess function to remove outliers. Defaults to True
            std_thresh (int, optional): number of standard deviations in (observation - fit) outside of which is considered an outlier. Defaults to 2.2
            lowess_frac (float, optional): fraction of total data points considered in local fits of LOWESS smoother. A smaller value captures more local behaviour, and may be required for large datasets. Defaults to 0.1
            show_difference (bool, optional): if True, plots (observation - LOWESS fit) at each measurement time
            x_label (str, optional): x axis label. Defaults to None
            y_label (str, optional): y axis label. Defaults to None. If left as None, will be set to analyte_name + f' [{units}]'
            year_interval (int, optional): plot by how many years to appear in the axis e.g.(1 = every year, 5 = every 5 years, ...). Defaults to 2
            log_transform (bool, optional): choose whether or not the data should be transformed to log base 10 values. Defaults to False
            y_zoom (bool, optional): if True, plot y axes will zoom in to [minimum y value, maximum y value] after removing outliers. Defaults to False
            return_data (bool, optional): if True, return the data used to make the plots. Defaults to False
            save (bool, optional): if True, save plots to file in save_dir. Defaults to False
            save_dir (str, optional): name of the directory you want to save the plot to. Defaults to 'plot_data_lowess'
            plot_inline (bool, optional): choose whether or not to show plot inline. Defaults to True
            save_as_pdf (bool, optional): if True, saves figure as vectorized pdf instead of png. Defaults to False
    
        Returns:
            if return_data:
                if rm_outliers:
                    pd.Series: concentration data with outliers set to np.nan, indexed by date
                    pd.Series: LOWESS fit of post-outlier-removal concentration data, indexed by date
                    pd.Series: outliers, indexed by date
                else:
                    pd.Series: concentration data, indexed by date
                    pd.Series: LOWESS fit of concentration data, indexed by date
            else:
                None
        """
    
        # import data and filter for the chosen analyte
        query = self.query_data(station_name, analyte_name)
        query = self.simplify_data(data=query)
        if(type(query)==int and query == 0):
            return 'No results found for {} and {}'.format(station_name, analyte_name)
        
        # reshape data so it can be passed into lowess() and remove_outliers_lowess()
        concentration_data = query.set_index('COLLECTION_DATE').RESULT
        time_data = concentration_data.index
        units = query.RESULT_UNITS.values[0]  # assumes same units for every entry
        if(log_transform):
            concentration_data = concentration_data.apply(np.log10)

        # if rm_outliers toggled, call remove_outliers_lowess() to set outliers to np.nan via lowess fit
        if rm_outliers:
            concentration_data_xOutliers, difference = self.remove_outliers_lowess(concentration_data, lowess_frac=lowess_frac, std_thresh=std_thresh, return_difference=True)
        else:
            concentration_data_xOutliers = concentration_data
            difference = pd.Series  # will not be used, just need something to pass into __plot_data_xOutliers_fit
        outliers_loc = concentration_data_xOutliers.compare(concentration_data).index  # dates at which outliers occurred
        outliers = concentration_data[outliers_loc]  # outlier values indexed by their date of occurrence
        num_outliers = outliers.size
        analyte_max_xOutliers, analyte_min_xOutliers = max(concentration_data_xOutliers), min(concentration_data_xOutliers)
                
        # now that outliers have been dropped, calculate a new lowess fit for plotting
        # lowess() reads datetime as nanoseconds and has issues with large numbers & low frac, so need to scale down x-axis during fitting
        scaleDown = 1e17
        x_readable = time_data.to_numpy().astype(int)/scaleDown
        y_data_lowess = lowess(concentration_data_xOutliers, x_readable, frac=lowess_frac, return_sorted=False)
        concentration_data_xOutliers_lowess = pd.Series(data=y_data_lowess, index=concentration_data_xOutliers.index, dtype=float, name=concentration_data_xOutliers.name)

        # plot lowess fit
        self.__plot_data_xOutliers_fit(station_name=station_name, analyte_name=analyte_name, concentration_data_xOutliers=concentration_data_xOutliers,
                                       concentration_data_xOutliers_fit=concentration_data_xOutliers_lowess, outliers=outliers, units=units, difference=difference, 
                                       fitColor='violet', fitName='LOWESS Fit', rm_outliers=rm_outliers, std_thresh=std_thresh,
                                       lowess_frac=lowess_frac, show_difference=show_difference, x_label=x_label, y_label=y_label,
                                       year_interval=year_interval, log_transform=log_transform, y_zoom=y_zoom, return_data=return_data,
                                       save=save, save_dir=save_dir, plot_inline=plot_inline, save_as_pdf=save_as_pdf)
        
        if return_data:
            if rm_outliers:
                return concentration_data_xOutliers, concentration_data_xOutliers_lowess, outliers
            else:
                return concentration_data_xOutliers, concentration_data_xOutliers_lowess



    ## by K. Whiteaker, kwhit@alum.mit.edu
    def time_average_all_stations(self, analyte, period='1W', rm_outliers=True, std_thresh=2.2, lowess_frac=0.1):
        """Transforms all analyte data from exact measurements, at potentially different dates for each station, into periodic averaged data (ex. weekly averages) at the same dates for all stations.
    
        Args:
            analyte (_type_): analyte name
            period (str, optional): {‘D’, ‘W’, ‘M’, ‘Y’} time period over which to average. See https://pandas.pydata.org/pandas-docs/stable/user_guide/timeseries.html for valid frequency inputs. (e.g. ‘2W’ = every 2 weeks). Defaults to '1W'
            rm_outliers (bool, optional): flag to remove outliers in the data via LOWESS fit. Defaults to True
            std_thresh (int, optional): number of standard deviations in (observation - fit) outside of which is considered an outlier. Defaults to 2.2
            lowess_frac (float, optional): fraction of total data points considered in local fits of LOWESS smoother. A smaller value captures more local behaviour, and may be required for large datasets. Defaults to 0.1
    
        Returns:
            pd.DataFrame: dataframe with columns = stations, rows = dates, and data representing the average for that station between each date
        """
        # import data and filter for the chosen analyte
        data_concentration = self.data
        data_concentration = data_concentration[data_concentration.ANALYTE_NAME == analyte]
        data_construction = self.construction_data
    
        # create reshaped_df with columns=all available stations and rows=all measurement dates, to iterate over when conducting periodic averaging
        stations = np.unique(data_concentration[data_concentration.ANALYTE_NAME == analyte].STATION_ID)
        times = pd.to_datetime(np.unique(data_concentration.COLLECTION_DATE))
        reshaped_df = pd.DataFrame(index=times, columns=stations)
        
        # remove outliers from each station's data using remove_outlier_lowess(), and save the remaining data in reshaped_df
        for station in stations:
            data_in = data_concentration[data_concentration.STATION_ID == station].set_index('COLLECTION_DATE').RESULT
            data_in.index = pd.to_datetime(data_in.index)
            if rm_outliers:
                data_xOutliers = self.remove_outliers_lowess(data_in, lowess_frac, std_thresh)
            else:
                data_xOutliers = data_in
            reshaped_df[station].loc[data_xOutliers.index] = data_xOutliers
        
        # find first and last measurement dates
        date_start = pd.Timestamp(times.min())
        date_end = pd.Timestamp(times.max())
        # make an array of datetimes, 1 period (ex. 1 week) apart, that spanning from before the first measurement date to after the last measurement date
        period_timesteps = pd.date_range(min(times) - pd.Timedelta(period), max(times) + pd.Timedelta(period), freq=period).to_numpy()
    
        # conduct periodic averaging: take periodic (ex. 1 week) averages of each station and save the result to averaged_df
        averaged_df = pd.DataFrame(index=period_timesteps, columns=stations,dtype=float)
        for period_step in range(period_timesteps.size-1):  # for each period...
            # if period_step%100==0:
            #     print(period_step)  # print progress
            period_step_bin = []
            relevant_times = times[(times > period_timesteps[period_step]) & (times < period_timesteps[period_step+1])]
            for relevant_time in relevant_times:  # for each existing measurement that occurred within this period...
                timestep_data = reshaped_df.loc[relevant_time]  # identify all data points at this timestep
                period_step_bin.append(timestep_data)
            if len(period_step_bin)>0:  # if there are any measurements in this period
                period_data = pd.concat(period_step_bin, axis=1)
                period_station_avg = period_data.mean(axis=1, skipna=True)  # for each station, take the average over this period
                averaged_df.loc[period_timesteps[period_step]] = period_station_avg
        return averaged_df



    ## by K. Whiteaker, kwhit@alum.mit.edu
    def get_all_analyte_data(self, analyte, rm_outliers=True, std_thresh=2.2, lowess_frac=0.1):
        """Returns a Pandas DataFrame containing all measurement data for a given analyte, with columns as measurement stations and rows as measurement dates.
    
        Args:
            analyte (_type_): analyte name
            rm_outliers (bool, optional): flag to remove outliers in the data via LOWESS fit. Defaults to True
            std_thresh (int, optional): number of standard deviations in (observation - fit) outside of which is considered an outlier. Defaults to 2.2
            lowess_frac (float, optional): fraction of total data points considered in local fits of LOWESS smoother. A smaller value captures more local behaviour, and may be required for large datasets. Defaults to 0.1
    
        Returns:
            pd.DataFrame: dataframe of analyte masurement data, with columns = stations and rows = measurement dates
        """
        # import data and filter for the chosen analyte
        data_concentration = self.data
        data_concentration = data_concentration[data_concentration.ANALYTE_NAME == analyte]
        data_construction = self.construction_data
    
        # create reshaped_df with columns=all available stations and rows=all measurement dates, to fill through iteration
        stations = np.unique(data_concentration[data_concentration.ANALYTE_NAME == analyte].STATION_ID)
        times = pd.to_datetime(np.unique(data_concentration.COLLECTION_DATE))
        # times = pd.DatetimeIndex(np.unique(data_concentration.COLLECTION_DATE))
        reshaped_df = pd.DataFrame(index=times, columns=stations)
        
        # remove outliers from each station's data using remove_outlier_lowess(), and save the remaining data in reshaped_df
        for station in stations:
            data_in = data_concentration[data_concentration.STATION_ID == station].set_index('COLLECTION_DATE').RESULT
            data_in.index = pd.to_datetime(data_in.index)
            if rm_outliers:
                data_xOutliers = self.remove_outliers_lowess(data_in, lowess_frac, std_thresh)
            else:
                data_xOutliers = data_in
            reshaped_df[station].loc[data_xOutliers.index] = data_xOutliers

        return reshaped_df
    
##################################################################################################
##################################################################################################
######### new functions for sensor data 2025 #########

    def sensor_calibration_adjustment(self, sensor_data, sensor_name, calibration_date):
        """
        Adjusts sensor measurements prior to a calibration date to correct for drift or offset
        based on average values before and after the calibration.

        The function estimates drift by comparing mean sensor values within defined time windows:
        - 10 days shortly after the first data point (for baseline),
        - 10 days before the calibration date (for pre-calibration trend),
        - 10 days after the calibration date (for post-calibration trend).

        A linear correction is applied to all measurements prior to the calibration date.

        Parameters:
            sensor_data (pd.DataFrame): Time series data with columns:
                - 'COLLECTION_DATE' (datetime): timestamp of each measurement
                - 'RESULT' (float): measured sensor value
            sensor_name (str): Name of the sensor being calibrated
            calibration_date (datetime): Date when the sensor was calibrated

        Returns:
            pd.DataFrame: The adjusted sensor data with a corrected 'RESULT' column before calibration.
                        Measurements on or after the calibration date remain unchanged.
        """
        
        def mean_in_window(start, end):
            mask = (sensor_data['COLLECTION_DATE'] >= start) & (sensor_data['COLLECTION_DATE'] < end)
            return sensor_data.loc[mask, 'RESULT'].mean()

        # get the start date of the sensor data
        t0 = sensor_data['COLLECTION_DATE'].iloc[0]

        # define calibration (sampling) windows for calculating drifts
        calib1_start, calib1_end = t0 + datetime.timedelta(days=3), t0 + datetime.timedelta(days=13)
        calib2_start, calib2_end = calibration_date - datetime.timedelta(days=13), calibration_date - datetime.timedelta(days=3)
        calib3_start, calib3_end = calibration_date + datetime.timedelta(days=3), calibration_date + datetime.timedelta(days=13)

        # compute mean values for each window
        v1 = mean_in_window(calib1_start, calib1_end)
        v2 = mean_in_window(calib2_start, calib2_end)
        v3 = mean_in_window(calib3_start, calib3_end)

        # calculate the drift of value and time
        delta = v2 - v3
        period = (calib2_end - calib1_start).total_seconds()

        # split data before and after calibration
        pre_cal = sensor_data[sensor_data['COLLECTION_DATE'] < calibration_date].copy()
        post_cal = sensor_data[sensor_data['COLLECTION_DATE'] >= calibration_date].copy()

        if sensor_name == "Level: Depth to Water - Sensor":
            # apply constant correction to post-calibration data (pre-calibration is assumed more accurate)
            post_cal['RESULT'] += delta

        elif sensor_name == "pH - Sensor":
            # apply linear correction to pre-calibration data
            pre_cal['RESULT'] -= delta * ((pre_cal['COLLECTION_DATE'] - t0).dt.total_seconds() / period)

        else:
            # apply constant correction to post-calibration data (post-calibration is assumed more accurate)
            pre_cal['RESULT'] -= delta

        
        # combine and return
        return pd.concat([pre_cal, post_cal])
    


    def hampel_filter(self, ts, window_size=7, n_sigmas=3, replace=False):
        """
        Robust Hampel filter with adaptive window size near edges.

        Parameters:
            ts (pd.Series): Time series data
            window_size (int): Size of the rolling window (should be odd)
            n_sigmas (float): Threshold in terms of MAD
            replace (bool): If True, replaces outliers with median, else sets to NaN

        Returns:
            pd.Series: Filtered series
        """
        ts_filtered = ts.copy()
        k = 1.4826  # Constant for Gaussian distribution
        half_window = window_size // 2

        for i in range(len(ts)):
            # Dynamically adjust window bounds
            start = max(i - half_window, 0)
            end = min(i + half_window + 1, len(ts))

            window = ts.iloc[start:end]
            median = window.median()
            mad = k * np.median(np.abs(window - median))

            if mad == 0:
                continue  # skip division by zero risk

            if abs(ts.iloc[i] - median) > n_sigmas * mad:
                ts_filtered.iloc[i] = median if replace else np.nan

        return ts_filtered


    def plot_sample_with_outliers(self, ax, sample_data, remove_outliers=False, **kwargs):
        """
        Plot sample data, optionally highlighting outliers.
        """
        if sample_data.empty:
            return

        if remove_outliers:
            # Detect outliers using Hampel filter but keep both sets
            filtered = sample_data.copy()
            filtered['RESULT_filtered'] = self.hampel_filter(filtered['RESULT'], window_size=5, n_sigmas=4)
            mask_outlier = filtered['RESULT_filtered'].isna()  # assume Hampel marks outliers as NaN

            # Plot normal points
            ax.plot(filtered.loc[~mask_outlier, 'COLLECTION_DATE'],
                    filtered.loc[~mask_outlier, 'RESULT'],
                    'o-', color='tab:blue', label='Sample Data (clean)', **kwargs)

            # Plot outliers in a different marker/color
            ax.plot(filtered.loc[mask_outlier, 'COLLECTION_DATE'],
                    filtered.loc[mask_outlier, 'RESULT'],
                    'x', color='tab:purple', markersize=8, label='Sample Outliers')
        else:
            ax.plot('COLLECTION_DATE', 'RESULT', 'o-', data=sample_data,
                    color='tab:blue', label='Sample Data', **kwargs)


    def plot_sensor_pre_post(self, ax, sensor_data, sensor_name,
                            calib_records=None, remove_outliers=False, **kwargs):
        """
        Plot raw vs processed sensor data (calibration + outlier removal).
        """
        if sensor_data.empty:
            return

        # --- raw data first ---
        ax.plot('COLLECTION_DATE', 'RESULT', data=sensor_data,
                color='tab:orange', alpha=0.7, label='Sensor Data (raw)', **kwargs)

        processed = sensor_data.copy()

        # Apply calibration adjustments
        if calib_records is not None:
            sensor_calibs = calib_records[
                calib_records['analytes_to_adj'].apply(lambda lst: sensor_name in lst)
            ].copy()
            for _, calib_row in sensor_calibs.iterrows():
                calib_date = calib_row['calibration_date']
                processed = self.sensor_calibration_adjustment(processed, sensor_name, calib_date)

        # Apply Hampel outlier filter after adjustment
        if remove_outliers:
            processed['RESULT'] = self.hampel_filter(processed['RESULT'],
                                                    window_size=30, replace=True)

        # --- plot processed data ---
        ax.plot('COLLECTION_DATE', 'RESULT', data=processed,
                color='tab:red', alpha=0.9, label='Sensor Data (processed)', **kwargs)


    def plot_station_data_time_series(self, station_name, time_start, time_end,
                                  remove_sample_outliers=False, remove_sensor_outliers=False,
                                  calib_records=None, show_preprocess=True,
                                  save=False, fig_file_name=None, output_dir=None):

        # filter station and time range
        station_data = self.data[
            (self.data['STATION_ID'] == station_name) &
            (self.data['COLLECTION_DATE'] >= time_start) &
            (self.data['COLLECTION_DATE'] < time_end)
        ]

        analyte_info = [
            ("Specific Conductivity (µS/cm)", "SPECIFIC CONDUCTANCE", "Specific Conductivity - Sensor"),
            ("Tritium Concentration (pCi/mL)", "TRITIUM", None),
            ("Depth to Water (m)", "WATER LEVEL DEPTH", "Level: Depth to Water - Sensor"),
            ("pH", "PH", "pH - Sensor")
        ]

        fig, axes = plt.subplots(nrows=4, figsize=(7.5, 10), sharex=True)

        for i, (ylabel, sample_name, sensor_name) in enumerate(analyte_info):
            ax = axes[i]

            # --- sample data ---
            sample_data = station_data[station_data['ANALYTE_NAME'] == sample_name].copy()
            if show_preprocess:
                self.plot_sample_with_outliers(ax, sample_data, remove_outliers=remove_sample_outliers)
            else:
                if remove_sample_outliers:
                    sample_data['RESULT'] = self.hampel_filter(sample_data['RESULT'], window_size=5, n_sigmas=4)
                ax.plot('COLLECTION_DATE', 'RESULT', 'o-', data=sample_data,
                        color='tab:blue', label='Sample Data')

            # --- sensor data ---
            if sensor_name:
                sensor_data = station_data[station_data['ANALYTE_NAME'] == sensor_name].copy()
                if show_preprocess:
                    self.plot_sensor_pre_post(ax, sensor_data, sensor_name,
                                        calib_records=calib_records,
                                        remove_outliers=remove_sensor_outliers)
                else:
                    # keep your old pipeline for "processed only"
                    if calib_records is not None:
                        for _, calib_row in calib_records.iterrows():
                            calib_date = calib_row['calibration_date']
                            sensor_data = self.sensor_calibration_adjustment(sensor_data, sensor_name, calib_date)
                    if remove_sensor_outliers:
                        sensor_data['RESULT'] = self.hampel_filter(sensor_data['RESULT'], window_size=30, replace=True)
                    ax.plot('COLLECTION_DATE', 'RESULT', data=sensor_data,
                            color='tab:red', label='Sensor Data')

            ax.set_ylabel(ylabel)
            ax.grid(True, which='both', linestyle='--', linewidth=0.7, alpha=0.7)
            ax.legend()

        fig.suptitle(f"Station: {station_name}", fontsize=12)
        fig.tight_layout(rect=[0, 0, 1, 0.96])

        if save:
            if not os.path.exists(output_dir):
                os.makedirs(output_dir)
            station_name_alt = station_name.replace("-", "_")
            suffix = "_with_preprocess" if show_preprocess else "_processed_only"
            path = os.path.join(output_dir, f"{fig_file_name}_{station_name_alt}{suffix}.png")
            fig.savefig(path, dpi=300, bbox_inches='tight')


    # def plot_station_data_time_series(self, station_name, time_start, time_end,
    #                               remove_sample_outliers=False, remove_sensor_outliers=False, calib_records=None,
    #                               save=False, fig_file_name=None, output_dir=None):
    #     """
    #     Plots time series data of selected analytes for a specific station, with options to:
    #     - remove sample-based outliers using a Hampel filter
    #     - apply calibration adjustment to selected sensor analytes
    #     - save the figure with a descriptive file name

    #     Parameters:
    #         pylenm_data (pd.DataFrame): Full dataset with 'STATION_ID', 'COLLECTION_DATE', 'ANALYTE_NAME', 'RESULT'
    #         station_name (str): Name of the station to filter and plot
    #         time_start, time_end (datetime): Time window for filtering data
    #         remove_sample_outliers (bool): If True, applies Hampel filter to sample-based data
    #         adj_calib (bool): If True, applies calibration adjustment to specified sensor analytes
    #         analytes_to_adj (list): List of sensor analyte names to adjust (e.g., ['Specific Conductivity - Sensor'])
    #         calib_date (datetime): Calibration date used in adjustment
    #         save (bool): Whether to save the plot
    #         fig_file_name (str): Base file name for saving (used if save=True)
    #         output_dir (str): Directory to save the figure in
    #     """

    #     # filter station and time range
    #     station_data = self.data[
    #         (self.data['STATION_ID'] == station_name) &
    #         (self.data['COLLECTION_DATE'] >= time_start) &
    #         (self.data['COLLECTION_DATE'] < time_end)
    #     ]

    #     # define analytes to plot: (label, sample analyte, sensor analyte)
    #     analyte_info = [
    #         ("Specific Conductivity (µS/cm)", "SPECIFIC CONDUCTANCE", "Specific Conductivity - Sensor"),
    #         ("Tritium Concentration (pCi/mL)", "TRITIUM", None),
    #         ("Depth to Water (m)", "WATER LEVEL DEPTH", "Level: Depth to Water - Sensor"),
    #         ("pH", "PH", "pH - Sensor")
    #     ]

    #     fig, axes = plt.subplots(nrows=4, figsize=(9, 12), sharex=True)

    #     any_sensor_adjusted = False    # flag to check if any sensor data was adjusted
    #     for i, (ylabel, sample_name, sensor_name) in enumerate(analyte_info):
    #         ax = axes[i]

    #         # sample data
    #         sample_data = station_data[station_data['ANALYTE_NAME'] == sample_name].copy()
    #         if remove_sample_outliers and not sample_data.empty:
    #             sample_data['RESULT'] = self.hampel_filter(sample_data['RESULT'], window_size=5, n_sigmas=4)
    #         ax.plot('COLLECTION_DATE', 'RESULT', 'o-', data=sample_data, color='tab:blue', label='Sample Data')

    #         # sensor data
    #         if sensor_name:
    #             sensor_data = station_data[station_data['ANALYTE_NAME'] == sensor_name].copy()

    #             # --- apply calibration adjustment if applicable ---
    #             adjusted = False
    #             if calib_records is not None:
    #                 sensor_calibs = calib_records[
    #                     (calib_records['analytes_to_adj'].apply(lambda lst: sensor_name in lst))
    #                 ].copy()

    #                 for j, calib_row in sensor_calibs.iterrows():
    #                     calib_date = calib_row['calibration_date']
    #                     sensor_data = self.sensor_calibration_adjustment(sensor_data, sensor_name, calib_date)
    #                     label = f'Adjusted Sensor Data' if j == sensor_calibs.index[0] else None
    #                     adjusted = True
    #                     any_sensor_adjusted = True

    #             # --- optionally remove sensor outliers AFTER adjustment ---
    #             if remove_sensor_outliers:
    #                 sensor_data['RESULT'] = self.hampel_filter(sensor_data['RESULT'], window_size=30, replace=True)

    #             # --- plot final processed sensor data ---
    #             final_label = label if adjusted else 'Sensor Data'
    #             final_color = 'tab:red' if adjusted else 'tab:orange'
    #             ax.plot('COLLECTION_DATE', 'RESULT', data=sensor_data, color=final_color, alpha=1, label=final_label)


    #         # style settings
    #         ax.set_ylabel(ylabel)
    #         ax.set_xlabel("Date Time")
    #         ax.grid(True, which='both', linestyle='--', linewidth=0.7, alpha=0.7)
    #         ax.legend()

    #     # title and layout
    #     fig.suptitle(f"Station: {station_name}", fontsize=12)
    #     fig.tight_layout(rect=[0, 0, 1, 0.96])

    #     # save figure if requested
    #     if save:
            
    #         if not os.path.exists(output_dir):
    #             os.makedirs(output_dir)
            
    #         station_name_alt = station_name.replace("-", "_")

    #         if remove_sample_outliers and any_sensor_adjusted:
    #             suffix = "_outliers_removed_sensor_data_adjusted"
    #         elif remove_sample_outliers:
    #             suffix = "_outliers_removed"
    #         # elif any_sensor_adjusted:
    #         #     suffix = "_sensor_data_adjusted"
    #         else:
    #             suffix = ""
             
    #         path = os.path.join(output_dir, f"{fig_file_name}_{station_name_alt}{suffix}.png")
    #         fig.savefig(path, dpi=300, bbox_inches='tight')

    
    def process_station_data(self, station_name, time_start, time_end,
                             remove_sample_outliers=False, remove_sensor_outliers=False, calib_records=None):
        """
        Process station data by filtering, applying calibration adjustments, and removing outliers.

        Returns:
            processed_data (dict): Dictionary with keys = analyte labels and values = dicts of sample and sensor DataFrames
            any_sensor_adjusted (bool): Whether any sensor analyte was adjusted
            processed_df (pd.DataFrame): Combined processed DataFrame in original long format
        """

        station_data = self.data[
            (self.data['STATION_ID'] == station_name) &
            (self.data['COLLECTION_DATE'] >= time_start) &
            (self.data['COLLECTION_DATE'] < time_end)
        ]

        analyte_info = [
            ("SPECIFIC CONDUCTANCE", "Specific Conductivity - Sensor"),
            ("TRITIUM", None),
            ("STRONTIUM-90", None),
            ("IODINE-129", None),
            ("URANIUM-238", None),
            ("WATER LEVEL DEPTH", "Level: Depth to Water - Sensor"),
            ("PH", "pH - Sensor")
        ]

        # processed_data = {}
        all_processed_dfs = []
        any_sensor_adjusted = False

        for sample_name, sensor_name in analyte_info:
            sample_data = station_data[station_data['ANALYTE_NAME'] == sample_name].copy()
            sensor_data = station_data[station_data['ANALYTE_NAME'] == sensor_name].copy() if sensor_name else None

            if remove_sample_outliers and not sample_data.empty:
                sample_data['RESULT'] = self.hampel_filter(sample_data['RESULT'], window_size=5, n_sigmas=4)
            if not sample_data.empty:
                all_processed_dfs.append(sample_data)

            if sensor_name and not sensor_data.empty:
                adjusted = False
                if calib_records is not None:
                    sensor_calibs = calib_records[
                        (calib_records['analytes_to_adj'].apply(lambda lst: sensor_name in lst))
                    ].copy()
                    for _, calib_row in sensor_calibs.iterrows():
                        calib_date = calib_row['calibration_date']
                        sensor_data = self.sensor_calibration_adjustment(sensor_data, sensor_name, calib_date)
                        adjusted = True
                        any_sensor_adjusted = True

                if remove_sensor_outliers:
                    sensor_data['RESULT'] = self.hampel_filter(sensor_data['RESULT'], window_size=30, replace=True)

                all_processed_dfs.append(sensor_data)

        processed_df = pd.concat(all_processed_dfs, axis=0).sort_values("COLLECTION_DATE")
        
        return any_sensor_adjusted, processed_df


    def calculate_decay_stats(self, YYs, n_grid, grid_shape):
        """
        Calculate log-linear decay trend statistics from a sequence of annual grids.

        Parameters
        ----------
        YYs : list of np.ndarray
            List of interpolated values (one array per year).
        n_grid : int
            Total number of grid cells.
        grid_shape : tuple
            Shape of the 2D grid (e.g., (285, 250)).

        Returns
        -------
        a : np.ndarray
            Slope of the log-linear trend at each grid cell.
        b : np.ndarray
            Intercept of the log-linear trend at each grid cell.
        var_y : np.ndarray
            Variance of residuals at each grid cell.
        """
        n_year = len(YYs)
        year = np.arange(n_year)  # use 0,1,... as indices

        # --- Initialize sums ---
        sumy = np.zeros(n_grid)
        sumxy = np.zeros(n_grid)
        sumyy = np.zeros(n_grid)
        sumer = np.zeros(n_grid)

        # --- Accumulate sums across years ---
        for i, YY in enumerate(YYs):
            YY_flat = YY.ravel()
            sumy += YY_flat
            sumxy += i * YY_flat
            sumyy += YY_flat**2

        # --- Precompute means and sums ---
        sumx = year.sum()
        sumxx = (year**2).sum()
        avy = sumy / n_year
        avx = sumx / n_year

        # --- Slope (a) and intercept (b) ---
        S_xx = sumxx - 2 * avx * sumx + n_year * avx**2
        S_xy = sumxy - avx * sumy - avy * sumx + n_year * avx * avy
        a = S_xy / S_xx
        b = avy - a * avx

        # --- Residual variance ---
        for i, YY in enumerate(YYs):
            YY_flat = YY.ravel()
            sumer += (YY_flat - (a * i + b))**2

        var_y = sumer / (n_year - 2)

        return a.reshape(grid_shape), b.reshape(grid_shape), var_y.reshape(grid_shape)


    def sub_sample_grid(self, YY0, xxi, yyi, a, b, var_y, grid_shape=(285, 250), step=3):
        """
        Sub-sample spatial grid data to reduce size.

        Parameters
        ----------
        YY0 : np.ndarray
            1D or 2D array of initial estimates (sample-based).
        XX : np.ndarray
            (n_points, 2) array of original grid coordinates [x, y].
        a, b, var_y : np.ndarray
            2D arrays (grid_shape) of decay statistics from calculate_decay_stats().
        grid_shape : tuple, optional
            Shape of the full grid (default = (285, 250)).
        step : int, optional
            Sub-sampling step size along each axis (default = 3).

        Returns
        -------
        YY0_sub : np.ndarray
            Flattened sub-sampled estimates.
        XX_sub : np.ndarray
            (n_sub_points, 2) sub-sampled grid coordinates.
        a_sub, b_sub, var_y_sub : np.ndarray
            Flattened sub-sampled decay statistics.
        sub_grid_shape : tuple
            Shape of the sub-sampled 2D grid.
        """

        # Reshape original data to full 2D grid
        YY0 = YY0.reshape(grid_shape)
        # XX_x = XX[:, 0].reshape(grid_shape)
        # XX_y = XX[:, 1].reshape(grid_shape)

        # Sub-sample by step
        YY0_sub = YY0[::step, ::step].flatten()
        xxi_sub = xxi[::step, ::step]
        yyi_sub = yyi[::step, ::step]
        XX_sub = np.column_stack((xxi_sub.flatten(), yyi_sub.flatten()))
        sub_grid_shape = xxi_sub.shape

        # Sub-sample a, b, var_y
        a_sub = a[::step, ::step].flatten()
        b_sub = b[::step, ::step].flatten()
        var_y_sub = var_y[::step, ::step].flatten()

        return YY0_sub, XX_sub, xxi_sub, yyi_sub, a_sub, b_sub, var_y_sub, sub_grid_shape


    def calculate_posterior_variance(self, XX, X_sensor, n_sensor, n_grid, gp_kernel, var_y, sigma2):
        """
        Calculate posterior variance matrix Q for Gaussian Process + sample-based interpolation.

        Parameters
        ----------
        XX : np.ndarray
            (n_grid, n_features) array of grid coordinates for predictions.
        X_sensor : np.ndarray
            (n_sensor, n_features) array of sensor coordinates.
        n_sensor : int
            Number of sensors.
        n_grid : int
            Number of grid cells.
        gp_kernel : sklearn.gaussian_process.kernels.Kernel
            Kernel function for Gaussian Process.
        var_y : np.ndarray
            (grid_shape) or flattened array of sample-based residual variances.
        sigma2 : float or np.ndarray
            Sensor error variance (scalar or length-n_sensor array).

        Returns
        -------
        Q : np.ndarray
            (n_grid, n_grid) posterior variance matrix.
        A : np.ndarray
            (n_sensor, n_grid) measurement matrix mapping grid to sensors.
        R : np.ndarray
            (n_sensor, n_sensor) sensor covariance matrix.
        S : np.ndarray
            (n_grid, n_grid) combined sample-based + GP covariance matrix.
        gp_kernel : sklearn.gaussian_process.kernels.Kernel
            Updated GP kernel after fitting.
        """

        # --- GP fit + covariance prediction ---
        gpr = GaussianProcessRegressor(kernel=gp_kernel, optimizer=False, normalize_y=True)
        _, YCOV = gpr.predict(XX, return_cov=True)
        gp_kernel = gpr.kernel  # updated kernel

        # --- Covariance matrix for the sample-based map ---
        S = np.zeros((n_grid, n_grid))
        np.fill_diagonal(S, var_y.flatten())
        S = S + YCOV

        # --- Measurement matrix (map grid <-> sensors) ---
        A = np.zeros((n_sensor, n_grid))
        for i in range(n_sensor):
            # Distance from this sensor to all grid cells
            d = (XX[0,:] - X_sensor[i, 0])**2 + (XX[1,:] - X_sensor[i, 1])**2
            idx = np.argmin(d)
            A[i, idx] = 1

        # --- Combined covariance ---
        R = np.zeros((n_sensor, n_sensor))
        np.fill_diagonal(R, sigma2)

        R = gp_kernel(X_sensor) + R
        Q = np.linalg.solve(R, A)
        Q = np.linalg.solve(S, np.eye(n_grid)) + A.T @ Q

        return Q, A, R, S, gp_kernel


    def calculate_posterior_mean(self, i, sc_interp, t0, YY0, a, b, A, S, R, Q, conversion_func):
        """
        Calculate posterior mean (updated estimates) at a given timestep.

        Parameters
        ----------
        i : int
            Time index into sc_interp.
        sc_interp : pd.DataFrame
            Interpolated specific conductivity values at sensor locations (rows = time, cols = wells).
        t0 : pd.Timestamp
            Reference start time.
        YY0 : np.ndarray
            Initial estimates (grid-based).
        a, b : np.ndarray
            Slope and intercept arrays from log-linear decay model.
        A : np.ndarray
            Measurement matrix mapping grid to sensor locations.
        S : np.ndarray
            Sample-based + GP covariance matrix (n_grid x n_grid).
        R : np.ndarray
            Sensor covariance matrix (n_sensor x n_sensor).
        Q : np.ndarray
            Posterior variance matrix (n_grid x n_grid).
        conversion_func : callable
            Function to convert specific conductivity to tritium (H3). 
            Signature: y_h3 = conversion_func(y_sc).

        Returns
        -------
        y_est : np.ndarray
            Posterior mean estimates at grid cells for this timestep.
        delta : pandas.Timedelta
            Time difference from t0 to sc_interp.index[i].
        """

        # --- Time delta ---
        delta = sc_interp.index[i] - t0

        # --- Sensor measurement (converted to H3) ---
        y = np.array(sc_interp.iloc[i, :])  # conductivity
        y = conversion_func(y)              # convert to H3

        # --- Prior mean estimate with log-linear decay ---
        YY = YY0 + a * delta.days / 365.0

        # --- Posterior update ---
        step1 = np.linalg.solve(R, y)
        step2 = np.linalg.solve(S, YY) + A.T @ step1
        y_est = np.linalg.solve(Q, step2)

        return y_est, delta


    # def plot_spatial_estimation_map(self, X, y, XX, xxi, yyi, reg_model, reg_features, contour_levels,
    #                                 station_names, basin_boundaries, fig_title, cbar_label, gp_kernel=None,
    #                                 save_path=None, fig_name=None,
    #                                 vmin=55.0, vmax=70.0,
    #                                 annotate_stations=True, add_flow_directions=False,
    #                                 fontsize=15, save=False, return_YY=False, f_area=False):
    #     """
    #     Performs spatial interpolation and generates a prediction map with elevation annotations.

    #     Parameters:
    #         X (pd.DataFrame): Training features with columns ['Easting', 'Northing', 'Elevation']
    #         y (np.ndarray): Target values corresponding to X
    #         XX (np.ndarray): Prediction grid locations with same feature structure as X
    #         xxi, yyi (np.ndarray): Meshgrid matching XX for plotting
    #         reg_model (str): Regression model to use ('gp' for Gaussian Process, 'regression' for regression + GP)
    #         reg_features (list of str): Features to use for regression (e.g., ['Easting', 'Northing', 'Elevation'])
    #         contour_levels (list of float): Levels for contour lines
    #         station_names (list of str): Names of the stations to be plotted
    #         basin_boundaries (list of np.ndarray): List of arrays of shape (N, 2) for basin outlines
    #         fig_title (str): Title (typically a date) to show on the plot
    #         cbar_label (str): Label for the color bar
    #         gp_kernel (sklearn kernel, optional): Kernel for Gaussian Process regression. If None, uses default.
    #         save_path (str): Path to save figure. If None and save=True, figure won't be saved.
    #         fig_name (str): Base name for the figure file
    #         vmin, vmax (float): Color range limits for the contour map
    #         annotate_stations (bool): Whether to annotate station names on the map
    #         fontsize (int): Font size for labels
    #         save (bool): Whether to save the figure
    #         return_YY (bool): If True, returns the interpolated values YY
    #         f_area (bool): If True, applies F-Area site-specific characteristics to the interpolated values
    #     """
        
    #     # Spatial interpolation
    #     if reg_model == 'gp':

    #         # Get predictions from Gaussian Process model
    #         gpr = GaussianProcessRegressor(kernel=gp_kernel, optimizer=None, normalize_y=True).fit(X, y)
    #         YY = gpr.predict(XX)

        
    #     else:
            
    #         # Get predictions from regression + GP model
    #         YY, _, _, _, _ = self.interpolate_topo(X, y, XX,
    #                                                ft=reg_features, gp_kernel=gp_kernel,
    #                                                regression=reg_model, smooth=True)


    #     # Adjust YY for F-Area site-specific characteristic if requested
    #     if f_area:

    #         # Create shapely LineString from river_line
    #         river_line = basin_boundaries[-1]  # assuming the last basin boundary is the river line
    #         river = LineString(river_line)  # river_line must be a (N,2) array

    #         # Flatten the grid
    #         xx_flat = xxi.ravel()
    #         yy_flat = yyi.ravel()
    #         YY_flat = YY.ravel()

    #         # Initialize mask
    #         mask = np.zeros_like(YY_flat, dtype=bool)

    #         # Loop through grid points
    #         for i, (x, y) in enumerate(zip(xx_flat, yy_flat)):
    #             point = Point(x, y)
    #             nearest_point_on_river = nearest_points(point, river)[1]  # second is on the river

    #             # If point is southeast of its nearest river point
    #             if point.x > nearest_point_on_river.x and point.y < nearest_point_on_river.y:
    #                 mask[i] = True

    #         # Apply mask
    #         YY_flat[mask] = 0    # np.log10(1)    (1 pCi/mL, log10 scale)
    #         YY = YY_flat.reshape(YY.shape)


    #     # Visualization
    #     fig, ax = plt.subplots(figsize=(5, 5), dpi=300)
    #     bounds = np.linspace(YY.min(), YY.max(), 50)
    #     norm = colors.BoundaryNorm(boundaries=bounds, ncolors=256, extend='both')
    #     # cmap = plt.cm.get_cmap('YlGnBu_r')
    #     cmap = plt.cm.get_cmap('jet')

    #     map1 = ax.pcolor(xxi, yyi, YY.reshape(xxi.shape), cmap=cmap, vmin=vmin, vmax=vmax)
    #     fig.colorbar(map1, extend='both', ax=ax).set_label(label=cbar_label, size=fontsize)

    #     # Plot boundaries
    #     for basin in basin_boundaries:
    #         ax.plot(basin[:, 0], basin[:, 1], 'w', zorder=10)

    #     # Plot contour lines
    #     ctr = ax.contour(xxi, yyi, YY.reshape(xxi.shape), levels=contour_levels, colors='black', alpha=0.7, vmin=vmin, vmax=vmax, linewidths=0.3)
    #     # ax.clabel(ctr, inline=True, fontsize=8, fmt='%1.1f', colors='white', use_clabeltext=True)

    #     # Plot station points
    #     sc = ax.scatter(X['Easting'], X['Northing'],
    #                     c='none', vmin=vmin, vmax=vmax,
    #                     edgecolor='tab:red', linewidth=0.5, s=40, zorder=3)
        
    #     # Annotate station names
    #     if annotate_stations:
    #         for j in range(len(X)):
    #             # annotation = ": ".join([station_names[j],
    #             #                        "{:.1f}".format(y[j])])
    #             ax.text(X.iloc[j, 0] + 4, X.iloc[j, 1] + 6,
    #                     station_names[j], fontsize=9, color='white',
    #                     ha='center', va='bottom')

    #     # Add flow directions if requested
    #     if add_flow_directions:
            
    #         GRAD_SPACING = 40    # Spacing for gradient arrows
            
    #         xx_grad, yy_grad = np.meshgrid(xxi[0,::GRAD_SPACING], yyi[::GRAD_SPACING,0])  # gradient spacing
    #         YY_grad = YY.reshape(xxi.shape)[::GRAD_SPACING, ::GRAD_SPACING]  # reshape to match gradient spacing
    #         y_grad, x_grad = np.gradient(YY_grad, yyi[::GRAD_SPACING,0], xxi[0,::GRAD_SPACING])  # gradients in y and x

    #         # Normalize gradients for quiver plot
    #         ax.quiver(xx_grad, yy_grad, -x_grad, -y_grad, color='gray', angles='xy', pivot='mid',
    #         scale=2e-4, scale_units='dots', alpha=0.6)  # standardize arrows so ex. 0.01 m/m looks the same on every plot (disable autoscaling)
            


    #     # Plot styling
    #     ax.set_title(str(fig_title))
    #     ax.set_xlabel("Easting (NAD83), m", fontsize=fontsize)
    #     ax.set_ylabel("Northing (NAD83), m", fontsize=fontsize)
    #     ax.set_xlim([xxi.flatten().min(), 4.373e5])
    #     ax.set_ylim([yyi.flatten().min(), 3.6824e6])
    #     ax.set_aspect('equal', 'box')
    #     ax.tick_params(axis='both', which='major', labelsize=fontsize)
    #     ax.ticklabel_format(style='sci', axis='both', scilimits=(-1, 1), useMathText=True)
    #     plt.locator_params(axis='both', nbins=4, tight=False)
    #     # plt.tight_layout()

    #     # Save or return
    #     if save and save_path:
    #         file_name = fig_name + "_" +fig_title + ".png"
    #         # print(os.path.join(save_path,file_name))
    #         fig.savefig(os.path.join(save_path,file_name), bbox_inches='tight')
    #         plt.close(fig)


    #     if return_YY:
    #         return YY


    def generate_spatial_estimation(self, X, y, XX, reg_model, reg_features=None, gp_kernel=None,
                                    save_data=False, save_path=None, file_name=None):
        """
        Performs spatial interpolation using Gaussian Process or regression + GP.

        Parameters:
            X (pd.DataFrame): Training features with columns ['Easting', 'Northing', 'Elevation']
            y (np.ndarray): Target values corresponding to X
            XX (np.ndarray): Prediction grid locations with same feature structure as X
            reg_model (str): Regression model to use ('gp' for Gaussian Process, 
                            otherwise regression + GP)
            reg_features (list of str, optional): Features to use for regression
            gp_kernel (sklearn kernel, optional): Kernel for Gaussian Process regression. 
                                                If None, uses default.
            save_data (bool, optional): If True, save YY as a .npy file (default=False).
            save_path (str, optional): Directory path to save the .npy file. Required if save_data=True.
            file_name (str, optional): File name for the saved .npy file (without extension). Required if save_data=True.

        Returns:
            YY (np.ndarray): Interpolated values at grid locations XX
        """
        
        if reg_model == 'gp':

            # Get predictions from Gaussian Process model
            gpr = GaussianProcessRegressor(kernel=gp_kernel, optimizer=None, normalize_y=True).fit(X, y)
            YY = gpr.predict(XX)
        
        else:
            
            # Get predictions from regression + GP model
            YY, _, _, _, _ = self.interpolate_topo(
                X, y, XX,
                ft=reg_features, gp_kernel=gp_kernel,
                regression=reg_model, smooth=True
            )

        # --- Save as .npy file if requested ---
        if save_data:
            if save_path is None or file_name is None:
                raise ValueError("Both save_path and file_name must be provided when save_data=True")
            np.save(os.path.join(save_path, file_name + ".npy"), YY)


        return YY
    

    def plot_spatial_estimation_map(self, X, YY, xxi, yyi,
                                    station_names, basin_boundaries,
                                    fig_title, cbar_label, contour_levels,
                                    save_path=None, fig_name=None,
                                    vmin=55.0, vmax=70.0,
                                    annotate_stations=True,
                                    add_barriers=False, barrier_lines=None,
                                    add_flow_directions=False,
                                    add_MCL_line=False, MCL_value=None,
                                    fontsize=15, save=False, f_area=False, f_area_value=0):
        """
        Generates a prediction map with optional F-Area site-specific adjustment.

        Parameters:
            X (pd.DataFrame): Training features with columns ['Easting', 'Northing']
            y (np.ndarray): Target values corresponding to X
            XX (np.ndarray): Prediction grid locations
            xxi, yyi (np.ndarray): Meshgrid for plotting
            YY (np.ndarray): Interpolated values from generate_spatial_estimation()
            station_names (list of str): Names of the stations
            basin_boundaries (list of np.ndarray): List of arrays (N,2) for basin outlines
            fig_title (str): Title (e.g., a date)
            cbar_label (str): Label for the color bar
            contour_levels (list of float): Contour levels
            save_path (str): Path to save figure
            fig_name (str): Base name for saved figure
            vmin, vmax (float): Colorbar limits
            annotate_stations (bool): Whether to annotate station names
            add_barriers (bool): Whether to add barrier boundaries
            barrier_lines (list of tuple): List of barrier line coordinates [(coord0, coord1), ...]
            add_flow_directions (bool): Whether to add flow arrows
            add_MCL_line (bool): Whether to add MCL contour line
            MCL_value (float): MCL value to plot (if add_MCL_line is True)
            fontsize (int): Font size
            save (bool): Whether to save the figure
            f_area (bool): Apply F-Area site-specific adjustment before plotting
            f_area_value (float): Value to set in F-Area adjustment (default=0, i.e., 1 pCi/mL in log10 scale)
        """

        # Adjust YY for F-Area site-specific characteristic if requested
        if f_area:
            river_line = basin_boundaries[-1]  # assuming last boundary is river line
            river = LineString(river_line)

            xx_flat = xxi.ravel()
            yy_flat = yyi.ravel()
            YY_flat = YY.ravel()

            mask = np.zeros_like(YY_flat, dtype=bool)
            for i, (x, y_) in enumerate(zip(xx_flat, yy_flat)):
                point = Point(x, y_)
                nearest_point_on_river = nearest_points(point, river)[1]
                if point.x > nearest_point_on_river.x and point.y < nearest_point_on_river.y:
                    mask[i] = True

            YY_flat[mask] = f_area_value  # site-specific adjustment
            YY = YY_flat.reshape(YY.shape)

        # === Plotting ===
        fig, ax = plt.subplots(figsize=(5, 5), dpi=300)
        # cmap = plt.cm.get_cmap('jet')
        cmap = plt.cm.get_cmap('YlGnBu_r')

        # Create spatial estimation map with color bar
        map1 = ax.pcolor(xxi, yyi, YY.reshape(xxi.shape), cmap=cmap, vmin=vmin, vmax=vmax)
        fig.colorbar(map1, extend='both', ax=ax).set_label(label=cbar_label, size=fontsize)

        # Plot boundaries
        for basin in basin_boundaries:
            ax.plot(basin[:, 0], basin[:, 1], 'w', zorder=5)

        # Plot contour lines
        ctr = ax.contour(xxi, yyi, YY.reshape(xxi.shape),
                         levels=contour_levels, colors='black', alpha=0.7,
                         vmin=vmin, vmax=vmax, linewidths=0.3)

        # Add MCL contour line if requested
        if add_MCL_line and MCL_value is not None:
            MCL_value_log = np.log10(MCL_value)
            mcl_ctr = ax.contour(xxi, yyi, YY.reshape(xxi.shape), levels=[MCL_value_log], colors='tab:red', linestyles='--', linewidths=1.2, zorder=12)
            # ax.clabel(mcl_ctr, inline=False, fontsize=fontsize-3,
            #           fmt={MCL_value_log: f"MCL={MCL_value} pCi/mL"}, colors='tab:red', use_clabeltext=False)

            # Add custom legend entry instead of inline label
            legend_line = Line2D([0], [0], color='tab:red', linewidth=1.2, linestyle='--')
            legend = ax.legend([legend_line], [f"MCL = {MCL_value} pCi/mL"], loc="lower right", labelcolor="tab:red", fontsize=15, frameon=False)
            legend.set_zorder(15)

        # Plot station points
        ax.scatter(X['Easting'], X['Northing'],
                   c='none', edgecolor='tab:red',
                   linewidth=0.5, s=40, zorder=3)

        # Annotate stations if requested
        if annotate_stations:
            for j in range(len(X)):
                ax.text(X.iloc[j, 0] + 4, X.iloc[j, 1] + 6,
                        station_names[j], fontsize=9, color='black',
                        ha='center', va='bottom')
                
        # Add barrier boundaries if requested
        if add_barriers and barrier_lines is not None:

            BARRIER_BUFFER = 50  # meters
            barrier_geoms = []

            for coord0, coord1 in barrier_lines:
                ax.plot([coord0[0], coord1[0]], [coord0[1], coord1[1]], color='tab:red', linestyle='-', linewidth=3.5, alpha=0.8, zorder=6)
                
                barrier_geoms.append(LineString([coord0, coord1]))
                barrier_multiline = MultiLineString(barrier_geoms)

            # Plot the buffer zone
            buffer_zone = barrier_multiline.buffer(BARRIER_BUFFER)
            # patches = [Polygon(np.array(poly.exterior.coords)) for poly in buffer_zone.geoms]
            # ax.add_collection(PatchCollection(patches, facecolor='tab:red', edgecolor='none', alpha=0.4, zorder=5))

        # Add flow directions if requested
        if add_flow_directions:

            GRAD_SPACING = 40    # Spacing for gradient arrows

            xx_grad, yy_grad = np.meshgrid(xxi[0, ::GRAD_SPACING], yyi[::GRAD_SPACING, 0])    # gradient spacing
            YY_grad = YY.reshape(xxi.shape)[::GRAD_SPACING, ::GRAD_SPACING]    # reshape to match gradient spacing
            y_grad, x_grad = np.gradient(YY_grad, yyi[::GRAD_SPACING, 0], xxi[0, ::GRAD_SPACING])    # gradients in y and x
            
            # Apply barrier effect on flow directions
            if add_barriers and barrier_lines is not None:
                points = np.column_stack((xx_grad.ravel(), yy_grad.ravel()))
                barrier_mask = np.array([Point(p).distance(barrier_multiline) < BARRIER_BUFFER for p in points])
                barrier_mask = barrier_mask.reshape(xx_grad.shape)

                # Mask gradients near barriers
                x_grad[barrier_mask] = 0
                y_grad[barrier_mask] = 0

            # Normalize gradients for quiver plot
            ax.quiver(xx_grad, yy_grad, -x_grad, -y_grad,
                      color='gray', angles='xy', pivot='mid',
                      scale=2e-4, scale_units='dots', alpha=0.6)

        # Plot styling
        ax.set_title(str(fig_title))
        ax.set_xlabel("Easting (NAD83), m", fontsize=fontsize)
        ax.set_ylabel("Northing (NAD83), m", fontsize=fontsize)
        ax.set_xlim([xxi.flatten().min(), 4.373e5])
        ax.set_ylim([yyi.flatten().min(), 3.6824e6])
        ax.set_aspect('equal', 'box')
        ax.tick_params(axis='both', which='major', labelsize=fontsize)
        ax.ticklabel_format(style='sci', axis='both', scilimits=(-1, 1), useMathText=True)
        plt.locator_params(axis='both', nbins=4, tight=False)

        # Save figure
        if save and save_path:
            file_name = fig_name + "_" + str(fig_title) + ".png"
            fig.savefig(os.path.join(save_path, file_name), bbox_inches='tight')
            plt.close(fig)

    
    def get_approx_predictions(self, X, y_map, XX):
        """
        Approximates the closest prediction grid points in XX for each training location in X.

        This function finds, for each (Easting, Northing) location in X, the closest 
        corresponding location in XX based on the maximum coordinate difference (Chebyshev distance). 
        It returns the matched coordinates from XX and their corresponding predicted y_map values.

        Parameters:
            X (pd.DataFrame): Training locations with columns ['Easting', 'Northing', 'Elevation'].
            y_map (np.ndarray or list): Predicted values corresponding to grid locations in XX.
            XX (pd.DataFrame): Prediction grid locations with columns ['Easting', 'Northing', 'Elevation'].

        Returns:
            X_approx (pd.DataFrame): Approximated coordinates from XX matching each row in X.
            y_approx (list): Corresponding predicted values from y_map for the approximated locations.
        """
        X_approx = []
        y_approx = []
        # idx = []
        XX_cols = XX.columns  # Get the columns of XX to match with X

        for _, row in X.iterrows():
            # Compute Chebyshev distance (max absolute difference in Easting and Northing)
            abs_east = np.abs(XX['Easting'] - row['Easting'])
            abs_north = np.abs(XX['Northing'] - row['Northing'])
            chebyshev_dist = np.maximum(abs_east, abs_north)

            # Find index of the closest prediction location
            index = chebyshev_dist.idxmin()

            # Append matched coordinates and predicted value
            X_approx.append(XX.loc[index, XX_cols].values)
            y_approx.append(y_map[index])
            # idx.append(index)

        # Convert results to DataFrame
        X_approx = pd.DataFrame(X_approx, columns=XX_cols)

        return X_approx, y_approx
    

    def leave_one_out_cv_spatial(self, X, XX, y_time_series, reg_model, reg_features, gp_kernel=None):
        """
        Performs Leave-One-Out Cross-Validation (LOO-CV) for spatial interpolation models 
        and computes model performance metrics over multiple time steps.

        Parameters:
            X (pd.DataFrame): Training features with columns like ['Easting', 'Northing', 'Elevation'].
            XX (pd.DataFrame): Prediction grid features matching X's columns.
            y_time_series (pd.DataFrame): Target time series (stations as columns, time as index).
            reg_model (str): Name of the regression model used (for logging results).
            reg_features (list): List of feature column names to use for interpolation.
            gp_kernel (sklearn kernel, optional): Kernel for Gaussian Process regression. If None, uses default.

        Returns:
            pd.DataFrame: DataFrame containing LOO-CV results with columns ['model', 'features', 'time', 'mse', 'r2'].
        """

        # Initialize results DataFrame
        model_results_loo = pd.DataFrame(columns=['model', 'features', 'time', 'mse', 'r2'])

        for i in range(y_time_series.shape[0]):
            
            # Extract target values for the current time step
            y = np.array(y_time_series.iloc[i, :])
            time = y_time_series.index[i].date()

            y_approx_loo = []
            loo = LeaveOneOut()

            for train_index, test_index in loo.split(X):
                
                X_train, X_test = X.iloc[train_index], X.iloc[test_index]
                y_train, y_test = y[train_index], y[test_index]

                # Spatial interpolation
                if reg_model == 'gp':

                    # Get predictions from Gaussian Process model
                    gpr = GaussianProcessRegressor(kernel=gp_kernel, optimizer=None, normalize_y=True).fit(X, y)
                    y_map = gpr.predict(XX)
                
                else:
                    
                    # Get predictions from regression + GP model
                    y_map, _, _, _, _ = self.interpolate_topo(X, y, XX,
                                                        ft=reg_features, gp_kernel=gp_kernel,
                                                        regression=reg_model, smooth=True)

                # Approximate prediction at the test location
                X_approx_test, y_approx_test = self.get_approx_predictions(X_test, y_map, XX)
                y_approx_loo.extend(y_approx_test)  # Append predictions for each test point

            # Compute error metrics

            mse_value = self.mse(y, y_approx_loo)  # Default MSE
            r2_value = r2_score(y, y_approx_loo)

            # Save results
            model_results_loo.loc[i] = [reg_model, reg_features, time, mse_value, r2_value]

        return model_results_loo



    def summarize_with_time_window(self, query, analyte_main, analytes, time_window_days=7, min_samples=10, 
                                   remove_outliers=False, z_threshold=4, log_transform=False, no_log=None):
        """
        Generalized summarization using a centered time window around a main analyte.

        For each measurement of the main analyte, computes mean values of other analytes 
        within ± time_window_days around the COLLECTION_DATE.

        Parameters:
            query (pd.DataFrame): Data with ['COLLECTION_DATE', 'STATION_ID', 'ANALYTE_NAME', 'RESULT'].
            analyte_main (str): The main analyte name to keep all data points for (e.g., 'tr-sample').
            analytes (list): Other analytes to calculate rolling means for (e.g., ['sc-sensor', 'sc-sample']).
            time_window_days (int): Time window in days (±). Default is 7.
            min_samples (int): Minimum number of samples required. Default is 10.
            remove_outliers (bool): Whether to remove outliers. Default is False.
            z_threshold (float): Z-score threshold for outlier removal. Default is 4.
            log_transform (bool): Whether to apply log10 transform to results. Default is False.
            no_log (list): List of analyte result columns to exclude from log transform. Default is None.

        Returns:
            pd.DataFrame or str: Result summary DataFrame or error message string.
        """
        query = query.copy()
        # query['COLLECTION_DATE'] = pd.to_datetime(query['COLLECTION_DATE'])

        # Validate required analytes
        required_analytes = [analyte_main] + analytes
        available_analytes = query['ANALYTE_NAME'].unique()
        missing_analytes = [a for a in required_analytes if a not in available_analytes]
        if missing_analytes:
            return f"ERROR: Missing analytes in data: {', '.join(missing_analytes)}"

        query = query[query['ANALYTE_NAME'].isin(required_analytes)]
        if query.empty:
            return "ERROR: No data available for the selected analytes."

        samples = query[['COLLECTION_DATE', 'STATION_ID', 'ANALYTE_NAME']].drop_duplicates().shape[0]
        if samples < min_samples:
            return f"ERROR: Dataset has fewer than {min_samples} unique samples."

        # Prepare datasets
        main_df = query[query['ANALYTE_NAME'] == analyte_main].copy()
        analyte_dfs = {a: query[query['ANALYTE_NAME'] == a].copy() for a in analytes}

        # Initialize result DataFrame
        result_df = main_df[['COLLECTION_DATE', 'STATION_ID', 'RESULT']].rename(columns={'RESULT': analyte_main})

        # Compute rolling means for each analyte
        for analyte in analytes:
            mean_results = []
            analyte_df = analyte_dfs[analyte]

            for _, row in result_df.iterrows():
                station = row['STATION_ID']
                date = row['COLLECTION_DATE']
                window_start = date - datetime.timedelta(days=time_window_days)
                window_end = date + datetime.timedelta(days=time_window_days)

                values_in_window = analyte_df[
                    (analyte_df['STATION_ID'] == station) & 
                    (analyte_df['COLLECTION_DATE'].between(window_start, window_end))
                ]['RESULT']

                mean_results.append(values_in_window.mean())

            result_df[analyte] = mean_results

        result_df = result_df.set_index(['STATION_ID','COLLECTION_DATE'])

        # Optional: Remove outliers
        if remove_outliers:
            result_df = self.remove_outliers(result_df, z_threshold=z_threshold)

        # Optional: Log transformation
        if log_transform:
            numeric_cols = result_df.select_dtypes(include=[np.number]).columns
            result_df[numeric_cols] = result_df[numeric_cols].replace(0, 1e-8)  # Avoid log(0)
            temp_result = result_df.copy()
            result_df[numeric_cols] = np.log10(result_df[numeric_cols])

            # Revert excluded columns from log transform
            if no_log:
                for col in no_log:
                    if col in temp_result:
                        result_df[col] = temp_result[col]

        return result_df
    


    def get_MCL_results(self, station_name, analyte_name):
        """Plots the linear regression line of data given the analyte_name and well_name. The plot includes the prediction where the line of best fit intersects with the Maximum Concentration Limit (MCL).

        Parameters:
            station_name (str): name of the station to be processed
            analyte_name (str): name of the analyte to be processed

        Returns:
            tuple: parameters of the linear regression line (slope, intercept) and the intersection point with MCL
            """
        
        data = self.data
        # finds the intersection point of 2 lines given the slopes and y-intercepts
        def line_intersect(m1, b1, m2, b2):
            if m1 == m2:
                print ('The lines are parallel')
                return None
            x = (b2 - b1) / (m1 - m2)
            y = m1 * x + b1
            return x,y

        # Gets appropriate data (well_name and analyte_name)
        query = self.query_data(station_name, analyte_name)
        query = query[query.RESULT > 0]  # drop non-positive values
        n_samples = query.shape[0]

        if(type(query)==int and query == 0):
            return 'No results found for {} and {}'.format(station_name, analyte_name)
        else:   

            test = query.groupby(['COLLECTION_DATE'])[['RESULT']].mean()
            test.index = pd.to_datetime(test.index)

            x = date2num(test.index)
            y = np.log10(test.RESULT)
            ylabel = 'log-Concentration (' + self.get_unit(analyte_name) + ')'
            y = y.rename(ylabel)

            # p, cov = np.polyfit(x, y, 1, cov=True)  # parameters and covariance from of the fit of 1-D polynom.
            # slope, intercept = p

            slope, intercept, rvalue, pvalue, _ = stats.linregress(x, y)

            f = np.poly1d([slope, intercept])  # polynomial function of the fit

            try:
                MCL = self.get_MCL(analyte_name)
                m1, b1 = f # line of best fit
                m2, b2 = 0, MCL # MCL constant

                intersection = line_intersect(m1, b1, m2, b2)
                predicted_MCL_date = num2date(intersection[0]).date()

                return slope, rvalue**2, pvalue, n_samples, predicted_MCL_date

            except Exception as e:
                print(f"{station_name}: Error calculating MCL for {analyte_name}: {e}")
                # intersection = None

            
    def plot_attribute_map(self, df, value_col, cbar_label, fig_title, vmin, vmax,
                           lon_col='longitude', lat_col='latitude', label_col='station_id',
                           crs="EPSG:4326", cmap='jet', flag_stations=False,
                           basemap_provider=cx.providers.OpenStreetMap.Mapnik,
                           output_path=None):

    
        # Create GeoDataFrame
        gdf = gpd.GeoDataFrame(df.copy(),
                               geometry=gpd.points_from_xy(df[lon_col], df[lat_col]),
                               crs=crs)

        # Convert to Web Mercator for basemap
        gdf_web = gdf.to_crs(epsg=3857)    
        

        # Plot colored points
        fig, ax = plt.subplots(figsize=(8,12))
        sc1 = gdf_web.plot(ax=ax, column=value_col, cmap=cmap, vmin=vmin, vmax=vmax,
                           markersize=100, edgecolor='black', alpha=0.9)
        
        if flag_stations:
            gdf_web_sub = gdf_web[(gdf_web['slope']>0)|(gdf_web['pvalue']>0.1)]  # Filter out rows with NaN in value_col
            sc2 = gdf_web_sub.plot(ax=ax, column=value_col, color="white", edgecolor="black",
                                   vmin=vmin, vmax=vmax, markersize=102, linewidth=1)

        # Add labels
        for _, row in gdf_web.iterrows():
            ax.text(row.geometry.x + 10, row.geometry.y + 9, row[label_col],
                    fontsize=9, ha='left', va='bottom', color='black')

        # Add basemap
        cx.add_basemap(ax, source=basemap_provider)

        # Normalize color range
        norm = Normalize(vmin=vmin, vmax=vmax)

        # Add side colorbar
        sm = ScalarMappable(norm=norm, cmap=cmap)
        sm.set_array([])
        cbar = fig.colorbar(sm, ax=ax, orientation='vertical', shrink=0.6, pad=0.1)
        cbar.set_label(cbar_label, fontsize=10)

        # Style the plot
        ax.set_axis_off()
        ax.set_title(fig_title, fontsize=14)
        
        # Save or show
        if output_path:
            plt.savefig(output_path, dpi=300, bbox_inches='tight')
        else:
            plt.show()


    def lstm_anomaly_detection(self, data, scaler=MinMaxScaler(), window_size=30,
                               train_portion=0.5, val_portion=0.1, anomaly_percentile=0.99):
        """
        Performs LSTM-based anomaly detection on time series data.

        Parameters:
            data (pd.Series): Time series data for anomaly detection.
            window_size (int): Size of the sliding window for LSTM input.
            scaler (sklearn Scaler): Scaler object for data normalization.
            train_portion (float): Proportion of data for training.
            val_portion (float): Proportion of data for validation.
            anomaly_percentile (float): Percentile threshold for anomaly detection.
        Returns:
            results_df (pd.DataFrame): DataFrame with columns ['date', 'y_test', 'y_pred', 'anomaly_score', 'is_anomaly'].
            threshold (float): Anomaly detection threshold.
        """

        # helper function: sliding window generator
        def make_sliding_windows(data, timesteps=30):
            X, y = [], []
            for i in range(len(data)-timesteps):
                X.append(data[i:i+timesteps])
                y.append(data[i+timesteps])
            return np.array(X), np.array(y)
        
        # Pre-process data
        data_scaled = scaler.fit_transform(data[['RESULT']])    # scale data
        X, y = make_sliding_windows(data_scaled, window_size)

        # training/testing/validation data splitting
        n = len(X)
        train_size = int(train_portion*n)
        val_size   = int(val_portion*n)

        X_train, y_train = X[:train_size], y[:train_size]
        X_val,   y_val   = X[train_size:train_size+val_size], y[train_size:train_size+val_size]
        X_test,  y_test  = X[train_size+val_size:], y[train_size+val_size:]

        # Build LSTM predictor
        model = Sequential([
            Input(shape=(window_size, 1)),
            LSTM(64, activation='relu'),
            Dense(1)
        ])
        model.compile(optimizer='adam', loss='mae')

        # Train the model
        history = model.fit(X_train, y_train, validation_data=(X_val,y_val),
                            epochs=100, batch_size=32,
                            callbacks=[tf.keras.callbacks.EarlyStopping(patience=5, restore_best_weights=True)])
        
        # Evaluate and anomaly detection
        y_pred = model.predict(X_test)
        y_pred_inv = scaler.inverse_transform(y_pred)
        y_test_inv = scaler.inverse_transform(y_test)

        errors = np.abs(y_test_inv.flatten() - y_pred_inv.flatten())

        # Threshold (e.g., default 99th percentile of training error)
        train_pred = model.predict(X_train)
        train_errors = np.abs(scaler.inverse_transform(y_train) - scaler.inverse_transform(train_pred))
        threshold = np.quantile(train_errors, anomaly_percentile)

        # anomalies = errors > threshold
        # dates = data["COLLECTION_DATE"].iloc[train_size+val_size+window_size:].values

        # Combine training and validation results for completeness
        train_pred_inv = scaler.inverse_transform(model.predict(X_train))  
        val_pred_inv = scaler.inverse_transform(model.predict(X_val))

        y_inv = scaler.inverse_transform(y)
        y_pred_all = np.concatenate([train_pred_inv, val_pred_inv, y_pred_inv], axis=0)
        errors_all = np.abs(y_inv.flatten() - y_pred_all.flatten())
        anomalies = errors_all > threshold
        dates = data["COLLECTION_DATE"].iloc[window_size:].values

        results_df = pd.DataFrame({'date': dates,
                                   'y': y_inv.flatten(),
                                   'y_pred': y_pred_all.flatten(),
                                   'anomaly_score': errors_all,
                                   'is_anomaly': anomalies,
                                   'data_split': ['train']*len(train_pred_inv) + ['val']*len(val_pred_inv) + ['test']*len(y_pred_inv)
        })

        return results_df, threshold



    def plot_anomaly_detection_results(self, results_df, threshold, station_name, analyte_name, ylabels, output_dir=None):
        """
        Plots the results of LSTM-based anomaly detection.

        Parameters:
            results_df (pd.DataFrame): DataFrame with columns ['date', 'y', 'y_pred', 'anomaly_score', 'is_anomaly'].
            threshold (float): Anomaly detection threshold.
            station_name (str): Name of the station.
            analyte_name (str): Name of the analyte.
            output_dir (str): Directory to save the plot. If None, displays the plot.
        """

        fig, axes = plt.subplots(figsize=(12,6), nrows=2, sharex=True)
        fontsizes = {"title": 18, "label": 16, "ticks": 15, "legend": 15}


        # Plot actual and predicted values
        axes[0].plot(results_df['date'], results_df['y'], label='Sensor Measurements', color='blue')
        axes[0].plot(results_df['date'], results_df['y_pred'], label='LSTM Predictions', color='orange')

        # Highlight anomalies
        anomalies = results_df[results_df['is_anomaly']]
        axes[0].scatter(anomalies['date'], anomalies['y'], color='red', marker='x', alpha=0.7, label='Anomalies', zorder=5)

        # Add a vertical line to separate training, validation, and test sets
        train_end_date = results_df[results_df['data_split'] == 'train']['date'].max()
        val_end_date = results_df[results_df['data_split'] == 'val']['date'].max()
        axes[0].axvline(x=train_end_date, color='green', linestyle='--', label='Train/Val Split')
        axes[0].axvline(x=val_end_date, color='purple', linestyle='--', label='Val/Test Split')

        # Styling
        axes[0].set_title(f'LSTM-Based Anomaly Detection for {analyte_name} at {station_name}', fontsize=fontsizes["title"])
        axes[0].set_ylabel(ylabels[0], fontsize=fontsizes["label"])
        axes[0].tick_params(axis='both', labelsize=fontsizes["ticks"])
        axes[0].legend(loc='upper left', fontsize=fontsizes["legend"])
        axes[0].grid(linestyle='--', alpha=0.7)

        # Plot errors (anomaly scores) and the threshold line
        axes[1].plot(results_df['date'], results_df['anomaly_score'], label='Absolute Error', color='blue')
        axes[1].axhline(y=threshold, color='red', linestyle='--', label='Anomaly Threshold')
        axes[1].scatter(anomalies['date'], anomalies['anomaly_score'], color='red', marker='x', alpha=0.7, label='Anomalies', zorder=5)

        # Add a vertical line to separate training, validation, and test sets
        axes[1].axvline(x=train_end_date, color='green', linestyle='--', label='Train/Val Split')
        axes[1].axvline(x=val_end_date, color='purple', linestyle='--', label='Val/Test Split')

        # Styling
        axes[1].set_xlabel('Date', fontsize=fontsizes["label"])
        axes[1].set_ylabel(ylabels[1], fontsize=fontsizes["label"])
        axes[1].tick_params(axis='both', labelsize=fontsizes["ticks"])
        axes[1].legend(loc='upper left', fontsize=fontsizes["legend"])
        axes[1].grid(linestyle='--', alpha=0.7)

        plt.tight_layout()

        # Save or show
        if output_dir:
            figure_name = f"anomaly_detection_{station_name}_{analyte_name.replace(' ', '_').lower()}.png"
            plt.savefig(os.path.join(output_dir, figure_name), dpi=300, bbox_inches='tight')
        else:
            plt.show()
