# Plot_nuebar_e_event_array.py
# Needs current_istope_event_calc_run.txt and nuebar_e_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

# setup
radius0 = 6.5
df_factor = 0.9/0.8  #   80%->90% DF for ES only
veto_livetime = 0.624  # veto livetime
radius = 5.0  
evis_min = 3.0
sinsqthw = 0.238  # option to change sinsqthw

print 'radius0,radius,df_factor,veto_livetime,evis_min,sinsqthw: '
print  radius0,radius,df_factor,veto_livetime,evis_min,sinsqthw

# 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("nuebar_e_event_array.dat","r") as f:
    contents = f.readlines()   # Read the file into a list of lines.
line = contents[0].strip()
fields = line.split()
nev = len(contents)-1  
sinsqthw0 = float(fields[1])
print 'nev, sinsqthw0 = ',nev," ",sinsqthw0

# option to change the input distribution
eminus0 = -0.5-sinsqthw0
eplus0 = -sinsqthw0
eminus = -0.5-sinsqthw
eplus = -sinsqthw

dat_array = np.zeros((5,nev))
enu_array = np.zeros(nev)  # true neutrino energy
evis_array = np.zeros(nev)  #  electron energy
levt_array = np.zeros(nev)
xsec_array = np.zeros(nev)
wt_array = np.zeros(nev)

for i in range(1,nev+1):
    iev = i-1 # Adjust for the first line of the file
    # dat_array[k,iev] where k = enu,levt,e_elect,xsec_nuebar_e,wtnuebar
    fields = contents[i].split()
    for j in range(0,5):
        dat_array[j,iev] = float(fields[j])
    # correct weight for new number of years
    dat_array[3,iev] *= nyrs/nyrs_dat_array

    evis = dat_array[2,iev]
    enu = dat_array[0,iev]
    y = evis/enu
    factor0 = eplus0**2 + eminus0**2*(1-y)**2
    factor = eplus**2 + eminus**2*(1-y)**2 \
        + eminus*eplus*y*0.511/enu   #  This term changes 0.008425 to 0.07960
    # rescale for radius,veto_livetime,df_factor,couplings
    dat_array[3,iev] *= factor/factor0
    dat_array[4,iev] *= veto_livetime*df_factor*(radius/radius0)**3 \
        * factor/factor0

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

    enu_array[iev] = max(dat_array[0,iev]+np.random.normal(0, 1.0, 1)*eres*sqrt(dat_array[0,iev]),0.0)
    levt_array[iev] = dat_array[1,iev]+np.random.normal(0, 1.0, 1)*sqrt(lres**2/dat_array[0,iev]+0.41**2)
    evis_array[iev] = max(dat_array[2,iev]+np.random.normal(0, 1.0, 1)*eres*sqrt(dat_array[2,iev]),0.0)
    xsec_array[iev] = dat_array[3,iev]
    wt_array[iev] = dat_array[4,iev]
    if evis_array[iev] < evis_min:
        wt_array[iev] = 0.0
    
    if iev % 100000 == 0:
        print iev,enu_array[iev],levt_array[iev],evis_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 = 40
pl.figure(1)

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

ax = pl.subplot(2,2,2)
title = 'nuebar-e events vs Path'
histog2w(levt_array,wt_array,title,nbins,5.,25.)

ax = pl.subplot(2,2,3)
title = 'nuebar-e events vs E_Elect'
histog2w(evis_array,wt_array,title,nbins,0.,20.)

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