from functions import Divideh

import ROOT, argparse, math, sys, os, copy


pandapath=os.environ.get('PANDAPATH')
sys.path.append(pandapath+'/macro/tpc/FOPI/mberger')
from functions import set_palette

parser=argparse.ArgumentParser(description='check pad response')
parser.add_argument('--padf',help='the padplane file to use',type=str,default='tpc/parfiles/padPlane_FOPI.dat')
parser.add_argument('--pshape',help='pad shape file',type=str,default='tpc/parfiles/PadShapePrototype.dat')
parser.add_argument('--padres',help='padresponse file anme',type=str,default='tpc/FOPI/par/padresponse.txt')
parser.add_argument('--old',help='use old padresponse function (gauss int)',action='store_true')

args=parser.parse_args()

ROOT.gROOT.ProcessLine(".x rootlogon.C") 
ROOT.gROOT.ProcessLine("gStyle->SetPalette(1)") 
ROOT.gROOT.LoadMacro("stlPYROOT.h+")

set_palette("name",99)

DigiMapper = []
file = open(args.padf, "r")
lines=file.readlines()
file.close()
lines.reverse()
for line in lines:
    koord = line.split(" ") 
    DigiMapper.append([float(koord[3]), float(koord[4])])


lbin=100
xmin=-1
xmax=1

spreadmult=3

gem=ROOT.TpcGem(2000,0.02)
if args.old:
    padShapePool=ROOT.TpcPadShapePool(args.pshape,gem,0.1,0.005,0.01)
    ltbins=99
else:
    padShapePool=ROOT.TpcPadShapePool(args.pshape,args.padres)
    pref=open(args.padres,'r')
    line=pref.readline()
    ltbins=int(line.split()[0])-1
    pref.close()
padplane=ROOT.TpcPadPlane(args.padf,padShapePool)

lookuptable=ROOT.TH2D("lookup","lookup",lbin,DigiMapper[200][0]+xmin,DigiMapper[200][0]+xmax,
                      lbin,DigiMapper[200][1]+xmin,DigiMapper[200][1]+xmax)


lookuptable2=ROOT.TH2D("lookup2","lookup2;Xpos (mm);Ypos (mm)",
                       ltbins,
                       -ltbins/2*0.125,ltbins/2*0.125,
                       ltbins,
                       -ltbins/2*0.125,ltbins/2*0.125)

weighttable=ROOT.TH2D("weight","weighttable",
                       ltbins,
                       0,ltbins,
                       ltbins,
                       0,ltbins)
weight1D=ROOT.TH1D("weight1D",'weight1D',10001,0,1)
weight1D2=ROOT.TH1D("weight1D2",'weight1D2',10001,0,1)

#print 'value=',padplane.GetPad(200).GetValue(DigiMapper[200][0],DigiMapper[200][1]-0.8)
print 'weight=',padplane.GetPad(200).GetWeight(0.000652694176518,-1.00000020546,-1.73205068895)
#exit()

print "getting lookuptable"
for i in range(ltbins+1):
    for j in range(ltbins+1):
        lookuptable2.SetBinContent(i,j,padplane.GetPad(20).GetLTValue(i,j))

c3=ROOT.TCanvas('clookup2','lookup2',500,500)
lookuptable2.Draw('contz fb')
#lookuptable2.Draw('cont1 same')
c3.Update()

xpos=DigiMapper[200][0]+xmin
step=abs((xmax-xmin)/float(lbin)) 

xbin=1

print DigiMapper[200][0],DigiMapper[200][1]
print DigiMapper[200][0]+xmax,DigiMapper[200][1]+xmax

print "getting lookuptable again (rotated)"
while xpos<=DigiMapper[200][0]+xmax:
    ypos=DigiMapper[200][1]+xmin
    ybin=1
    while ypos<=DigiMapper[200][1]+xmax:
        '''
        if padplane.GetPad(200).Contains(xpos,ypos):
            lookuptable.SetBinContent(xbin,ybin,1)
        else:
            lookuptable.SetBinContent(xbin,ybin,0)
           ''' 
        #print xpos,ypos
        lookuptable.SetBinContent(xbin,ybin,padplane.GetPad(200).GetValue(xpos,ypos))
        xdir=0-(xpos-DigiMapper[200][0])
        ydir=0-(ypos-DigiMapper[200][1])
        xdir/=math.sqrt(xdir**2+ydir**2)
        ydir/=math.sqrt(xdir**2+ydir**2)
        #print "bin:",xbin,ybin
        #weighttable.SetBinContent(xbin,ybin,padplane.GetPad(200).GetWeight(0.5,xdir,ydir))

        ypos+=step
        ybin+=1
    xpos+=step
    xbin+=1

c2=ROOT.TCanvas("lookup","lookup",500,500)
lookuptable.Draw('colz')
padplane.GetPad(200).Draw(1)
c2.Update()

amp=0
ampstep=1./10000.

print "getting weighting"
for i in range(0,10000):

    ydir=0.
    xdir=1.
    weight1D.Fill(amp,padplane.GetPad(200).GetWeight(amp,xdir,ydir))

    
    ydir=1.
    xdir=0.
    weight1D2.Fill(amp,padplane.GetPad(200).GetWeight(amp,xdir,ydir))

    if abs(amp-0.5)<1e-5:
        print 'amp=',amp,':'
        print 'along x:',padplane.GetPad(200).GetWeight(amp,1,0)
        print 'along y:',padplane.GetPad(200).GetWeight(amp,0,1)

    amp+=ampstep

c4=ROOT.TCanvas("weight","weight",1000,500)
c4.Divide(2,1)
c4.cd(1)
weighttable.Draw('colz')
c4.cd(2)
weight1D.Draw()
weight1D2.SetLineColor(2)
weight1D2.Draw('same')
padplane.GetPad(200).Draw(1)
c4.Update()

c1=ROOT.TCanvas()
u='q'
while u!='q': 

    midX=ROOT.gRandom.Uniform(-15,15)
    midY=ROOT.gRandom.Uniform(-15,15)
    avpos=ROOT.TVector3(midX,midY,0)
    hitpads=[]

    numsigs=ROOT.gRandom.Uniform(1,10)
    print "getting padlist"
    hitpads=padplane.GetPadList(midX,midY,gem.spread()*spreadmult)
    if len(hitpads)==0:
        continue

    sumamp=0
    meanpos=ROOT.TVector3(0,0,0)
    meanpos_unweighted=ROOT.TVector3(0,0,0)
    print "getting amps"
    for pad in hitpads:
        amp=padplane.GetPad(pad).GetValue(midX,midY)
        sumamp+=amp

        print 'padamp=',amp,'(',DigiMapper[pad][0],DigiMapper[pad][1],')'
        meanpos.SetX(meanpos.X()+amp*DigiMapper[pad][0])
        meanpos.SetY(meanpos.Y()+amp*DigiMapper[pad][1])

        meanpos_unweighted.SetX(meanpos_unweighted.X()+DigiMapper[pad][0])
        meanpos_unweighted.SetY(meanpos_unweighted.Y()+DigiMapper[pad][1])

    if sumamp==0:
        continue
        
    print 'sumamp=',sumamp
    meanpos*=1./sumamp
    print 'Weighted:'
    print 'POS: (',meanpos.X(),',',meanpos.Y(),')'
    print 'residual: (',(meanpos-avpos).X(),",",(meanpos-avpos).Y(),')'

    meanpos_unweighted*=1./len(hitpads)
    print 'Unweighted'
    print 'POS: (',meanpos_unweighted.X(),',',meanpos_unweighted.Y(),')'
    print 'residual (unweighted): (',(meanpos_unweighted-avpos).X(),',',(meanpos_unweighted-avpos).X(),')'

    #calculate reweighted pos
    meanpos_reweighted=ROOT.TVector3(0,0,0)
    sumamp_re=0
    weights=[]
    weightsum=0
    '''
    for pad in hitpads:
        amp=padplane.GetPad(pad).GetValue(midX,midY)

        xdir=DigiMapper[pad][0]-meanpos.X()
        ydir=DigiMapper[pad][1]-meanpos.Y()
        print 'get weight',xdir,ydir
        weights.append(padplane.GetPad(200).GetWeight(amp/sumamp,xdir,ydir))
        print 'got weight:',weights[-1]
        weightsum+=weights[-1]
        '''
    pcount=-1
    weightampsum=0
   # meanpos_reweighted.SetXYZ(meanpos_unweighted.X(),meanpos_unweighted.Y(),0)
    print 'meanpos_reweighted=(',meanpos_reweighted.X(),',',meanpos_reweighted.Y(),')'
    print "reweighting"
    for pad in hitpads:
        pcount+=1
        #thisweight=1./weights[pcount]#/weightsum
        amp=padplane.GetPad(pad).GetValue(midX,midY)
        
        #thisweight=1-thisweight
        xdir=meanpos_unweighted.X()-DigiMapper[pad][0]
        ydir=meanpos_unweighted.Y()-DigiMapper[pad][1]
        print "getting weight for",amp/sumamp,xdir,ydir
        weight=padplane.GetPad(pad).GetWeight(amp/sumamp,xdir,ydir)
        weightampsum+=abs(weight)

        print 'amp=',amp,'weight=',weight,'xdir=',xdir,'ydir=',ydir
        print 'padpos: (',DigiMapper[pad][0],',',DigiMapper[pad][1],')'
        
        #meanpos_reweighted.SetX(meanpos_reweighted.X()+(weight*xdir))
        #meanpos_reweighted.SetY(meanpos_reweighted.Y()+(weight*ydir))

        meanpos_reweighted.SetX(meanpos_reweighted.X()+amp*(weight*xdir+meanpos_unweighted.X()))
        meanpos_reweighted.SetY(meanpos_reweighted.Y()+amp*(weight*ydir+meanpos_unweighted.Y()))
        
        
        print 'reweightes pos: X=',meanpos_reweighted.X()," Y=",meanpos_reweighted.Y()
        print ''
        #weightampsum+=thisweight

    meanpos_reweighted*=1./sumamp
    print 'final reweightes pos: X=',meanpos_reweighted.X()," Y=",meanpos_reweighted.Y()
    circ=ROOT.TEllipse(midX,midY,gem.spread()*spreadmult)
    circ.SetFillStyle(4000)
    
    circ_unw=ROOT.TEllipse(meanpos_unweighted.X(),meanpos_unweighted.Y(),gem.spread())
    circ_unw.SetLineColor(3)
    circ_unw.SetFillStyle(4000)
    

    hist=ROOT.TH2D("pad","pad",300,-15,15,300,-15,15)
    hist.GetXaxis().SetRangeUser(midX-0.5,midX+0.5)
    hist.GetYaxis().SetRangeUser(midY-0.5,midY+0.5)
    hist.Draw()

    avpos=ROOT.TPolyMarker()
    avpos.SetPoint(0,midX,midY)
    avpos.SetMarkerStyle(7)
    avpos.SetMarkerColor(1)
    avpos.Draw('same')

    meanpospos=ROOT.TPolyMarker()
    meanpospos.SetPoint(0,meanpos.X(),meanpos.Y())
    meanpospos.SetMarkerStyle(7)
    meanpospos.SetMarkerColor(2)
    meanpospos.Draw('same')

    meanpospos_unw=ROOT.TPolyMarker()
    meanpospos_unw.SetPoint(0,meanpos_unweighted.X(),meanpos_unweighted.Y())
    meanpospos_unw.SetMarkerStyle(7)
    meanpospos_unw.SetMarkerColor(3)
    meanpospos_unw.Draw('same')

    meanpospos_rew=ROOT.TPolyMarker()
    meanpospos_rew.SetPoint(0,meanpos_reweighted.X(),meanpos_reweighted.Y())
    meanpospos_rew.SetMarkerStyle(7)
    meanpospos_rew.SetMarkerColor(4)
    meanpospos_rew.Draw('same')


    for pad in hitpads:
        print pad
        padplane.GetPad(pad).Draw(2)


    circ.Draw()
    circ_unw.Draw()
    c1.Update()
    u=raw_input();

u=raw_input()