import aa, ROOT
import os,sys, itertools, copy, collections, pickle
from array import array
from math import *



def integrate( func , x0, x1 , dx = 0.01 , jacobian = lambda dx : dx):

    c = 0
    for x in aa.frange( x0, x1, dx ) :
        c += func( x ) * jacobian( dx ) 
    return c



class BackgroundModel :
   
    """
    The background model consists = N x F(sindecl) x F(logE) [ sr-1 ]
    The two functions are derived from a histogram of the data. The former is a spline, the
    latter consists of two Gaussians that are fit to the distribution. 
    The constant N is chosen so that the integral over the energy range and the over the
    sphere (2 pi * \int decl from -1 to 1 ) is exactly the number of event in the data-sample.
    """
    
    
    def dp_dOmega( self, sindecl ) :
        "pdf of sin(decl) per sr"
        return self.decl_norm * self.decl_spline.Eval( sindecl )

    def dp_dlogE( self, logE ) :
        "pdf per unit log10(E)"
        return self.loge_norm * self.loge_func.Eval( logE )

    def eval( self, sindecl, logE ) :
        "number of backgournd events per sr and per unit logE"
        r = self.norm * self.dp_dOmega( sindecl ) * self.dp_dlogE( logE )
        return r



    def __init__( self , hist ) :

        "histogram has decl,ra,logE on x,y,z axes "

        self.hdec  = hist.ProjectionX()
        self.hra   = hist.ProjectionY(); self.hra.SetMinimum(1e-3)
        self.hloge = hist.ProjectionZ(); self.hloge.SetMinimum(1e-3)

        # extract spline for sin-decl

        xs,ys = array('d'), array('d')

        hdec = self.hdec
        # for b in range(1,hdec.GetNbinsX()+1):
          
        #     if   b == 1                : x = hdec.GetXaxis().GetBinLowEdge(b)
        #     elif b == hdec.GetNbinsX() : x = hdec.GetXaxis().GetBinUpEdge(b)
        #     else                       : x = hdec.GetXaxis().GetBinCenter(b)
        #     xs.append ( x )
        #     ys.append ( hdec.GetBinContent( b ))
        
        #self.decl_spline = ROOT.TSpline3( "decl_dependence", xs, ys, len(xs) , "e1b1"  )
        self.decl_spline =  ROOT.TSpline3()
        ROOT.IntSpline(hdec,self.decl_spline)
        self.decl_spline.SetLineWidth(2)
        self.decl_spline.SetLineColor(2) 

        dectest = hdec
        dectest.Scale(1/hdec.Integral()/hdec.GetBinWidth(0))
        dectest.Draw() 
        self.decl_spline.Draw("SAME")
        self.hdec.Write()

        # Define and fit the function for the energy dependence

        self.loge_func = ROOT.TF1("double_gaus", "gaus(0) + gaus(3)", 1, 8)
        self.loge_func.SetParNames("norm1", "mean1", "sigma1","norm2", "mean2", "sigma2")
        self.loge_func.SetParameters( 9e2,4,0.5, 100,2,1)
        self.hloge.Fit( self.loge_func ,"0")

        #pffff fix root on the spot
        ROOT.TSpline.Integral = lambda s,a,b : sum( s.Eval(x)*1e-3 for x in aa.frange(a,b,1e-3) )         

        # normalisation 

        # decl_norm * spline = pdf (sr-1)
        print("check int is (1):", self.decl_spline.Integral(-1,1))
        self.decl_norm = 1.0 / ( 2 * pi * self.decl_spline.Integral(-1,1) )

        # loge_norm * loge_func = pdf ( per unit of log10(E) )
        self.loge_norm = 1.0 / ( self.loge_func.Integral(1,8) )
       
        self.norm = hist.Integral()

        print ( "                 ---------------- check integral loge : " , integrate( self.dp_dlogE, 1,8,0.01 ) )
        print ( "                 ---------------- check integral decl : " , integrate( self.dp_dOmega , -1,1,0.01, lambda dsindecl : 2*pi*dsindecl ) )


    def plot( self , canvname = "c1") :

        c = ROOT.TCanvas( canvname , canvname ,1200,400)    
        c.Divide(3,1); 
        
        c.cd(1)
        self.hdec.SetMinimum(0)
        self.hdec.Draw()
        self.decl_spline.Draw("same")

        c.cd(2) 
        self.hra.Draw()
        self.hra.Fit("pol0") # for illustration only

        c.cd(3)
        ROOT.gPad.SetLogy()
        self.hloge.Draw()
        self.loge_func.Draw("same")
            
        return c
    



    def write_to_webpage( self, W  ) :
        W.write ("background model <br>")
        W.write( self.plot() )
        W.write( " <br> check integral loge : " +str( integrate( self.dp_dlogE, 1,8,0.01 ) ))
        W.write( " <br> check integral decl : " +str( integrate( self.dp_dOmega, -1,1,0.01, lambda dsindecl : 2*pi*dsindecl ) ))


class DataSample :

    def __init__ ( self ) :

        self.name           = "noname"
        self.events         = []
        self.binning        = ( 10, -1,1, 20,0,2*pi, 14, 1,8 )
        self.version_to_use = 7


    def gettrk( self, evt ) :

        if self.version_to_use == 6 : return evt.trks[0]
        if self.version_to_use == 7 : return evt.trks[2] if evt.trks.size()>2 else None
        print ("unknown version")
        sys.exit()

    def missing_runs( self ) :

        if self.version_to_use == 6 : return []
        if self.version_to_use == 7 : return self.missing_v7runs
        print ("unknown version")
        sys.exit()


    def read_dst_file( self, filename ) :

        self.det = ROOT.Det()
        self.det.longitude                  = 0.278819;
        self.det.latitude                   = 0.633407;
        self.det.meridian_convergence_angle = 0.0100733;

        ax = ";sin(#delta);ra;log(E_{rec}/GeV)"
        self.hist = ROOT.TH3D("dhist"+self.name,"data-"+self.name+ax, *self.binning)

        f = ROOT.EventFile( filename )
        self.src = filename

        self.events = [ copy.copy( evt ) for evt in f ]


  
    def compute_coords(self, scramble = False ) :

        print ("computing celestical coordinates using version ", self.version_to_use )
        self.data_is_scrambled = scramble


        for i,evt in enumerate( self.events ):
            trk = self.gettrk( evt )
            if not trk or trk.dir.z < -0.1 : 
                evt.eq = None
            else :
                t = self.events[i-1].t if scramble else evt.t
                evt.eq = ROOT.EquatorialCoords( self.det, trk.dir , t )


    def fill_hist( self ) :

        self.hist.Reset()

        for evt in self.events :
            # nb: always use the v6 version of energy, since this is what is in the IRFs
            if evt.eq == None : continue            
            self.hist.Fill( sin(evt.eq.dec()), evt.eq.ra(), log10( evt.trks[0].E ) )

        print ( "read",len( self.events ),"events. histogram integral = ", self.hist.Integral() )



    def events_table(self ) :

        T = ROOT.Table("number","run","time","dir.z","E", "ra","dec")
        for i,evt in enumerate( self.events ) :
            T.add( i, evt.run_id, evt.t.AsString()[:26], evt.trks[0].dir.z, evt.trks[0].E , evt.eq.dec(), evt.eq.ra()  )
        return T


    def skymap( self , scramble = False ): 

        csky = ROOT.TCanvas("csky","csky",900,500)
        s = ROOT.SkyMap()
        m = ROOT.TMarker()
        m.SetMarkerSize(0.6)
        m.SetMarkerStyle(20)
        for i,e in enumerate( self.events ) :
            if not e.eq : continue 
            s.add( m, e.eq )
        if self.data_is_scrambled : s.add ( ROOT.TLatex( 0,0, "scrambled "))
        s.Draw()
        return csky



    def plot( self , canvname = "cx") :

        hx1 = ROOT.TH1D("hx1","zenith angle;cos(zenith);events",20,-1,1)
        hx2 = ROOT.TH1D("hx2","phi angle;#phi;events",20,-pi,pi)
    
        for evt in self.events:
            hx1.Fill( evt.trks[0].dir.z )
            hx2.Fill( evt.trks[0].dir.phi() )


        cx = ROOT.TCanvas( canvname , canvname ,800,400)    
        cx.Divide(2,1); 
        
        for h in [hx1,hx2] :
            h.SetLineWidth(2)
            h.SetLineColor(4)
            h.SetMinimum(0)

        cx.cd(1)
        hx1.DrawCopy()

        cx.cd(2) 
        hx2.DrawCopy()
            
        return cx


    def info_table( self ) :

        self.nevents_with_v7 = sum( evt.trks.size()>2 for evt in self.events )
        self.nevents         = len( self.events )
        self.nmissing_v7     = len (self.missing_events_in_v7)
        self.v7downgoing     = sum ( 1 for e in self.events if self.gettrk(e) and self.gettrk(e).dir.z<-0.1)

        t = ROOT.Table("attribute","value")
        for k,v in self.__dict__.items():
            t.add( k ).add( str(v)[:99] )
        t.title = 'DataSample' + self.name 
        return t


    def write_to_webpage( self, W , scramble = True ) :

        T = ROOT.Table("attribuute","value")
        T.add( "sample name", self.name )
        T.add(" number of events", len( self.events ))
        T.add(" histogram integral", self.hist.Integral() )
        T.add(" events with coords", sum ( 1 for e in self.events if e.eq != None ))
        T.title ="data sample"
        W.write(T)

        W.write( self.info_table() )

        W.write( self.plot() )

        W.write( self.skymap( scramble ) )







def get_v7( dset, v7dir ) :

    """
    this augments the events with the jpp muon track from another (v7) production.
    The new track is goind at position trks[2]

    the datasample object also get the following atrributes:
    dset.missing_v7runs     # to be used for livetime corrections later
    dset.processed_v7runs
    dset.missing_events_in_v7

    """

    def idef( evt ) :
        return str(evt.frame_index) + " " + str(evt.trigger_counter) + " " + evt.t.AsString()


    dset.missing_v7runs       = []
    dset.processed_v7runs     = []
    dset.missing_events_in_v7 = []  # only those not already in missing v7_runs


    M = collections.defaultdict( list ) 
    for evt in dset.events :
        M[evt.run_id].append(evt)

    runs = M.keys()
    fns = {}

    for run in runs :
        retr = "retr." if dset.name == 'P75.1' else ""
        rr   = "%08d" % (run,)
        fn   = f"{v7dir}/datav7.0.{retr}jchain.arca.aashower.{rr}.root"

        ok = os.path.exists( fn  ) 
        if ok : 
            fns[run] = fn

        else  : 
            fns[run] = None
            dset.missing_v7runs.append(run)

        print ( run, "ok" if ok else " MISSING ", fn )



    for i,run in enumerate( M ) :

       

        fn = fns[run] 
        print ("run ", i, "of", len(M), run, fn )
        if not fn: continue

        f = ROOT.EventFile( fn )
        print ("building index for", run )
        f.roottree().BuildIndex("frame_index","trigger_counter")
        print ("done")

        for evt in M[run] :
            print ( 'search', idef(evt))
            nb = f.roottree().GetEntryWithIndex( evt.frame_index , evt.trigger_counter  )
            if nb < 1 : 
                print ( "failed to find event with", evt.frame_index , evt.trigger_counter  ) 
                dset.missing_events_in_v7.append( evt )
            else :

               evt.trks.resize(3)
               try :
                   evt.trks[2] = ROOT.get_best_reconstructed_jppmuon( f.evt )
                   evt.trks[2].comment = "jpp muon from v7 production"
               except ROOT.Exception:
                   print ( "no jpp muon found in event ", evt )
                   dset.missing_events_in_v7.append( evt )



             

               print ( ' found', idef(evt) )
               print()

        dset.processed_v7runs.append( run )
        


f0 = "/sps/km3net/users/rmuller/production/dst/arca6_data/v6.2/datav6.2.jchain.aashower.dst.merged_9635_10005_pre_upm01_antinoise_upaam01.root"
f1 = "/sps/km3net/users/rmuller/production/dst/arca6_data/v6.2/datav6.2.retr.jchain.aashower.dst.merged_10046_10286_pre_upm01_antinoise_upaam01.root"




def prepare_dataset( name , dst_file ):

    "after this, you can simply load the pickle to have your datasamples ready to go"

    D1 = DataSample()
    D1.name = name
    D1.read_dst_file ( dst_file )
    get_v7( D1, "/sps/km3net/repo/data_processing/tag/v6_ARCA6_run_by_run/prod/data/KM3NeT_00000075/v7.0/reco/datav7.0.jchain.arca.00009743.root" )
    
    with open ( D1.name+".datasample.pickle", 'wb' ) as ofile :
        pickle.dump ( D1  , ofile )



if __name__ == "__main__" :
    D1 = prepare_dataset("P75.1", f1 )
    D0 = prepare_dataset("P75.0", f0 )
    
    sys.exit()