# Read the Fermi datafiles generated by gtmktime and perform additional cuts
# Here I select galactic latitude |b| > 50 deg

# Create new fits files with the filtered photon data

import numpy as np
import pyfits

# I commented the line below because I will define a function
# gtmfile = 'gtmktime10_20GeVp130.fits'

def bcut(gtmfile,outfile):
    fd = pyfits.open(gtmfile)

    # first HDU is a generic Fermi Header
    # data is on second HDU
    fdata = fd[1].data

    # we are filtering column 'b', galactic latitude.
    # for a list of column names define
    # fcols = fd[1].columns
    # and see the output of fcols.names

    # Get the photons that satisfy abs(b)>50deg

    index = np.where(abs(fdata.field('b'))>50)

    # We are now going to create a new fits file with E, RA, DEC, L, B, t, Id

    cole = pyfits.Column(name='ENERGY',format='E',array=fdata.field('energy')[index])
    cora = pyfits.Column(name='RA',format='E',array=fdata.field('ra')[index])
    codec = pyfits.Column(name='DEC',format='E',array=fdata.field('dec')[index])
    coll = pyfits.Column(name='L',format='E',array=fdata.field('l')[index])
    colb = pyfits.Column(name='B',format='E',array=fdata.field('b')[index])
    colt = pyfits.Column(name='TIME',format='D',array=fdata.field('time')[index])
    colid = pyfits.Column(name='EVENT_ID',format='J',array=fdata.field('event_id')[index])

    # create a ColDefs object for all columns
    cols = pyfits.ColDefs([cole,cora,codec,coll,colb,colt,colid])

    # create new binary table HDU
    tbhdu = pyfits.new_table(cols)

    tbhdu.writeto(outfile)

    # return event_id array for comparison
    #return colid.array