# Cut away 3deg regions from sources

# Create new fits files with the filtered photon data

# Usage: we had the data stored in directories under ./data/period, where
# period was e.g. jan14z100 or sep13z100 for data until Jan 2014 with 
# zenith angle cut 100deg
# The script takes one argument: period
# For example the files jan14z100eE1_E2GeVp202.fits generated by bcut.py
# where stored under ./data/jan14z100. Then to generate data with regions
# around the sources cut, we issued
# python nosource.py jan14z100
# the result are the files contained in nosourceoutput.tar.gz
# which are stored under ./sourcecuts/jan14z100
# their filenames have 'nosource' added to the beggining of the input filename

import numpy as np
import pyfits
import sys,os

# 1st HE source catalogue
fermicatalogue='/srv/data/fermi/sources/1fhel/gll_psch_v07.fit'

# 2nd source catalogue
#fermicatalogue='/srv/data/fermi/sources/2yrps/gll_psc_v08.fit'

fsource=pyfits.open(fermicatalogue)

# deg to rad
degrad=np.pi/180.

# Size of mask around each source. This is the cosine.
msize=np.cos(1.5*degrad)

# Create an array of unit vectors pointing in the direction of the sources

srcdata = fsource[1].data
numsrc = srcdata.shape[0]

cb = np.cos(srcdata.field('glat') * degrad)
sb = np.sin(srcdata.field('glat') * degrad)

cl = np.cos(srcdata.field('glon') * degrad)
sl = np.sin(srcdata.field('glon') * degrad)

srcvec=np.column_stack((cl*cb,sl*cb,sb))

# obtain the files with the event data

datadir='data/'+sys.argv[1]+'/'
outdir ='sourcecuts/'+sys.argv[1]+'/'


file_list=[]


for file in [doc for doc in os.listdir(datadir) if doc.startswith(sys.argv[1])]:
    infile=datadir+file
    evfile=pyfits.open(infile)
    evdata=evfile[1].data
    ecb = np.cos(evdata.field('b') * degrad)
    esb = np.sin(evdata.field('b') * degrad)
    ecl = np.cos(evdata.field('l') * degrad)
    esl = np.sin(evdata.field('l') * degrad)
    datavec=np.column_stack((ecl*ecb,esl*ecb,esb))
    angdist=np.zeros(len(ecb))
    # store the maximum value of the cos(source-event) for each event
    for i in range(len(ecb)):
        angdist[i]=np.dot(srcvec,datavec[i]).max()
    # keep events that have cos < cos_max
    index = np.where(angdist < msize)

    # now save only these events
    # 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=evdata.field('energy')[index])
    cora = pyfits.Column(name='RA',format='E',array=evdata.field('ra')[index])
    codec = pyfits.Column(name='DEC',format='E',array=evdata.field('dec')[index])
    coll = pyfits.Column(name='L',format='E',array=evdata.field('l')[index])
    colb = pyfits.Column(name='B',format='E',array=evdata.field('b')[index])
    colt = pyfits.Column(name='TIME',format='D',array=evdata.field('time')[index])
    colid = pyfits.Column(name='EVENT_ID',format='J',array=evdata.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(outdir+'nosource'+file)

    # print file name and # events  
    print 'nosource'+file, len(index[0])