#!/usr/bin/env python
#

"""
This file is part of ESA's XMM-Newton Scientific Analysis System (SAS).
SAS is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

SAS is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with SAS.  If not, see <http://www.gnu.org/licenses/>.


Name    = 'mergeEPICflareBkgTS'
VERSION = '1.1';
author  = 'Jose Vicente Perea';
date    = '2021-02-19';


ChangeLog
=========

Version 1.1 - 2021-02-19 (JVP)
------------

+ Version output string fixed

Version 1.0 - 2020-12-02 (JVP)
------------

+ Initial version
  Task to merge EPIC flare background time series
  from different exposures for the same instrument

"""

__version__   = "1.1"
__author__    = "Jose V. Perea"

__copyright__ = "ESA (C) 2000-2020"
__license__   = "GPL"


version_num   = __version__


###############################
#  Input files
#
import sys, os, argparse

parser = argparse.ArgumentParser(description='Merge EPIC flare background timeseries',
                         formatter_class=argparse.ArgumentDefaultsHelpFormatter)

parser.add_argument('--version', '-v', action='version', version=f'%(prog)s (%(prog)s-{version_num}) [-]')

parser.add_argument('--ratesFiles', help='Input EPIC flare bkg timeseries files', default='Exp1FBKTSR.fits,Exp2FBKTSR.fits')
parser.add_argument('--mergedRates', help='Output EPIC flare bkg timeseries files', default='X000FBKTSR.fits')

args = parser.parse_args()


files = []
for i in sorted(args.ratesFiles.split(',')):
    files.append(i)

for f in files:
    if not os.path.isfile(f):
       print(f'Some of the input files does not exist: {f}. Exiting')
       sys.exit(1)

outrates  = args.mergedRates

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

from astropy.table import Table, vstack
from astropy.io import fits

instrument = []
all   = {}
times = {}
for file in files:
    
    with fits.open( file ) as hdul:
        h0 = hdul[0].header
        r  = hdul['RATE'].data
        h  = hdul['RATE'].header
        
        inst = h['INSTRUME']
        exp  = h['EXPIDSTR']
        flare_cut = h['FLCUTTHR']
            
        instrument.append(inst)
        if instrument[0] != instrument[-1]:
            print('You are trying to merge timeseries from different instruments. Exiting.')
            sys.exit(1)
        
        bs   = h['BACKSCAL']
        texp = h['EXPOSURE']

        tdel = h['TIMEDEL']
        print(f'TIMEDEL : {tdel}')

        
        tsr = Table(r)
 
        ntsr = tsr
        ntsr['RATE']  = r.RATE/bs
        ntsr['ERROR'] = r.ERROR/bs

        all[exp]   = {'bs': bs, 'texp': texp, 'table': ntsr, 't0': ntsr['TIME'][0], 'flare_cut': flare_cut}
        times[exp] = ntsr['TIME'][0]

 
bs   = sum([ s['bs']*s['texp'] for s in all.values() ]) / sum([ s['texp'] for s in all.values() ])
Texp = sum([ s['texp'] for s in all.values() ])

for exp,vals in all.items():
    bsi   = vals['bs']
    texpi = vals['texp']
    print(f'{inst}\t{exp}\tbckscal: {bsi:.2f}\ttexp: {texpi:.2f}\t\tGlobal bckscal: {bs:.2f}')


sortedExps = sorted( times, key=times.get )

allTables = [ all[exp]['table'] for exp in sortedExps ]

##
#for exp in sortedExps:
#    T0 = all[exp]['table']['TIME'][0]
#    Tn = all[exp]['table']['TIME'][-1]
#    T = Tn - T0
#    print(f'{exp} --> T0, Tn: {T0}, {Tn}\tTotal : {T}')    
##

allTable  = vstack( allTables )    # Optional option : ,metadata_conflicts='silent'

allTable['RATE']  = allTable['RATE']*bs
allTable['ERROR'] = allTable['ERROR']*bs
## print(allTable)
#
TSTART = allTable['TIME'][0]                                    # First and last time values
TSTOP  = allTable['TIME'][-1]
#
# Convert 'Table object' to 'BinTableHDU object'
allTableHDU = fits.table_to_hdu( allTable )

h['TSTART'] = TSTART
h['TSTOP']  = TSTOP

h['BACKSCAL'] = bs
h['EXPOSURE'] = Texp
h.pop('EXP_ID')
h.pop('EXPIDSTR')
h.pop('FLCUTTHR')
allTableHDU.header = h

# Create FITS:
hdul = []
hdu0 = fits.PrimaryHDU(header = h0)
hdul.append(hdu0)
hdul.append(allTableHDU)

## outrates = 'merged-fbktsr.fits'
hdul = fits.HDUList( hdul )
hdul.writeto(outrates, overwrite = True )


sys.exit(0)