# Plot_IBD_event_array.py
# Needs current_istope_event_calc_run.txt and IBD_event_array.dat files

import numpy as np
import matplotlib.pyplot as pl
from iminuit import Minuit 
from math import sqrt,isnan
from random import random

global nev
global dat_array

# Number of years for this run
nyrs = 5.0*0.8

# print out current run parameters
with open("current_istope_event_calc_run.txt","r") as f:
    contents = f.readlines()   # Read the file into a list of lines.
    
for i in range(0,5):
    print contents[i]

#Number of years for the input data file is second line
line = contents[1].strip()
fields = line.split()
nyrs_dat_array = float(fields[5]) 

print "nyrs_dat_array: {:5.2f} nyrs run: {:5.2f}".format(nyrs_dat_array,nyrs)

# read in data file
with open("IBD_event_array.dat","r") as f:
    contents = f.readlines()   # Read the file into a list of lines.
# The first line contains the number of lines in the rest of the file.
line = contents[0].strip()
fields = line.split()
nev = int(fields[0])

dat_array = np.zeros((4,nev)) # Define 2D array
enu_array = np.zeros(nev)
levt_array = np.zeros(nev)
xsec_array = np.zeros(nev)
wt_array = np.zeros(nev)

# Loop over event lines in data file
for iev in range(0,nev):
    line = contents[iev+1].strip()
    # enu,levt,xsecIBD,n_nuebar_IBD
    fields = line.split()
    for i in range(0,4):
        dat_array[i,iev] = float(fields[i])
    # correct weight for new number of years
    dat_array[3,iev] *= nyrs/nyrs_dat_array

    eres = 0.03 # energy resolution = 3%/sqrt(E)
    lres = 0.12 # length resolution = 12cm/sqrt(E)

    enu_array[iev] = dat_array[0,iev]+np.random.normal(0, 1.0, 1)*eres*sqrt(dat_array[0,iev])
    levt_array[iev] = dat_array[1,iev]+np.random.normal(0, 1.0, 1)*sqrt(lres**2/dat_array[0,iev]+0.41**2)
    xsec_array[iev] = dat_array[2,iev]
    wt_array[iev] = dat_array[3,iev]
    #print iev,enu_array[iev],levt_array[iev],xsec_array[iev],wt_array[iev]

# Unweighted histogram
def histog2(data,title,nbins,xmin0,xmax0):
    delta = (xmax0-xmin0)/1000.0
    delta = 0.
    Xhist,Xedges = np.histogram(data,bins=nbins,range=(xmin0-delta,xmax0+delta))
    x = np.array( [ (Xedges[i]+Xedges[i+1])/2. for i in range(len(Xedges)-1) ] )
    pl.errorbar(x,Xhist,yerr=Xhist**0.5,marker='.',drawstyle='steps-mid',capsize=0)
    pl.xlabel(title, fontsize = 12,labelpad=8)
    pl.ylabel(r'Events', fontsize = 12)
    pl.locator_params(nbins=6)
    nevents = len(data)
#    print 'nevents: ',nevents
    mu = np.mean(data)
    sigma = np.std(data)
    textstr = '$N$=%5.0f \n $\mu$=%.3f \n $\sigma=$%.3f'% (nevents,mu,sigma)
    props = dict(boxstyle='round', facecolor='white', alpha = 0.5)
    ax.text(0.80,0.95,textstr,transform=ax.transAxes, fontsize=10,verticalalignment='top',bbox=props)

# Weighted histogram
def histog2w(data,wts,title,nbins,xmin0,xmax0):
    delta = (xmax0-xmin0)/1000.0
    delta = 0.
    Xhist,Xedges = np.histogram(data,weights=wts,bins=nbins,range=(xmin0-delta,xmax0+delta))
    x = np.array( [ (Xedges[i]+Xedges[i+1])/2. for i in range(len(Xedges)-1) ] )
    pl.errorbar(x,Xhist,yerr=Xhist**0.5,marker='.',drawstyle='steps-mid',capsize=0)
    pl.xlabel(title, fontsize = 12,labelpad=8)
    pl.ylabel(r'Events', fontsize = 12)
    pl.locator_params(nbins=6)
    #nevents = len(data)
    nevents = np.sum(wts)
    #print 'nevents: ',nevents
    mu = np.mean(data)
    sigma = np.std(data)
    textstr = '$N$=%5.0f \n $\mu$=%.4g \n $\sigma=$%.4g'% (nevents,mu,sigma)
    props = dict(boxstyle='round', facecolor='white', alpha = 0.5)
    ax.text(0.80,0.95,textstr,transform=ax.transAxes, fontsize=10,verticalalignment='top',bbox=props)

nbins = 30
pl.figure(1)

ax = pl.subplot(2,1,1)
title = 'IBD events vs Enu'
histog2w(enu_array,wt_array,title,nbins,0.,20.)

ax = pl.subplot(2,1,2)
title = 'IBD events vs Path'
histog2w(levt_array,wt_array,title,nbins,5.,30.)

A_txt = 'IBD Event Distributions'
pl.suptitle(A_txt, fontsize = 12)
pl.tight_layout(pad=2.5)
pl.show()