# Calculate Q statistic from MC samples generated by mc.py 

# Usage: same as mc.py.  

import numpy as np
import pyfits
import sys,os


# obtain the files with the event data. We only need them to have the counts

datadir='montecarlo/'+sys.argv[1]+'/'

file_list=[]

# galactic latitude cut
galbcut = np.int(sys.argv[2])

for file in [doc for doc in os.listdir(datadir) if 
        doc.startswith('northsouthbgt'+sys.argv[2])]:
    file_list.append(file)
    
file_list=np.sort(file_list)

# This is the angular region that we are going to consider 1-20 deg.
angstep=1
minangle=1
maxangle=20

region=np.linspace(minangle,maxangle,(maxangle-minangle+1.)/angstep)

# we loop over the different combinations of E1 and E2

import matplotlib
matplotlib.use('pdf')

import matplotlib.pyplot as mp
fig, ax = mp.subplots(nrows=2,ncols=3,sharex='col')


for i in range(len(file_list)):
    qmc=np.load(datadir+file_list[i])
    qav=qmc[:,0].mean(0)
    qstd=qmc[:,0].std(0)
   
    qn=qmc[:,1].mean(0)
    qstdn=qmc[:,1].std(0)
 
    qs=qmc[:,2].mean(0)
    qstds=qmc[:,2].std(0)



    # We now have all the values of q for this particular choice of E1,E2
    # Lets plot them
    ax[i/3,i%3].set_xlim((0,21))
    ax[i/3,i%3].plot([0,21],[0,0],'k-')
    ax[i/3,i%3].errorbar(region,qav,yerr=qstd,fmt='o',ecolor='g',)
    ax[i/3,i%3].errorbar(region,qn,yerr=qstdn,fmt='>',ecolor='r',)
    ax[i/3,i%3].errorbar(region,qs,yerr=qstds,fmt='<',ecolor='b',)
    ax[i/3,i%3].set_title(r'$(E_1,E_2) = $'+'('+file_list[i][5:7]+
            ', '+file_list[i][8:10]+')')


# Fine-tune figure
mp.setp([a.get_xticklabels() for a in ax[0,:]], visible=False)
fig.suptitle(sys.argv[1]+' Gal latidude > '+str(galbcut),fontsize=22)
fig.savefig(datadir+'/northsouthbgt'+str(galbcut)+sys.argv[1]+'.pdf')