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

from ROOT import Trk, Det, Hit, Pmt, Mat, Vec, Det, Evt

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


L.pdf.set_nsamples( 20 )

det = Det( aa.aadir + "/data/KM3NeT_-00000001_20171212.detx" )
#det.shift( - det.get_cg() )

def sim( det, trk ) :

    hits = L.generate( trk, det.pmts )
    print ( "generated", len(hits),"hits")
    return hits

pmts = ROOT.vector(Pmt)()
for i,pmt in det.pmts  : pmts.push_back( pmt )

t = Trk();
#1904: Trk: id=0 pos=Vec:65.848 -22.603 291.387 dir=Vec:-0.566477 0.237638 0.78907 t=0 E=63668.4 pdg-type=12
#t.E = 63668.4
#t.pos.set( 65.848, -22.603, 291.387 )
#t.dir = Vec( -0.566477, 0.237638, 0.78907 )
#1855: Trk: id=0 pos=Vec:-204.444 100.579 497.136 dir=Vec:0.00308 0.987024 -0.160543 t=0 E=93890.9 pdg-type=12
t.E = 93890.9
t.pos.set( -204.444, 100.579, 497.136 )
t.dir = Vec( 0.00308, 0.987024, -0.160543)

hits = sim( det, t )

L.set_event( pmts, hits )

f = ROOT.EventFile()
f.set_output("aart_toy_simulation_1855_1e5GeV_elong.root")
f.evt.clear()
f.evt.mc_trks.push_back( t )
for hit in hits: f.evt.mc_hits.push_back( hit )
f.write()

print("wrote")

names = "x,y,z,theta,phi,logE".split(",")


def plot_event ( trk, hits ) :


    names = "d z theta phi t".split()

    c = splitcanvas( 6 )

    for i,n in enumerate( names ) : 
        c.cd(i+1)
        hist = ROOT.TH1D("_h"+n,"_h"+n, 30,0,-1 )

        for h in hits :
            pp = ROOT.Paramst(trk, h)
            hist.Fill( getattr( pp , n ) )

        hist.DrawCopy()

plot_event(t, hits)        


def to_trk( x,y,z,theta,phi,logE ) : 
    t = Trk();
    t.pos.set(x,y,z)
    t.dir.set_angles( theta, phi )
    t.E = 10**logE
    return t


def f( x, y, z, theta, phi, logE ) :
    r = -L.eval( to_trk( x,y,z,theta,phi,logE ) )
    print( t.dir.len() , x,y,z, theta, phi, logE, " ===> ", r )
    return r

def fvec( x, y, z, theta, phi, logE ) :
    r = L.eval( to_trk( x,y,z,theta,phi,logE ) )
    print( "res: ", r );
    return [ L.L_hit +L.L_unhit, +L.L_time, r ]



#no trk in function? bug?
def vals ( trk ) : return [ t.pos.x, t.pos.y, t.pos.z, t.dir.theta(), t.dir.phi(), log10( t.E ) ]



def print_hessian( ) :


    X = vals( t )
    H = hessian( fvec )

    print ("computing hessian") 

    hessian = H( X )


    total_hessian = [ [ r[2] for r in row] for row in hessian ]

    for r in hessian : print ( r )

    C = np.matrix( total_hessian ).I

    print ( C )

    for i,n in enumerate( names ) :
        print ( n, "\t ", sqrt(-C[i,i]),  "info. from hit presence / time ", hessian[i][i][0], "/", hessian[i][i][1] )


d = 1
da = 0.04

ranges = [ (t.pos.x-d, t.pos.x+d),(t.pos.y-d, t.pos.y+d), (t.pos.z-d, t.pos.z+d), 
            (t.dir.theta()-da, t.dir.theta()+da) , 
            (t.dir.phi()-da, t.dir.phi()+da) , 
            (log10(t.E)-2, log10(t.E)+2) ]



def plot_slices ( trk ) :


    ranges = [ (t.pos.x-d, trk.pos.x+d),(trk.pos.y-d, trk.pos.y+d), (trk.pos.z-d, trk.pos.z+d), 
            (trk.dir.theta()-da, trk.dir.theta()+da) , 
            (trk.dir.phi()-da, trk.dir.phi()+da) , 
            (log10(trk.E)-2, log10(trk.E)+2) ]

    print ( 'tkr pos', trk.pos )

    c = splitcanvas( 6 )

    V0 = vals(trk)
    F0 = fvec(*V0) 
    N  = len( F0 )

    for i in range( 6 ) :
        
        print("trk for {}: {}".format( names[i], trk ) )

        V = vals( trk )
        c.cd(i+1)
        
        h = [ ROOT.TH1D("h"+names[i]+str(j),names[i], 15, *ranges[i] )    for j in range(N) ]
        
        def ff( q ) : 
            V[i] = q
            return fvec( *V )


        for bs in zip( *[ bins(hh) for hh in h] ):

            r = ff( bs[0].x )

            print ( V, " \t " , names[i], bs[0].x , r )

            for j, val in enumerate( r ) :
                bs[j].content = val


        equalize(h)

        h[2].SetLineWidth(3)

        for term in range(N) :
            h[term].SetLineColor( term + 1 )
            h[term].DrawCopy( "L same" if term> 0 else "L" )



    return c


#plot_slices(t)
                


#best = minimize( f , start_values = [ t.pos.x, t.pos.y, t.pos.z, t.dir.theta(), t.dir.phi(), log10( t.E )] )

#tbest = to_trk( *best )

plot_slices( t )


#for i,n in enumerate( names ) :#
#print ( n, "\t ", sqrt(-C[i,i]),  "info. from hit presence / time ", hessian[i][i][0], "/", hessian[i][i][1] )