import aa, ROOT
from math import *
from util.pyroot_util import *
from loadlik import L # this will load the pdfs

from ROOT import CascadePdf, Params, Paramst, Trk, Det, Hit, Pmt, Mat, Vec

def again() : 
	exec(open("test_CascadePdf.py").read())



def test_angles() :

	theta , phi = 0.3 , 0.4 

	t = Trk(); t.dir.set(0,0,1)

	p = Pmt(); 
	p.dir.set_angles( theta, phi )
	print (p.dir)
	p.pos.set( 100, 0, 100 )

	M = Mat()
	M.set_rotation( Vec(1,2,3).normalize() , 0.12345 )

	p.pos = M * p.pos;
	p.dir = M * p.dir;
	t.pos = M * t.pos;
	t.dir = M * t.dir;

	p = Params( t, p )

	dtheta = p.theta - theta 
	dphi   = p.phi - phi
	print(dtheta, dphi)

	print(p)

	return abs(dtheta) + abs(dphi) < 1e-14


def theta_phi_map( pdf ) :

	h = ROOT.TH2D("h2d","h2d", 40, 0, pi, 40, 0, pi )

	for bin in bins(h) :
		bin.content = pdf.Ntot( Params( 100, 100, 0.7, bin.x, bin.y ))

	h.DrawCopy("colz")



def plot_pdf ( pdf , E, what = "dN_dt", pos = Vec( 20,0,120) , theta = 2.5, phi = 3.0 ) :

	trk = Trk(); trk.dir.set(0,0,1)
	trk.E =E
	p = Pmt(); p.pos = pos; p.dir.set_angles( theta, phi )
	pars = Params( trk, p )

	h0 = ROOT.TH1D("h0","h0",200,-50,350 )
	f = lambda t : getattr( pdf, what )( Paramst( pars, t ))

	fill_hist( f, hist=h0 )

	return h0


def plot_ntot( pdf , E, pos = Vec( 20,0,220) , theta = 2.5, phi = 3.0 , npts = 200 ) :

	trk = Trk(); trk.dir.set(0,0,1)
	trk.E =E
	p = Pmt(); p.pos = pos; p.dir.set_angles( theta, phi )
	pars = Params( trk, p )

	h0 = ROOT.TH1D("h0","h0",npts,0.72,0.76) 

	for b in bins(h0) :
		pars.z = b.x
		b.content = pdf.Ntot( pars )

	h0.DrawCopy()



c = splitcanvas( 1 )

plot_ntot( L.pdf, 1e5 )

import sys
# sys.exit(0)




def plot_elong( pdf , E = 1e5 ):

	c = ROOT.TCanvas("c","c",500,500)
	ROOT.SetOwnership(c, False )

	pdf.set_nsamples(0)
	h = plot_pdf( pdf, E)
	h.SetLineColor(1)
	h.DrawCopy("")

	pdf.set_nsamples(1)
	h=plot_pdf( pdf, E)
	h.SetLineColor(2)
	h.DrawCopy("same")

	pdf.set_nsamples(20)
	h=plot_pdf( pdf, E)
	h.SetLineColor(4)
	h.DrawCopy("same")

def tst() :
	ps = Params( 1000, 100, 1,1,1 )
	L.pdf.Ntot(ps)
	L.pdf.Ntot(Paramst(ps,10) )
	L.pdf.N(Paramst(ps,10) )
	L.pdf.dN_dt(Paramst(ps,10) )
	#L.pdf.dP1_dt(Paramst(ps,10) )

p = Paramst(1000, 100, 1,1,1,10)
z = L.pdf.dP1_dt( p  )



plot_elong( L.pdf )
theta_phi_map( L.pdf )
test_angles()

def plts (E = 1e6, pos = Vec( 10,0,100), dr = Vec( -1, 0, -1).normalize() ):

	h0 = ROOT.TH1D("h0","h0",300,-50,250)

	trk = Trk(); 
	trk.dir.set( 0,0,1)
	trk.E = E
	pp = Pmt()
	pp.pos, pp.dir  = pos, dr
	ps = Params( trk, pp )

	funcs = [   lambda t : L.pdf.Ntot   ( ps )             ,
		        lambda t : L.pdf.N      ( Paramst( ps, t )),
	            lambda t : L.pdf.dN_dt  ( Paramst( ps, t )),
	            lambda t : L.pdf.dP1_dt ( Paramst( ps, t )) ]

	c = splitcanvas( len(funcs) , Title = str(ps))

	L.pdf.set_nsamples( 0 )

	draw_into( c, [ fill_hist(f, xrange=[-50,200 ] , LineColor = 1 ) for f in funcs] )

	L.pdf.set_nsamples( 20 )

	draw_into( c, [ fill_hist(f, xrange=[-50,200 ] , LineColor = 2 ) for f in funcs] , opt = "same" )

	c.Update()

plts()