from elena import *
from array import array
from collections import OrderedDict
import ROOT

# Custom sorting
def sortNodesBy(nodes, coord='z', what='pos', printdict=True):
    d = OrderedDict(sorted(nodes.items(), key = lambda item: item[1][coord][what]))
    if printdict:
        print
        print
        print "Nodes by %s %s:"%(coord, what)
        for node in d.keys():
            print node, '\t\t', (d[node][coord]['pos']-d[node][coord]['dim'])/100., ' m', '\t', 2*d[node][coord]['dim']/100., ' m'
    return d

def retrieveMCParticleInfo(sTree, key):
    mcPartKey = sTree.fitTrack2MC[key]
    mcPart    = sTree.MCTrack[mcPartKey]
    if not mcPart:
        print "Error in  retrieveMCParticleInfo"
        sys.exit()
    mother    = sTree.MCTrack[mcPart.GetMotherId()]
    #print mcPartKey, mcPart, mcPart.GetMotherId(), mother
    try:
        mumId = mother.GetPdgCode()
    except:
        mumId = 0
    return int(mcPart.GetPdgCode()), int(mumId), mcPart.GetStartX(), mcPart.GetStartY(), mcPart.GetStartZ()


def AddVect(t,name,vectType):
    vect =  ROOT.vector(vectType)()
    t.Branch( name, vect )
    return t, vect

def AddVar(t,name,varType):
    var = array(varType[0],[-999])
    t.Branch(name,var,name+"/"+varType.upper())
    return t, var

def PutToZero(var):
    if not isinstance(var, array):
        var.clear()
    else:
        var[0] = 0

def Push(leaf, value):
    if not isinstance(leaf, array):
        leaf.push_back(value)
    else:
        leaf[0] = value


def wasFired(indexKids, detPoints, detPos, pointsVects=None, Ethr=0):
    def lookingForHits(detPoints,flag,nstat,nstat_eff,indexKids,Eloss,Ethr):
        for hit in detPoints:
            if (indexKids is None) or (hit.GetTrackID() in indexKids):
                #print RPChit.GetTrackID(), t_ID
                # check if it is in one of the considered active layer
                for pos in detPos:
                    if pos[0]<=hit.GetZ()<=pos[1]:
                        Eloss += hit.GetEnergyLoss()
                        #print pos, hit.GetZ()
                        if pointsVects is not None:
                            pointsVects[0].push_back(hit.GetX())
                            pointsVects[1].push_back(hit.GetY())
                            pointsVects[2].push_back(hit.GetZ())
                            
                        flag = True
                        nstat+=1
                        #eff_val = gRandom.Uniform(0.,1.)
                        #if eff_val<0.9:
                        #    flag_eff = flag_eff or True
                        #    nstat_eff+=1
                        #else:
                        #    flag_eff = flag_eff or False
        flag_Ethr = Eloss>=Ethr
        return flag,  nstat, nstat_eff, Eloss, flag_Ethr
    # Now in partKidTrackID I should have the trackID connected to my charged particles
    #flag_eff = False
    flag = False
    nstat = 0
    nstat_eff = 0
    Eloss = 0
    flag,nstat,nstat_eff,Eloss,flag_Ethr = lookingForHits(detPoints,flag,nstat,nstat_eff,indexKids,Eloss,Ethr)
    #if flag==False and checkOn:
        #print "To be study event %s"%entry
        #PrintEventPart(particles,pdg)
        #PrintRPChits(detPoints,pdg)
        #print
        #print
        #assert(False)        
    #return flag, flag_eff, nstat, nstat_eff, flag_Ethr
    return flag, flag_Ethr


#ff = ROOT.TFile("histoForWeights.root","read")
#h_GioHans = ff.Get("h_Gio")

def calcWeight(E, w, entries, nuName, weightHist):
    if "bar" in nuName:
        reduction = 0.5
        flux = 6.98e+11 * 2.e+20 / 5.e+13
    else: 
        reduction = 1.
        flux = 1.09e+12 * 2.e+20/ 5.e+13
    crossSec = 6.7e-39*E * reduction
    NA = 6.022e+23
    binN = weightHist.GetXaxis().FindBin(E)    
    return crossSec * flux * weightHist.GetBinContent(binN) * w * NA / entries


def calcWeightOldNtuple(x,y,z, E, nodes, entries, weightHist, datatype):
    if 'sig' in datatype:
        return -999
    params = {'OPERA': {'material':'Fe','lenght':60.,},
          'window1': {'material':'Fe','lenght':nodes["lidT1O_1"]['z']['dim']},
          'window2': {'material':'Al','lenght':nodes["lidT1I_1"]['z']['dim']},
          'volumeOut': {'material':'Fe','lenght':30.,},
          'volumeIn': {'material':'Al','lenght':8.}}
    materialPars = {'Fe':{'A':55.847, #g/mol # fuori
                          'rho': 7874 * 1.e+3/1.e+6,#g/cm3
                          },
                    'Al':{'A':26.98, #g/mol # dentro
                          'rho': 2700 * 1.e+3/1.e+6 #g/cm3
                          }}
    perc = {'OPERA':0.33,
            'window1':0.015,
            'window2':0.015,
            'volumeOut':0.32,
            'volumeIn':0.32}
    intElName = interactionElement(x,y,z, nodes)
    #print intElName
    NA = 6.022e+23
    material = params[intElName]['material']
    crossSec = 8.8e-39*E #1.5e-38*fabs(E) ##3./2*(1.2+0.35)*1.e-38*fabs(E)
    flux = 2.e+18#5.e+16
    L = params[intElName]['lenght']
    rho = materialPars[material]['rho']
    n_element = perc[intElName]*entries
    binN = weightHist.GetXaxis().FindBin(E)
    weight = crossSec * NA * rho * L * flux / n_element * weightHist.GetBinContent(binN)
    return weight

def interactionElement(x,y,z, nodes):
    # window1
    if nodes["lidT1O_1"]['z']['pos']-nodes["lidT1O_1"]['z']['dim'] < z < nodes["lidT1O_1"]['z']['pos']+nodes["lidT1O_1"]['z']['dim']:
        return "window1"
    # window2
    if nodes["lidT1I_1"]['z']['pos']-nodes["lidT1I_1"]['z']['dim'] < z < nodes["lidT1I_1"]['z']['pos']+nodes["lidT1I_1"]['z']['dim']:
        return "window2"
    # vacuum tank borders
    if nodes["lidT1O_1"]['z']['pos']+nodes["lidT1O_1"]['z']['dim'] < z < nodes["lidT6O_1"]['z']['pos']+nodes["lidT6O_1"]['z']['dim']: # include le lid x la parte fuori ma non x la parte dentro
        if inInnerLid(x,y,z, nodes):
            return "volumeIn"
        else:
            return "volumeOut"
    # opera
    try:
        minz, maxz = nodes["volIron_12"]['z']['pos']-nodes["volIron_12"]['z']['dim'], nodes["volIron_24"]['z']['pos']+nodes["volIron_24"]['z']['dim']
    except KeyError:
        try:
            minz, maxz = nodes["volLayer_11"]['z']['pos']-nodes["volLayer_11"]['z']['dim'], nodes["volLayer_0"]['z']['pos']+nodes["volLayer_0"]['z']['dim']
        except KeyError:
            # Pazienza esaurita, ora lo hanno chiamato "volMagneticSpectrometer_1" ed e' lungo piu' di 9m!!!
            try:
                minz, maxz = nodes["volMagneticSpectrometer_1"]['z']['pos']-nodes["volMagneticSpectrometer_1"]['z']['dim'], nodes["volMagneticSpectrometer_1"]['z']['pos']+nodes["volMagneticSpectrometer_1"]['z']['dim']
            except KeyError:
                print "Cannot find OPERA"
                sys.exit()
    #if nodes["volIron_12"]['z']['pos']-nodes["volIron_12"]['z']['dim'] < z < nodes["volIron_24"]['z']['pos']+nodes["volIron_24"]['z']['dim']:
    #if nodes["volLayer_11"]['z']['pos']-nodes["volLayer_11"]['z']['dim'] < z < nodes["volLayer_0"]['z']['pos']+nodes["volLayer_0"]['z']['dim']:
    if minz < z < maxz:
        return "OPERA"


def inInnerLid(x,y,z, nodes):
    if nodes["lidT1I_1"]['z']['pos']-nodes["lidT1I_1"]['z']['dim'] < z < nodes["lidT6I_1"]['z']['pos']+nodes["lidT6I_1"]['z']['dim']:
        if z < nodes["Veto_5"]['z']['pos']-nodes["Veto_5"]['z']['dim']:
            a,b = nodes['T1I_1']['x']['dim'], nodes['T1I_1']['y']['dim'] # smaller part of the tank (first 5 m)
        else:
            a,b = nodes['T2I_1']['x']['dim'], nodes['T2I_1']['y']['dim'] # larger part of the tank
        if inEllipse(x,y,a,b):
            return True
    return False


def inEllipse(x,y,a,b):
    if ( (x*x)/(a*a) + (y*y)/(b*b) ) < 1.:
        return True
    return False

def pointInEllipse(point,a,b):
    x = point.X()
    y = point.Y()
    return inEllipse(x,y,a,b)