import os
from collections import OrderedDict

#os.system("cd "+os.path.expandvars("$FAIRSHIP")+"/../FairShipRun/")
#os.system(". config.sh")
#os.system("cd "+os.path.expandvars("$FAIRSHIP")+"/../Barbara")
#
## Make debug ntuple
#if not os.path.isfile("nuNutple_debug691.root"):
#    os.system("python ntuple_nuVeto.py")
#
## For my output
#if not os.path.exists("elena_plots"):
#    os.system("mkdir elena_plots")
#
## Import amenities
#from analyseNtupleForVeto import *
#
## Reminder
#print ""
#print "Interaction Elements:"
#print "0\t Last passive element of Opera"
#print "1\t Entrance window of the vacuum tank"
#print "2\t Along the vacuum tank"
#print "3\t Along Opera"
#print "4\t Between windows (in the liquid scintillator)"
#print ""

from lookAtGeo import *

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 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()}

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 searchForNodes3_xyz_dict(inputFile, verbose=False):
    d = {}
    r = loadGeometry(inputFile)
    fGeo = r['fGeo']
    ## Get the top volume
    #fGeo = ROOT.gGeoManager
    tv = fGeo.GetTopVolume()
    topnodes = tv.GetNodes()
    for (j,topn) in enumerate(topnodes):
        # top volumes
        if verbose:
            print j, topn.GetName()
            print "            x: ", findPositionElement(topn)['x'],findDimentionBoxElement(topn)['x']
            print "            y: ", findPositionElement(topn)['y'],findDimentionBoxElement(topn)['y']
            print "            z: ", findPositionElement(topn)['z'],findDimentionBoxElement(topn)['z']
        d[topn.GetName()] = {'x': {}, 'y':{}, 'z':{}, 'r':{}}
        d[topn.GetName()]['x']['pos'] =     findPositionElement(topn)['x']
        d[topn.GetName()]['x']['dim'] = findDimentionBoxElement(topn)['x']
        d[topn.GetName()]['y']['pos'] =     findPositionElement(topn)['y']
        d[topn.GetName()]['y']['dim'] = findDimentionBoxElement(topn)['y']
        d[topn.GetName()]['z']['pos'] =     findPositionElement(topn)['z']
        d[topn.GetName()]['z']['dim'] = findDimentionBoxElement(topn)['z']
        if topn.GetVolume().GetShape().IsCylType():
            d[topn.GetName()]['r']['pos'] =     findPositionElement(topn)['r']
            d[topn.GetName()]['r']['dim'] = findDimentionBoxElement(topn)['r']
        else:
            d[topn.GetName()]['r']['pos'] = 0.
            d[topn.GetName()]['r']['dim'] = 0.
        # First children
        nodes = topn.GetNodes()
        if nodes:
            topPos = topn.GetMatrix().GetTranslation()
            for (i,n) in enumerate(nodes):
                if verbose:
                    print j, topn.GetName(), i, n.GetName()
                    print "            x: ", findPositionElement2(n,topPos)['x'],findDimentionBoxElement(n)['x']
                    print "            y: ", findPositionElement2(n,topPos)['y'],findDimentionBoxElement(n)['y']
                    print "            z: ", findPositionElement2(n,topPos)['z'],findDimentionBoxElement(n)['z']
                d[n.GetName()] = {'x': {}, 'y':{}, 'z':{}, 'r':{}}
                d[n.GetName()]['x']['pos'] =     findPositionElement2(n,topPos)['x']
                d[n.GetName()]['x']['dim'] = findDimentionBoxElement(n)['x']
                d[n.GetName()]['y']['pos'] =     findPositionElement2(n,topPos)['y']
                d[n.GetName()]['y']['dim'] = findDimentionBoxElement(n)['y']
                d[n.GetName()]['z']['pos'] =     findPositionElement2(n,topPos)['z']
                d[n.GetName()]['z']['dim'] = findDimentionBoxElement(n)['z']
                if n.GetVolume().GetShape().IsCylType():
                    d[n.GetName()]['r']['pos'] =     findPositionElement2(n,topPos)['r']
                    d[n.GetName()]['r']['dim'] = findDimentionBoxElement(n)['r']
                else:
                    d[n.GetName()]['r']['pos'] = 0.
                    d[n.GetName()]['r']['dim'] = 0.
                # Second children
                cnodes = n.GetNodes()
                if cnodes:
                    localpos = n.GetMatrix().GetTranslation()
                    localToGlobal = []
                    for i in xrange(3):
                        localToGlobal.append(localpos[i]+topPos[i])
                    for (k,cn) in enumerate(cnodes):
                        if verbose:
                            print j, topn.GetName(), i, n.GetName(), k, cn.GetName()
                            print "            x: ", findPositionElement2(cn,localToGlobal)['x'],findDimentionBoxElement(cn)['x']
                            print "            y: ", findPositionElement2(cn,localToGlobal)['y'],findDimentionBoxElement(cn)['y']
                            print "            z: ", findPositionElement2(cn,localToGlobal)['z'],findDimentionBoxElement(cn)['z']
                        d[cn.GetName()] = {'x': {}, 'y':{}, 'z':{}, 'r':{}}
                        d[cn.GetName()]['x']['pos'] =     findPositionElement2(cn,localToGlobal)['x']
                        d[cn.GetName()]['x']['dim'] = findDimentionBoxElement(cn)['x']
                        d[cn.GetName()]['y']['pos'] =     findPositionElement2(cn,localToGlobal)['y']
                        d[cn.GetName()]['y']['dim'] = findDimentionBoxElement(cn)['y']
                        d[cn.GetName()]['z']['pos'] =     findPositionElement2(cn,localToGlobal)['z']
                        d[cn.GetName()]['z']['dim'] = findDimentionBoxElement(cn)['z']
                        if cn.GetVolume().GetShape().IsCylType():
                            d[cn.GetName()]['r']['pos'] =     findPositionElement2(cn,localToGlobal)['r']
                            d[cn.GetName()]['r']['dim'] = findDimentionBoxElement(cn)['r']
                        else:
                            d[cn.GetName()]['r']['pos'] = 0.
                            d[cn.GetName()]['r']['dim'] = 0.


    return d


def searchForNodes2_xyz_dict(inputFile, verbose=False):
    d = {}
    r = loadGeometry(inputFile)
    fGeo = r['fGeo']
    ## Get the top volume
    #fGeo = ROOT.gGeoManager
    tv = fGeo.GetTopVolume()
    nodes = tv.GetNodes()
    for (i,n) in enumerate(nodes):
        if verbose:
            print i,n.GetName()
            print "            x: ", findPositionElement(n)['x'],findDimentionBoxElement(n)['x']
            print "            y: ", findPositionElement(n)['y'],findDimentionBoxElement(n)['y']
            print "            z: ", findPositionElement(n)['z'],findDimentionBoxElement(n)['z']
        d[n.GetName()] = {'x': {}, 'y':{}, 'z':{}, 'r':{}}
        d[n.GetName()]['x']['pos'] =     findPositionElement(n)['x']
        d[n.GetName()]['x']['dim'] = findDimentionBoxElement(n)['x']
        d[n.GetName()]['y']['pos'] =     findPositionElement(n)['y']
        d[n.GetName()]['y']['dim'] = findDimentionBoxElement(n)['y']
        d[n.GetName()]['z']['pos'] =     findPositionElement(n)['z']
        d[n.GetName()]['z']['dim'] = findDimentionBoxElement(n)['z']
        if n.GetVolume().GetShape().IsCylType():
            d[n.GetName()]['r']['pos'] =     findPositionElement(n)['r']
            d[n.GetName()]['r']['dim'] = findDimentionBoxElement(n)['r']
        else:
            d[n.GetName()]['r']['pos'] = 0.
            d[n.GetName()]['r']['dim'] = 0.
    return d

def sortNodesBy(dd, coord='z', what='pos', printdict=True):
    d = OrderedDict(sorted(dd.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




if __name__ == '__main__':

    # Search nodes
    nodesDict = searchForNodes2_xyz_dict('geofile_full.10.0.Genie-TGeant4.root')
    # Sort nodes by z
    from collections import OrderedDict
    sortedNodes = OrderedDict(sorted(nodesDict.items(), key = lambda item: item[1]['z']['pos']))
    
    def printByZ():
        # Print nodes by increasing z
        print
        print
        print "Nodes by z:"
        for node in sortedNodes.keys():
            print node, '\t\t', (sortedNodes[node]['z']['pos']-sortedNodes[node]['z']['dim'])/100., ' m', '\t', 2*sortedNodes[node]['z']['dim']/100., ' m'
    
    # Sort nodes by dimension along z (the walls of the vacuum tank should be ~60 m long)
    sortedNodesByZdim = OrderedDict(sorted(nodesDict.items(), key = lambda item: item[1]['z']['dim']))
    
    def printByZdim():
        # Print nodes by increasing dimension along z
        print
        print
        print "Nodes by z dimension:"
        for node in sortedNodesByZdim.keys():
            print node, '\t\t', (sortedNodesByZdim[node]['z']['pos']-sortedNodesByZdim[node]['z']['dim'])/100., ' m', '\t', 2*sortedNodesByZdim[node]['z']['dim']/100., ' m'
    
    # Custom sorting
    def sortNodesBy(coord='z', what='pos', printdict=True):
        d = OrderedDict(sorted(nodesDict.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
    
    
    
    # Load geometry (again?)
    #r = loadGeometry('geofile_full.10.0.Genie-TGeant4.root')
    #fGeo = r['fGeo']
    #someOtherNodes = ["Tr%s_%s"%(i,i) for i in xrange(1,5)]
    #someOtherNodes += ["volLateralS1_d_1", "volLateralS2_d_1", "volLateralS1_c_1", "volLateralS2_c_1"]
    #someOtherNodes += ["Target%s_1"%i for i in (0,1,2,3,5)]
    
    
    # Make plots
    from ROOT import *
    cSaver = []
    cSaver.append(TCanvas("entrance_windows","entrance_windows"))
    t2.Draw("startY_nu:startX_nu","isPrimary==1 && (interactionElement==2 || interactionElement==4)")
    
    # Veto
    vetoIn = nodesDict['Veto_5']['z']['pos'] - nodesDict['Veto_5']['z']['dim']
    vetoOut = nodesDict['Veto_5']['z']['pos'] + nodesDict['Veto_5']['z']['dim']
    cSaver.append(TCanvas("veto","veto"))
    t2.Draw("startY_nu:startX_nu","isPrimary==1 && interactionElement==2 && startZ_nu>%s && startZ_nu<%s"%(vetoIn,vetoOut))
    grVeto = TGraph(t2.GetSelectedRows(), t2.GetV2(), t2.GetV1())
    grVeto.SetMarkerColor(kGreen)
    grVeto.SetMarkerStyle(6)
    
    # Decay volume (big one)
    cSaver.append(TCanvas("decay_volume_big","decay_volume_big"))
    volIn = nodesDict['T2Startplate_1']['z']['pos'] + nodesDict['T2Startplate_1']['z']['dim']
    volOut = nodesDict['lidT6I_1']['z']['pos'] - nodesDict['lidT6I_1']['z']['dim']
    print 
    print "\tBig volume: in = %s cm, out = %s cm"%(volIn, volOut)
    print
    t2.Draw("startY_nu:startX_nu","isPrimary==1 && interactionElement==2 && startZ_nu>%s && startZ_nu<%s"%(volIn,volOut))
    grBig = TGraph(t2.GetSelectedRows(), t2.GetV2(), t2.GetV1())
    grBig.SetMarkerColor(kBlue)
    grBig.SetMarkerStyle(6)
    
    # Decay volume (small one)
    cSaver.append(TCanvas("decay_volume_small","decay_volume_small"))
    volInSmall = nodesDict['lidT1I_1']['z']['pos'] + nodesDict['lidT1I_1']['z']['dim']
    #volOutSmall = nodesDict['T1Endplate_1']['z']['pos'] - nodesDict['T1Endplate_1']['z']['dim']
    volOutSmall = nodesDict['Veto_5']['z']['pos'] - nodesDict['Veto_5']['z']['dim']
    print 
    print "\tSmall volume: in = %s cm, out = %s cm"%(volInSmall, volOutSmall)
    print
    print  ' veto ', nodesDict['Veto_5']['z']['pos'] - nodesDict['Veto_5']['z']['dim'], ' endplate ', nodesDict['T1Endplate_1']['z']['pos'] - nodesDict['T1Endplate_1']['z']['dim']
    #volOutSmall = volInSmall+500.
    t2.Draw("startY_nu:startX_nu","isPrimary==1 && interactionElement==2 && startZ_nu>%s && startZ_nu<%s"%(volInSmall,volOutSmall))
    grSmall = TGraph(t2.GetSelectedRows(), t2.GetV2(), t2.GetV1())
    grSmall.SetMarkerColor(kRed)
    grSmall.SetMarkerStyle(6)
    cSaver.append(TCanvas('whole_decay_volume','whole_decay_volume'))
    grBig.Draw('ap')
    grSmall.Draw('p same')
    grVeto.Draw('p same')
    
    # Save plots
    for c in cSaver:
        c.Print('elena_plots/'+c.GetName()+'.eps')
    
    # Elements of the vacuum tank:
    print
    endprevious = None
    startnew = None
    for i in xrange(1,7):
        print 'T'+str(i)+':'
        if 'T'+str(i)+'Startplate_1' in nodesDict:
            print '\tStartplate (z = %s to %s)'%(nodesDict['T'+str(i)+'Startplate_1']['z']['pos']-nodesDict['T'+str(i)+'Startplate_1']['z']['dim'], nodesDict['T'+str(i)+'Startplate_1']['z']['pos']+nodesDict['T'+str(i)+'Startplate_1']['z']['dim'])
            startnew = nodesDict['T'+str(i)+'Startplate_1']['z']['pos']-nodesDict['T'+str(i)+'Startplate_1']['z']['dim']
        else:
            startnew = nodesDict['T'+str(i)+'I_1']['z']['pos']-nodesDict['T'+str(i)+'I_1']['z']['dim']
        if endprevious:
            print '\tDelta z (%s to %s) = '%(i-1, i) + str(startnew-endprevious)
        print '\tI (z = %s to %s)'%(nodesDict['T'+str(i)+'I_1']['z']['pos']-nodesDict['T'+str(i)+'I_1']['z']['dim'], nodesDict['T'+str(i)+'I_1']['z']['pos']+nodesDict['T'+str(i)+'I_1']['z']['dim'])
        print '\tO (z = %s to %s)'%(nodesDict['T'+str(i)+'O_1']['z']['pos']-nodesDict['T'+str(i)+'O_1']['z']['dim'], nodesDict['T'+str(i)+'O_1']['z']['pos']+nodesDict['T'+str(i)+'O_1']['z']['dim'])
        print '\t\tx diff: ' + str(nodesDict['T'+str(i)+'O_1']['x']['dim']- nodesDict['T'+str(i)+'I_1']['x']['dim'])
        print '\t\ty diff: ' + str(nodesDict['T'+str(i)+'O_1']['y']['dim']- nodesDict['T'+str(i)+'I_1']['y']['dim'])
        print '\t\tr diff: ' + str(nodesDict['T'+str(i)+'O_1']['r']['dim']- nodesDict['T'+str(i)+'I_1']['r']['dim'])
        if 'T'+str(i)+'Endplate_1' in nodesDict:
            print '\tEndplate (z = %s to %s)'%(nodesDict['T'+str(i)+'Endplate_1']['z']['pos']-nodesDict['T'+str(i)+'Endplate_1']['z']['dim'], nodesDict['T'+str(i)+'Endplate_1']['z']['pos']+nodesDict['T'+str(i)+'Endplate_1']['z']['dim'])
            endprevious = nodesDict['T'+str(i)+'Endplate_1']['z']['pos']+nodesDict['T'+str(i)+'Endplate_1']['z']['dim']
        else:
            endprevious = nodesDict['T'+str(i)+'I_1']['z']['pos']+nodesDict['T'+str(i)+'I_1']['z']['dim']
    
    print 'Lenght of the whole decay volume (small + big): ' + str(nodesDict['lidT6I_1']['z']['pos'] + nodesDict['lidT6I_1']['z']['dim'] - (nodesDict['T2Startplate_1']['z']['pos'] - nodesDict['T2Startplate_1']['z']['dim']))