#!/usr/bin/env python3

# authors M.Appleby T. Mason J Beale M. Kepla 

"""
correlate statistics from bsdata, in aprticular JFJoch statistics to events or other data
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.cm import get_cmap
from scipy import stats
from scipy.optimize import curve_fit
from scipy.signal import peak_widths, find_peaks
from scipy import asarray as ar,exp
from math import sqrt, pi
import argparse
import os

from sfdata import SFDataFiles, SFScanInfo

def load_data(directory): # for data taken by slic
    num_of_acqs=len([name for name in os.listdir(os.path.join(directory,'data')) if "BSDATA" in name])
    #print(num_of_acqs)
    #bkg_estimate_data = []
    #bkg_estimate_pulse_id = []
    events_all = []
    bkgs_estimates_all = []
    pulse_ids_all = []
    #fig, ax1 = plt.subplots()
    #fig, axs = plt.subplots(2, 1, figsize=(6, 6))
    #ax2 = ax1.twinx() 
    num_lines = 28//2
    cmap = get_cmap('plasma') 
    colors = [cmap(i / num_lines) for i in range(num_lines)]
    for acquisition in range(1,num_of_acqs+1):
        events = []
        bkgs_estimates = []
        pulse_ids = []
        beam_ints=[]
        
        print(f"Processing acquisition: {acquisition}")

        filename = os.path.join(directory,'data','acq{:04}.BSDATA.h5'.format(acquisition))
        try:
            with SFDataFiles(filename) as bsdata:
            
                bkg_estimate = bsdata[f"JF17T16V01j:bkg_estimate"].data
                pulseids_JF = bsdata[f"JF17T16V01j:bkg_estimate"].pids
                pulseids_BS = bsdata[f"SAR-CVME-TIFALL6:EvtSet"].pids
                evt_set=bsdata[f"SAR-CVME-TIFALL6:EvtSet"].data
                az_int_profil=bsdata[f"JF17T16V01j:az_int_profile"].data
                beam_int=bsdata[f"SAROP31-PBPS149:INTENSITY"].data
                #pulseids_BS=bsdata[f"SAROP31-PBPS149:INTENSITY"].pids
                #az_int_profile_ps=bsdata[f"JF17T16V01j:az_int_profile"].pids
                #if az_int_profile_ps.all() == pulseids_JF.all():
                #   print("yes")
                   
                print("finding intersection")
                _inters, ind_JF, ind_BS = np.intersect1d(pulseids_JF, pulseids_BS, return_indices=True)
               
                print("pulling data")
                for idx_JF, idx_BS in zip(ind_JF, ind_BS):
                    pulse_id = pulseids_BS[idx_BS]
                    
                    if pulse_id != pulseids_JF[idx_JF]:
                        print(f"Mismatch in pulse IDs: BS={pulse_id}, JF={pulseids_JF[idx_JF]}")
                        continue 
                        
                    bkg_estimate_plt = bkg_estimate[idx_JF]
                    event=int(bool(evt_set[idx_BS][216]))-acquisition/100
                    
                    #if event == 1:
                    events.append(event)
                    events_all.append(event)
                    bkgs_estimates.append(bkg_estimate_plt)
                    bkgs_estimates_all.append(bkg_estimate_plt)
                    pulse_ids.append(pulse_id-24358602677)
                    pulse_ids_all.append(pulse_id)
                    #beam_ints.append(beam_int[idx_BS])
                    #if idx_JF % 2 == 0:
                    #    axs[0].plot(az_int_profil[idx_JF][:94], label=f'{idx_JF}', color=colors[idx_JF//2])
                    #else:
                    #    axs[1].plot(az_int_profil[idx_JF][:94], label=f'{idx_JF}', color=colors[idx_JF//2])
                    
                    #if idx_JF > 25:
                    #    break
                #axs[0].legend()
                #axs[1].legend()
                #plt.show()
                plt.scatter(events, bkgs_estimates, label=acquisition, color=colors[acquisition//2])
                    
        except FileNotFoundError:
            print(f"File not found: {filename}, skipping...")
        except KeyError as e:
            print(f"Missing key in file {filename}: {e}, skipping...")
        except Exception as e:
            print(f"Unexpected error in file {filename}: {e}")
    #print(events)
    #ax1.scatter(pulse_ids, events, color='orange')
    #ax2.set_ylim(5,18)
    #ax2.scatter(pulse_ids, bkgs_estimates, color='blue')
    df = pd.DataFrame(data={'events':events_all, 'bkg_estimate':bkgs_estimates_all}, index=pulse_ids_all)
    #print(pulse_ids)
    plt.legend()
    print(df.corr(method ='pearson'))
    plt.show()

    return df

#directory="/sf/cristallina/data/p22547/raw/run0040-hewl_droplets_hit_and_return_12x11_1p21s/"
#directory="/sf/cristallina/data/p22547/raw/run0039-hewl_droplets_stop_go_10_10"
#directory="/sf/cristallina/data/p22547/raw/run0081-hewl_droplets_stop-go_10w_10m_NaI_delay_changed"
#directory="/sf/cristallina/data/p22547/raw/run0048-hewl_NaI_9p72s"
#directory="/sf/cristallina/data/p22547/raw/run0058-hewl_NaI_1p32s"
#directory="/sf/cristallina/data/p22547/raw/run0072-hewl_NaI_0p560s"
#directory="/sf/cristallina/data/p22547/raw/run0046-hewl_apo"
#directory="/sf/cristallina/data/p22547/raw/run0082-hewl_droplets_2mylar"
#directory="/sf/cristallina/data/p22547/raw/run0087-hewl_droplets_stopgo_5um_beam_0p05trans"
#directory="/sf/cristallina/data/p22547/raw/run0036-hewl_droplets_stop_go_10_10"
#directory="/sf/cristallina/data/p22547/raw/run0075-hewl_droplets_hit_and_return_6x162_9p72s"
#directory="/sf/cristallina/data/p22547/raw/run0077-hewl_droplets_stop-go_10w_10m_NaI"
directory="/sf/cristallina/data/p22547/raw/run0094-hewl_100mM_NaI_9p72s"
load_data(directory)