# -*- coding: utf-8 -*-
"""
author: b.oerlemans@vu.nl
modified by: aniket.mita@vub.be

This script processes raw linescan files (csv / excel) from Agilent 8900 ICP-MS Triple Quad and per file:
- Identifies sections with background signals (blanks)
- Applies nearest background correction to linescan measurements
- Identifies standards and creates calibration from standards
- Applies calibration to samples

Required:
- raw_files: linescan file from agilent, provided with manually located standard/sample linescans
- sample_list_file: list of samples to look up in the MS-data
- standards_accepted: file with standards and accepted values used for calibration.

"""

#%%
import pandas as pd
import numpy as np
import os
import matplotlib as mpl
from matplotlib import pyplot as plt
from datetime import datetime

pd.options.mode.copy_on_write = True
os.chdir(filepath) # Change to local dir where files are located

#List of samples
sample_list_file = 'samples.xlsx' # file with samples to look up in MS-data
sample_list = pd.read_excel(sample_list_file)
sample_names = sample_list['sample'].tolist()
samples_dict = {} # as of yet empty dict to hold linescan-data per sample
foundtotal = [] # as of yet empty list to hold samples found in raw datafiles                              

# load raw data
raw_files = ['20241129_Aniket_Oyster_afterZerror_FINAL2.xlsx'] # list of raw MS-files to process
for session, file in enumerate(raw_files):

    data_raw = pd.read_excel(file, header=3)
    data = data_raw.rename(columns={
        'Time [Sec]': 'time', 
        'Unnamed: 20': 'sample', # column with manually located sample or standard identifiers
        'U238': 'U',
        'Th232': 'Th',
        'Pb208': 'Pb',
        'Sm147': 'Sm',
        'Nd146': 'Nd',
        'Pr141': 'Pr',
        'Ce140': 'Ce',
        'La139': 'La',
        'Ba138': 'Ba',
        'Sr86': 'Sr',
        'Zn66': 'Zn',
        'Cu63': 'Cu',
        'Ni60': 'Ni',
        'Fe56': 'Fe',
        'Mn55': 'Mn',
        'Ti47': 'Ti',
        'Ca43': 'Ca',
        'Mg24': 'Mg',
        'Na23': 'Na'}) # Columns renamedd for easier scripting below


    # Add blank names to data
    last_time = data.at[0, 'sample']
    max_time_gap = 1 # define max allowable time gap (s) between two measurement to be still considered the same blank
    j = 1 # start of blank counter
    k = 1 # max. no of measurments in a blank - counter
    l = 1 # margin between sample and blank - counter

    for i, row in data.iterrows():
        if pd.isna(row['sample']):
            if row['time'] - last_time > max_time_gap: # finding new blank
                j = j + 1 # increase the number of the blank by 1
                k = 1 # reset k when a new blank is found
                l = 1 # reset l when a new blank is found
            else:
                l = l + 1 # increase the margin between sample/standard and blank by 1
            if (k < 100) & (l > 5): # maximum consecutive measurements that constitute a blank AND margin is met
                data.loc[i, 'sample'] = 'blank_' + str(j) 
                k = k + 1 
            last_time = row['time']
   
    # Make grouped blank-data and sample-data (including standards)
    blanks = data[data['sample'].str.contains('blank', na=False)]
    # remove outliers somehow? https://stackoverflow.com/questions/35480109/remove-outliers-before-aggregate-in-python-pandas/35484620
    blanks_means = blanks.groupby('sample').agg(['mean', 'std', 'count'])
    print('\n', blanks_means.shape[0], 'blanks determined in', file)

    samples = data[data['sample'].str.contains('blank') == False]
    samples_means = samples.groupby('sample', as_index=False).mean()

    # find correction factor per sample per element
    elements = ['U', 'Th', 'Pb', 'Sm', 'Nd', 'Pr', 'Ce', 'La', 'Ba', 'Sr', 'Zn', 'Cu', 'Ni', 'Fe', 'Mn', 'Ti', 'Mg', 'Na'] # list elements for which to apply background correction and calibration. Only listed relevant ones to reduce calculation times.
    elements = ['Ca'] + elements # Ca always necessary as other elements expressed as ratio against Ca (shell-specific)
    bl_corr_factors = [str(element)+'_bl_corr_fact' for element in elements]

    for i, row in samples_means.iterrows(): # more efficient to iterate over means than over entire dataset
        # calculate diff between sample mean time and each blank mean time. Differentiate between after and before (plus/min). Gets overwritten for each sample mean.
        blanks_means['t_diff_plus'] = blanks_means.time['mean'] - row['time']
        blanks_means['t_diff_min'] = -1 * (blanks_means.time['mean'] - row['time'])
        blanks_means['t_diff_plus'] = blanks_means['t_diff_plus'].mask(blanks_means['t_diff_plus'] < 0, np.nan)
        blanks_means['t_diff_min'] = blanks_means['t_diff_min'].mask(blanks_means['t_diff_min'] < 0, np.nan)
        # find blank elemental concentration for nearest two blanks for each sample measurement series
        for j, element in enumerate(elements):
            element_blank_plus = blanks_means.loc[blanks_means['t_diff_plus'].idxmin(), (element, 'mean')]
            element_blank_min = blanks_means.loc[blanks_means['t_diff_min'].idxmin(), (element, 'mean')]
            element_blank_avg = (element_blank_plus + element_blank_min) / 2
            samples_means.at[i, bl_corr_factors[j]] = element_blank_avg
    blanks_means.drop(columns=['t_diff_min','t_diff_plus']) # no longer needed

    # apply per-measurement per-element blank correction (both to samples and standards)
    samples_means = samples_means[['sample'] + bl_corr_factors]
    samples = pd.merge(samples, samples_means, left_on='sample', right_on='sample', how='left')

    for i, row in samples.iterrows():
        for j, element in enumerate(elements):
            samples.at[i, str(element)+'_blank_corr'] = row[element] - row[bl_corr_factors[j]]

    # load accepted values for standards and identifty standards from measurements
    standards_accepted = pd.read_excel('la-icpms_standards.xlsx')
    standards_accepted = standards_accepted.assign(
        U_m = standards_accepted['U'] / 238.0289,
        Th_m = standards_accepted['Th'] / 232.037,
        Pb_m = standards_accepted['Pb'] / 207.2,
        Sm_m = standards_accepted['Sm'] / 150.4,
        Nd_m = standards_accepted['Nd'] / 144.2,
        Pr_m = standards_accepted['Pr'] / 140.9,
        Ce_m = standards_accepted['Ce'] / 140.1,
        La_m = standards_accepted['La'] / 138.9,
        Ba_m = standards_accepted['Ba'] / 137.327,
        Sr_m = standards_accepted['Sr'] / 87.62,
        Zn_m = standards_accepted['Zn'] / 65.38,
        Cu_m = standards_accepted['Cu'] / 63.55,
        Ni_m = standards_accepted['Ni'] / 58.69,
        Fe_m = standards_accepted['Fe'] / 55.845,
        Mn_m = standards_accepted['Mn'] / 54.9380,
        Ti_m = standards_accepted['Ti'] / 47.87,
        Ca_m = standards_accepted['Ca'] / 40.078,
        Mg_m = standards_accepted['Mg'] / 24.305,
        Na_m = standards_accepted['Na'] / 22.9897) # convert massconcentrations to molarities

    
    standard_ids = standards_accepted['identifier'].tolist()
    pattern = r'(' + '|'.join(standard_ids) + ')'
    samples['standard'] = samples['sample'].str.extract(pattern, expand=False) 
    standards = samples[samples['standard'].notnull()]
   
    standards_means = standards.drop(columns=['sample']).groupby('standard', as_index=False).mean()
        
    # Couple measured intensities to known molar abundances
    standards_means = pd.merge(standards_means, standards_accepted, left_on='standard', right_on='identifier', how='left', suffixes=('_int', '_acc'))
    print('\n', standards_means.shape[0], 'standards used for calibration')

    # Create calibrations of element/Ca_intensities against accepted element/Ca_concentrations for this session
    for element in elements[1:]: # skip Ca/Ca 
        for i, row in standards_means.iterrows():
            standards_means.at[i, str(element)+'/Ca_m'] = row[str(element)+'_m']/row['Ca_m']
            standards_means.at[i, str(element)+'/Ca_int'] = row[str(element)+'_blank_corr']/row['Ca_blank_corr']
        x = standards_means[str(element)+'/Ca_int']
        x = np.vstack(x)
        y = standards_means[str(element)+'/Ca_m'] 
        slope, _, _, _ = np.linalg.lstsq(x, y, rcond=None)
        fig, axs = plt.subplots(figsize=(5,2)) 
        axs.scatter(standards_means[str(element)+'/Ca_int'], standards_means[str(element)+'/Ca_m'])
        axs.plot(standards_means[str(element)+'/Ca_int'], standards_means[str(element)+'/Ca_int']*slope)
        plt.xlabel(str(element)+'/Ca intensity')
        plt.ylabel(str(element)+'/Ca (mol/mol)')
        plt.title(element + ' calibrations for ' + str(file))
        plt.show()
        print('Calibration for session:', file)
        print(str(element)+'/Ca (mol/mol) = ', round(slope[0], 6), '* ', str(element)+'/Ca intensity')

        # Apply calibration to samples
        samples[str(element)+'/Ca_m'] = (samples[str(element)+'_blank_corr']/samples['Ca_blank_corr'])*slope

    # Loop over raw data file for every sample name and place per-sample-dataframe in sample_dict
    found = []
    for i, sample in enumerate(sample_names):
        df = samples[samples['sample'] == sample]
        if not df.empty:
            found.append(sample)
            foundtotal.append(sample)
            df['dist'] = (df['time'] - df['time'].min()) * 20 # convert time to relative distance from start linescan using speed = 20 um/s
            # find shell boundaries
            Ca_max = df['Ca_blank_corr'].max()
            Ca_int_90 = df.nlargest(int(0.5 * df.shape[0]), 'Ca_blank_corr')['Ca_blank_corr'].mean() # average of 90% largest Ca intensities
            Ca_int_tresh = Ca_int_90 * 0.7 # treshold for shell determination set at 70% of Ca_int_90
            start = df[df['Ca_blank_corr'] > Ca_int_tresh]['time'].idxmin()
            end = df[df['Ca_blank_corr'] > Ca_int_tresh]['time'].idxmax()
            df['shell'] = np.where(np.logical_and(start <= df.index, df.index <= end), True, False) # Boolean value can be used to clip data
            # add sample to dict
            samples_dict[sample] = df
    if found:
        print('\nSamples', ', '.join(found), 'found in', file)
    if not found:
        print('No samples from', sample_list_file, 'found in', file)

notfound = []
for name in sample_names:
    if name not in foundtotal:
        notfound.append(name)

print('\nFiles', ', '.join(raw_files), 'loaded.')
if foundtotal:
    print('\nSamples', ', '.join(foundtotal), 'found in loaded files.')
if notfound:
    print('\nSamples', ', '.join(notfound), 'not found in loaded files.')
    
    
with pd.ExcelWriter('processed_samples_2.xlsx') as writer:  # Correct: No indentation
    for sample, df in samples_dict.items():      # Correct: Indented (inside 'with')
        df.to_excel(writer, sheet_name=sample)  # Correct: Indented (inside 'with')
    
        # Code that comes *after* the Excel writing should be at the same level as the 'for' loop
        print("Finished processing.")

  