import sys, os, copy
from math import sqrt
pandapath = os.environ.get('PANDAPATH')
sys.path.append(pandapath + '/macro/tpc/FOPI/python/argparse-1.2.1')
sys.path.append(pandapath + '/macro/tpc/FOPI/mberger')
import argparse, ROOT
from functions import set_palette, openTree, thisIsTheEnd, getSlices
import random


ROOT.gROOT.ProcessLine(".x rootlogon.C") 
ROOT.gROOT.ProcessLine('gROOT->SetStyle("Plain")')
ROOT.gROOT.LoadMacro("stlPYROOT.h+")
set_palette()


hpoints=ROOT.TH2D("hpointsVsZ","points pos vs Z",200,0,200,1000,-1,1)
hpoints_binned=ROOT.TH2D("hpointsVsZ_binned","points pos vs Z (binned)",200,0,200,1000,-1,1)
hmean=ROOT.TH2D("hmeanVsZ","Mean pos vs Z",200,0,200,1000,-1,1)
hmean_binned=ROOT.TH2D("hmeanVsZ_binned","Mean pos vs Z (binned)",200,0,200,100,-1,1)

nbins=10
binwidth=2./float(nbins)
ndigis=10
for i in range(1000):
    for z in range(0,200):
        #generate gauss distribution with small number of entries
        points=[]
        points_binned=[]
        for i in range(ndigis):
            points.append(random.gauss(0,sqrt(z)*0.02))
            hpoints.Fill(z,points[-1])
            #print z,points[-1],int(float(points[-1])/float(binwidth)),binwidth
            points_binned.append(binwidth*int(float(points[-1])/float(binwidth)))
            hpoints_binned.Fill(z,points_binned[-1])
        #calculate mean
        mean=sum(points)/len(points)
        hmean.Fill(z,mean)
        mean_binned=sum(points_binned)/len(points_binned)
        hmean_binned.Fill(z,mean_binned)

c1=ROOT.TCanvas("c1","C1",1500,3000)
c1.Divide(4,2)
c1.cd(1)
hpoints.Draw('colz')
hpoints.FitSlicesY(0,0,-1,0,"QNRG3S")
hpoints_sigma=ROOT.gDirectory.Get("hpointsVsZ_2")
c1.cd(5)
hpoints_sigma.Draw()

c1.cd(2)
hmean.Draw('colz')
hmean.FitSlicesY(0,0,-1,0,"QNRG3S")
hmean_sigma=ROOT.gDirectory.Get(hmean.GetName()+"_2")
c1.cd(6)
hmean_sigma.Draw()

c1.cd(3)
hpoints_binned.Draw('colz')
hpoints_binned.FitSlicesY(0,0,-1,0,"QNRG3S")
hpoints_binned_sigma=ROOT.gDirectory.Get("hpointsVsZ_binned_2")
c1.cd(7)
hpoints_binned_sigma.Draw()

c1.cd(4)
hmean_binned.Draw('colz')
hmean_binned.FitSlicesY(0,0,-1,0,"QNRG3S")
hmean_binned_sigma=ROOT.gDirectory.Get(hmean_binned.GetName()+"_2")
c1.cd(8)
hmean_binned_sigma.Draw()

thisIsTheEnd()