#!/usr/bin/env python 

from __future__ import print_function

""" This example is a complete analysis, computing the effective area
	from both neutrino and anti-neutrino event files. The output is
	automatically written to a web-page, which has both plots and tables.
"""

# This script shows how to use the event weights to compute event rates
# for a given flux, and the neutrino effective area.

from math import *
import ROOT, aa
aa.i()

ROOT.gStyle.SetOptStat(0)
ROOT.gROOT.SetBatch(True)
ROOT.gStyle.SetTitleOffset(1.5, "xyz")
from util.aautil import table

www = aa.WebFile("effective area", "./aeff_webdir" )

copy_cmd = "scp -r aeff_webdir t61@login.nikhef.nl:~/public_html" # change if you're not t61


www.add( """ This is the effective area of one KM3NeT building block (there will be
			 two building blocks, so multiply by 2). This is at "trigger level" 
			 (which means that all detected neutrinos contribute,
			 and no cuts are applied to remove backgrounds. (so it's optimistic). <br><br>

			 the cos(zenith) is defined such that -1 is a down-going neutrino, while +1 is
			 an upgoing neutrino (through the Earth -- note earth absorption is accounted
			 for and the effect visiable at high E and cos(zenith)=1). <br><br>

			 The tables refer to the average affective area ( = (Anu+Anubar)/2 ). This average
			 can be multiiplied with the total (nu+nubar) flux, if the flux consists of an
			 equal amount of nu and nubar.
			   <br><br>""" )
		 




def fill_histograms() :

	"this functions produces histograms of the Aeff from MC input files"

	dr = "/sps/km3net/repo/mc/atm_neutrino/KM3NeT_-00000001_20171212/v5.1/reco/mcv5.1.genhen_numuCC.sirene.jte.jchain.aashower.NNN.root"

	infiles_nu  = [ dr.replace("NNN",str(i)) for i in range(1,4) ]
	infiles_anu = [ n.replace("numu","anumu") for n in infiles_nu ]

	def fillhist( infiles , name = "h"):

		h = ROOT.TH2D( name, name ,28,1,8,20,-1,1)

		ngen_tot = 0
		for infile in infiles :

			f = ROOT.EventFile( infile )
			ngen_tot += header.ngen()

			for evt in f :

				#require at least one reconstructed track
				if evt.trks.size() == 0 : continue

				# add other quality cuts / upgoing etc here.

				nu = evt.mc_trks[0]
				h.Fill( log10(nu.E), nu.dir.z , evt.w[1]/nu.E )

		h.Scale( 1.0/ngen_tot)
		
		print (" ngen tot = " , ngen_tot )
		return h
		

	h_nu  = fillhist( infiles_nu,  "h_nu"  )
	h_anu = fillhist( infiles_anu, "h_anu" )
	fout  = ROOT.TFile("aeff.root","RECREATE"   )

	h_nu.Write()
	h_anu.Write()
	del fout

	
import os, sys
if 'gen' in sys.argv or os.path.exists("aeff.root") == False :
	fill_histograms()





	
def sindecl_vs_coszenith( lattitude = ( 36 + 16 / 60.0 ) * pi / 180 ): # ARCA: 36deg 16 sec

	"return a histogram of sin(declination) vs cos(zenith)"
	

	sindecl_vs_coszenith = ROOT.TH2D("sindecl_vs_coszenith","sindecl_vs_coszenith;cos(theta);sin(decl)",
						   20, -1,1, 
						   20, -1,1) 

	for idz in range( 1, sindecl_vs_coszenith.GetNbinsX()+1 ) :
		
		dz = sindecl_vs_coszenith.GetXaxis().GetBinCenter( idz )

		for az in aa.frange( 0, pi, 0.001 ):

			el = pi/2 - acos( dz )
			eq = ROOT.sla.dh2e( az, el , lattitude )        
			sindecl_vs_coszenith.Fill( dz , sin(eq.DEC) )

		# normalize the dz slice -- in a very painful way
		c =0
		for idecl in range ( 1,  sindecl_vs_coszenith.GetNbinsY()+1 ) :
			bin = sindecl_vs_coszenith.GetBin( idz, idecl )
			c += sindecl_vs_coszenith.GetBinContent( bin )
			
		for idecl in range ( 1,  sindecl_vs_coszenith.GetNbinsY()+1 ) :
			bin = sindecl_vs_coszenith.GetBin( idz, idecl )
			sindecl_vs_coszenith.SetBinContent( bin , sindecl_vs_coszenith.GetBinContent( bin) / c )
	
	return sindecl_vs_coszenith



def to_eq( h , name = "decl" , lattitude = ( 36 + 16 / 60.0 ) * pi / 180 ): 

	"turn a histogram with cos(zenith) on the y-axis into one with sin(decl)"

	r = h.Clone( name )
	r.Scale(0)
	
	for ix in range( 1, h.GetNbinsX()+1 ) :

		x = h.GetXaxis().GetBinCenter( ix )
		
		for iy in range( 1, h.GetNbinsY()+1 ) :

			dz = h.GetYaxis().GetBinCenter( iy )
			b = h.GetBin( ix, iy )

			N = 1000
			
			for i in range (N):
				az = 2*pi*i/N
				el = pi/2 - acos( dz )
				eq = ROOT.sla.dh2e( az, -el , lattitude )        

				cc = h.GetBinContent(b)
				
				r.Fill( x, sin( eq.DEC) , cc /N )

	return r


def scale ( h ):
	
	"normalize the effective area histograms"

	binsize_e   = h.GetXaxis().GetBinWidth(1)
	binsize_ang = h.GetYaxis().GetBinWidth(1)
	omega = 2*pi *binsize_ang # each bin covers this solid angle
	
	# ngen scaling has already been applied.
	S = 1 / (  log(10) * binsize_e * 3600.0 * 24 * 365 * omega )
	h.Scale(S)
	



f = ROOT.TFile("aeff.root")

h_nu = f.Get("h_nu")
h_anu = f.Get("h_anu")


scale(h_nu)
scale(h_anu)

for h in [h_nu, h_anu] :
	h.SetXTitle("logE")
	h.SetYTitle("cos(zenith)")


h_nu_decl  = to_eq ( h_nu,  "h_nu_decl" )
h_anu_decl = to_eq ( h_anu, "h_anu_decl" )

h_nu_decl.SetTitle("h_nu_decl")
h_anu_decl.SetTitle("h_anu_decl")
h_nu_decl.SetYTitle("sin(decl)")
h_anu_decl.SetYTitle("sin(decl)")

h_ave = h_nu.Clone("h_ave")
h_ave.Add( h_anu )

h_ave_decl = h_nu.Clone("h_ave_decl")
h_ave_decl.Add( h_anu_decl )

h_ave.SetTitle("average nu+nubar/2 vs cos(zenith)")
h_ave_decl.SetTitle("average nu+nubar/2 vs sin(decl)")



if True :

	from ROOT import gPad, TCanvas

	c = TCanvas("c","c",800,800)
	c.Divide(2,3)

	opt = "lego2"

	c.cd(1)
	gPad.SetLogz()
	h_nu.UseCurrentStyle()
	h_nu.Draw(opt)

	c.cd(2)
	gPad.SetLogz()
	h_anu.UseCurrentStyle()
	h_anu.Draw(opt)

	c.cd(3)
	gPad.SetLogz()
	h_nu_decl.UseCurrentStyle()
	h_nu_decl.Draw(opt)

	c.cd(4)
	gPad.SetLogz()
	h_anu_decl.UseCurrentStyle()
	h_anu_decl.Draw(opt)

	c.cd(5)
	gPad.SetLogz()
	h_ave.UseCurrentStyle()
	h_ave.Draw(opt)

	c.cd(6)
	gPad.SetLogz()
	h_ave_decl.UseCurrentStyle()
	h_ave_decl.Draw(opt)
	
	www.add( c)

	
def to_table( h2 ):

	T = table( ["binnr",
				h2.GetXaxis().GetTitle()+"-min",
				h2.GetXaxis().GetTitle()+"-max",
				h2.GetYaxis().GetTitle()+"-min",
				h2.GetYaxis().GetTitle()+"-max",
				"Aeff(m2)"] )

	for ix in range( 1, h2.GetNbinsX()+1 ) :

		for iy in range ( 1, h2.GetNbinsY()+1 ) :

			b = h2.GetBin( ix, iy )

			T.append( [ b,
						h2.GetXaxis().GetBinLowEdge( ix ),
						h2.GetXaxis().GetBinUpEdge( ix ),
						h2.GetYaxis().GetBinLowEdge( iy ),
						h2.GetYaxis().GetBinUpEdge( iy ),
						"%.4f" % (h2.GetBinContent(b),)
						] )

	return T

s = to_table( h_ave ).sprint_ascii()
www.add( "<pre>"+ s + "</pre>" )
www.add( "<pre>"+ to_table( h_ave_decl ).sprint_ascii() + "</pre>" )

	
def get_slice ( h , yval1, yval2 ):

	"return a x-slice of the 2d-histogrma h, for y-values between yval1 and yval2 "
	bin1 = h.GetYaxis().FindBin( yval1 )
	bin2 = h.GetYaxis().FindBin( yval2 )
	r =  h.ProjectionX( h.GetName()+str(bin1)+str(bin2), bin1, bin2 )
	r.Scale( 1.0 / (1+bin2-bin1) )

	return r


to_table( h_ave )


c2 = TCanvas("c2","c2",800,800)
c2.Divide(2,2)

c2.cd(1)

gPad.SetLogy()


h__0 = get_slice( h_ave, -.99, -.8 )
h__0.Draw("l hist")

h__1 = get_slice( h_ave, -0.1, 0.1 )
h__1.Draw("l hist")

h__2 = get_slice( h_ave, .8, .99 )
h__2.Draw("l hist same")

gPad.Update()
c2.cd(2)
gPad.SetLogy()

hd__0 = get_slice( h_ave_decl, -.99, -.8 )
hd__0.Draw("l hist")

hd__1 = get_slice( h_ave_decl, -0.1, 0.1 )
hd__1.Draw("l hist")

hd__2 = get_slice( h_ave_decl, .8, .99 )
hd__2.Draw("l hist same")

gPad.Update()
www.add( c2 )
del www

if copy_cmd : os.system( copy_cmd )

print ('bye!')