# %%
import numpy as np
from astropy.io import fits
from astropy.time import Time

# example lightcurve file
lc_fn="c:/Users/pmakela/OneDrive - NASA/Documents/Okezie/python/lat_LC_20230119.fits"
# corresponding spectrum file
spec_fn="c:/Users/pmakela/OneDrive - NASA/Documents/Okezie/python/lat_spectrum_20230119.fits"

# time conversion function
# MET - mission elapsed time - should be double precision for full accuracy
# Convert Fermi Mission Elapsed time (seconds since 1-Jan-2001 with
# no leap sec) to seconds since 1-Jan-1979 
# PM: MET=157766400 (leap second) goes to year 2005 in IDL function fermi_met2tim.pro
# On xTime time conversion site (heasarc.gsfc.nasa.gov/cgi-bin/Tools/xTime/xTime.pl) in year 2006 
# convert MET to time to seconds from the 1979-01-01 epoch as done in IDL
def fermi_met2tim(times):
    # IDL ospex uses 1979-01-01 as an epoch, here we use Terrestrial Time TT (TT=TAI+32.184 sec)
    epoch_1979_tt=Time('1979-01-01 00:00:00',scale='tt')
    epoch_1979_tt.format='cxcsec'
    epoch_met_tt=Time('2001-01-01 00:00:00',scale='tt')
    epoch_met_tt.format='cxcsec'
    epoch_met_tai=Time('2001-01-01 00:00:00',scale='tai')
    epoch_met_tai.format='cxcsec'
    epoch_met_utc=Time('2001-01-01 00:00:00',scale='utc')
    epoch_met_utc.format='cxcsec'

    times0=Time(times + epoch_met_tt.value,scale='tt',format='cxcsec')
    times0.format='isot'
    tt=times0.utc.value
    t1=Time(tt,scale='tai')
    t2=Time(tt,scale='utc')
    t1.format='cxcsec'
    t2.format='cxcsec'

    # times0 doesn't match what the IDL fermi_met2tim function returns, because
    # fermi_met2tim extracts the number of leap seconds between the MET epoch and tim 
    # (maybe not exactly correct if time shifter over leap second addition time)
    # this should match fermi_met2tim output
    nleap=np.round((epoch_met_tai.value-epoch_met_utc.value) - (t1.value - t2.value))
    tim=Time(times + epoch_met_tt.value - nleap,scale='tt',format='cxcsec')
 
    # return seconds from the 1979-01-01 epoch
    return tim.value - epoch_1979_tt.value

# Time converiosn from seconds from the 1979 epoch to UTC (default format=string)
def fermi_tim2utc(times,format='iso'):
    # IDL ospex uses 1979-01-01 as epoch, here we use Terrestrial Time TT (TT=TAI+32.184 sec)
    epoch_1979_tt=Time('1979-01-01 00:00:00',scale='tt')
    epoch_1979_tt.format='cxcsec'
    tim=Time(times+epoch_1979_tt.value,scale='tt',format='cxcsec')
    tim.format=format
    return tim

# Function to find consecutive sequencies (breaks at gapsize >1 as default)
def consecutive(data, gapsize=1):
    return np.split(data, np.where(np.diff(data) > gapsize)[0]+1)

# read data from the example light curve file
hdul = fits.open(lc_fn)
# hdul.info()
data_rate=hdul['rate'].data
hdul.close()

ts=fermi_met2tim(data_rate['time'])
te=fermi_met2tim( data_rate['time'] + data_rate['timedel'])

expos=data_rate['exposure'] #1-minute exposure values
cts=data_rate['counts'] #1-minute counts

# (2,ntimes) array for times
t2d=np.vstack((ts,te))
lc_dict={'times':t2d,'expos':expos,'cts':cts}

times=lc_dict['times']
twidth=times[1,0]-times[0,0]
#expos is exposure in cm^2 sec.  Our bins are 1-minute, so expos is 60 times effective area.
expos = lc_dict['expos'] / twidth  # now expos is in cm^2

# read data from the corresponding spectrum file
hdul = fits.open(spec_fn)
# hdul.info()
spec=hdul['spectrum'].data
shdr=hdul['spectrum'].header
ebounds=hdul['ebounds'].data
ehdr=hdul['ebounds'].header
hdul.close()

spectrum_count=spec['counts'] #observed counts
exposure=spec['exposure'] #exposure
spectrum_count_errors=np.sqrt(spectrum_count)
livetime=np.array([exposure,]*10).T #expand exposure vector to (ntimes,10) array (10 energy channels)
spectrum_rate=np.nan_to_num(np.divide(spectrum_count,livetime)) #this will set div by zero to 0.0
spectrum_rate_errors=np.nan_to_num(np.divide(spectrum_count_errors,livetime))

spec_stime=fermi_met2tim(spec['tstart'])
spec_etime=fermi_met2tim(spec['tstart'] + spec['telapse'])

spec_ut_edges=np.stack((spec_stime,spec_etime))
spec_file_time=[np.min(spec_ut_edges),np.max(spec_ut_edges)]
spec_ct_edges=np.stack((ebounds['e_min'],ebounds['e_max']))

if times.size != spec_ut_edges.size:
    if times.size > spec_ut_edges.size:
        ii=(spec_ut_edges[0,]<=times[0,]) and (spec_ut_edges[1,]>=times[0,])
        spec_ut_edges=spec_ut_edges[ii]
        spectrum_rate=spectrum_rate[ii]

totrate=spectrum_rate.sum(axis=1)
totcounts=spectrum_count.sum(axis=1)

# some cutoff limits for exposure (not sure if I used excatly these values)
exp_cutoff=26000. #good data points with texp >= exp_cutoff 
errfrac_cutoff=0.8 #good data points with error/tflux < errfrac_cutoff
exp1min_cutoff=650 #code using only this limit for the 1-minute exposure values

# find consecutive sequencies of indeces
jj=(np.arange(0,expos.size))[(expos>exp1min_cutoff)]
consec_ind=consecutive(jj)
consec_ind_len=len(consec_ind)
# create a data point for each sequency (should correspond observations during 1 orbit)
acc=[]
for k in range(consec_ind_len):
    l=consec_ind[k]
    texp=expos[l[expos[l]>exp1min_cutoff]].sum() # total exposure
    ave_exp=(expos[l[expos[l]>exp1min_cutoff]].mean())/1e4 #convert from cm^2 to m^2
    trate=totrate[l[expos[l]>exp1min_cutoff]].sum() # total rate
    tflux=np.nan_to_num(np.divide(trate,texp)) # total flux
    err=np.nan_to_num(np.divide(tflux,np.sqrt(trate*twidth))) # flux error
    ut=[times[0,l[0]],times[1,l[-1]]] # observation period start and end times in seconds from the 1979 epoch
    ut_mean=fermi_tim2utc(np.array(ut).mean()) # midpoint time of observation period as a string
    acc.append({'ut':ut,'ut_mean':ut_mean.value,'texp':texp,'trate':trate,'ave_exp':ave_exp,'tflux':tflux,'error':err})

print(acc[0])
# %%