import os

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 ""

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

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