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

def findDimentionBoxElement(node):
    sh = node.GetVolume().GetShape()
    if sh.IsCylType():
        #print node.GetName(), " is a cylinder", sh.GetRmin(), sh.GetRmax()
        if sh.HasRmin():
            r = (sh.GetRmax() - sh.GetRmin())
        else:
            r = 0.
        return {'x':sh.GetDX(),
                'y':sh.GetDY(),
                'z':sh.GetDZ(),
                'r':r}
    return {'x':sh.GetDX(),
            'y':sh.GetDY(),
            'z':sh.GetDZ()}

def findPositionElement(node):
    pos = node.GetMatrix().GetTranslation()
    sh = node.GetVolume().GetShape()
    if sh.IsCylType():
        if sh.HasRmin():
            r = (sh.GetRmax() - sh.GetRmin())/2.
        else:
            r = sh.GetRmax()
        return {'x':pos[0],
                'y':pos[1],
                'z':pos[2],
                'r':r}
    return {'x':pos[0],
            'y':pos[1],
            'z':pos[2]}

def findPositionElement2(node,topPos):
    local_pos = node.GetMatrix().GetTranslation()
    pos = []
    for i in xrange(3):
        pos.append(local_pos[i]+topPos[i])
    sh = node.GetVolume().GetShape()
    if sh.IsCylType():
        if sh.HasRmin():
            r = (sh.GetRmax() - sh.GetRmin())/2.
        else:
            r = sh.GetRmax()
        return {'x':pos[0],
                'y':pos[1],
                'z':pos[2],
                'r':r}
    return {'x':pos[0],
            'y':pos[1],
            'z':pos[2]}

def findDimentionBoxElement2(node,top):
    sh = node.GetVolume().GetShape()
    if sh.IsCylType():
        #print node.GetName(), " is a cylinder", sh.GetRmin(), sh.GetRmax()
        if sh.HasRmin():
            r = (sh.GetRmax() - sh.GetRmin())
        else:
            r = 0.
        return {'x':sh.GetDX(),
                'y':sh.GetDY(),
                'z':sh.GetDZ(),
                'r':r}
    return {'x':sh.GetDX(),
            'y':sh.GetDY(),
            'z':sh.GetDZ()}

# 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)

listOfHitLSSegments = None
numberOfHitLSSegments = None
zstarts = []
zends = []
lsNames = []
Nz, Nphi, nBins = None, None, None
Nphi = 16
dphi = 2.*math.pi/Nphi
phistart = dphi/2.
zin, zout = None, None
zstepsize = None
# To hold the energy loss
EnergyDeposits = []
abig = None
asmall = None
b = None

def makeLSsegments(nodes, tree, savelist=False):
    global zstarts, zends, lsNames, Nz, Nphi, nBins, zin, zout, zstepsize
    global abig, asmall, b
    abig = 250.#nodes['T2LiSc_1_1']['x']['dim']
    asmall = 186.#nodes['T1LiSc_1_1']['x']['dim']
    b = 500.#nodes['T2LiSc_1_1']['y']['dim']
    #print abig, asmall, b
    for node in nodes.keys():
        if 'LiSc' in node:
            zstarts.append(nodes[node]['z']['pos']-nodes[node]['z']['dim'])
            zends.append(nodes[node]['z']['pos']+nodes[node]['z']['dim'])
            lsNames.append(node)
    # ATTENTION! there is some superposition in z
    # To keep it clean & easy, we will do the segmentation
    # as Martin first suggested, i.e. by hands!!!
    # It cannot be so bad...
    zstarts.sort()
    zends.sort()
    #Nz = len(zstarts)
    zin = min(zstarts)
    zout =max(zends)
    #print 'zstart: ', zin, 'zend: ', zout
    zstepsize = 100.
    Nz = int(math.ceil((zout-zin)/zstepsize))
    #print Nphi, Nz, nBins
    nBins = Nphi*Nz
    # Init energy deposit to 0.
    for i in xrange(nBins):
        EnergyDeposits.append(0.)
    global listOfHitLSSegments, numberOfHitLSSegments, EnergyDeposits
    if savelist:
        tree, listOfHitLSSegments = AddVect(tree, 'listOfHitLSSegments', 'int')
    tree, numberOfHitLSSegments = AddVar(tree, 'numberOfHitLSSegments', 'int')

def hitSegments(vetoPoints, nodes, threshold, savelist=False):
    global listOfHitLSSegments, numberOfHitLSSegments, EnergyDeposits
    global abig, asmall, b, phistart, dphi, Nphi, nBins, zin, zout
    # Init energy deposit to 0.
    for i in xrange(nBins):
        EnergyDeposits[i]=0.
    endsmall = nodes['T1Endplate_1']['z']['pos']+nodes['T1Endplate_1']['z']['dim']
    #print 'endsmall: ', endsmall
    #print 'nBins: ', nBins, 'len(list): ',len(EnergyDeposits)
    for hit in vetoPoints:
        z = hit.GetZ()
        if (z<(nodes['VetoTimeDet_1']['z']['pos']+nodes['VetoTimeDet_1']['z']['dim'])
            or z > zout):
            continue
        x = hit.GetX()
        y = hit.GetY()
        a = abig
        if z < endsmall:
            a = asmall
        phi = ROOT.TMath.ATan2(a*y,b*x)+2*math.pi
        zbin = ROOT.TMath.FloorNint((z-zin)/zstepsize)
        phibin = ROOT.TMath.FloorNint((phi+phistart)/dphi)%Nphi
        iBin = zbin*Nphi+phibin
        try:
            EnergyDeposits[iBin] += hit.GetEnergyLoss()
        except:
            print 'hitSegments: could not find bin ', iBin
            print x, y, z
            sys.exit()
    nHit = 0
    for iBin in xrange(nBins):
        if EnergyDeposits[iBin] > threshold:
            nHit +=1
            if savelist:
                Push(listOfHitLSSegments, int(iBin))
    Push(numberOfHitLSSegments, nHit)