#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Feb  3 10:38:31 2023

@author: yannickburkard
"""

from ROOT import *
import numpy as np


#creating histograms

NData = np.array([29, 49, 84, 127, 157, 135, 120, 77, 54, 49, 43, 32, 25, 26, 17, 51])
NSM =   np.array([28, 46, 89, 133, 141, 122,  99, 72, 56, 41, 34, 30, 22, 21, 19, 43])

edges = np.linspace(50,450,17)
centers=(edges[1:]+edges[:-1])/2

distData = np.zeros(sum(NData))
distSM = np.zeros(sum(NSM))

#filling distributions
count= 0
for i in range(len(NData)):
    for j in range(NData[i]):
        distData[count]=centers[i]
        count+=1
count= 0
for i in range(len(NSM)):
    for j in range(NSM[i]):
        distSM[count]=centers[i]
        count+=1

histData = TH1D("", "", 16, 50, 450)
histSM = TH1D("", "", 16, 50, 450)

for x in distData:
    histData.Fill(x)
for x in distSM:
    histSM.Fill(x)


#---------------------------------------------------------------------------------------------

#plotting histograms

canvas = TCanvas("Histogram", "Histogram", 100, 200, 700, 600)
pad1 = TPad("pad1", "pad1", 0, 0.3, 1, 1.0)
pad2 = TPad("pad2", "pad2", 0, 0.05, 1, 0.3)

canvas.cd()
pad1.SetBottomMargin(0)
pad1.Draw()
pad1.cd()

#customizing the histograms
histData.SetLineColor(kBlue)
histData.SetLineWidth(2)
histData.SetLineStyle(1)
histSM.SetLineColor(kGreen)
histSM.SetFillColor(kGreen)

histSM.Draw("hist")
histSM.GetYaxis().SetTitle("Frequency")
histSM.SetStats(0)
histData.Draw("same E1P")

#customizing axes
histSM.GetYaxis().SetTitleSize(15)
histSM.GetYaxis().SetTitleFont(43)
histSM.GetYaxis().SetLabelFont(43)
histSM.GetYaxis().SetLabelSize(15)
histSM.SetMaximum(170)

pad1.SetGridx()
pad1.SetGridy()

ratiosSM=histData.Clone("ratiosSM")
ratiosSM.Divide(histSM)

canvas.cd()
pad2.SetTopMargin(0)
pad2.SetBottomMargin(0.25)
pad2.SetGridx()
pad2.SetGridy()
pad2.Draw()
pad2.cd()

#yburkard: calling "hist" here plots bars starting from 0, and my idea would be to make them start from 1
ratiosSM.Draw("hist") 

ratiosSM.SetTitle("")
ratiosSM.SetStats(0)

ratiosSM.GetYaxis().SetTitle("Data/SM")
ratiosSM.GetXaxis().SetTitle("m")

#customizing histograms
ratiosSM.SetLineWidth(2)
ratiosSM.SetMarkerStyle(8)
ratiosSM.SetMarkerSize(0.5)
ratiosSM.SetMarkerColor(kGray)
ratiosSM.SetLineColor(kGray)
ratiosSM.SetFillColor(kGray)

#customizing axes
ratiosSM.GetYaxis().SetNdivisions(505)
ratiosSM.GetYaxis().SetTitleSize(15)
ratiosSM.GetYaxis().SetTitleFont(43)
ratiosSM.GetYaxis().SetLabelFont(43)
ratiosSM.GetYaxis().SetLabelSize(15)
ratiosSM.GetYaxis().SetTitleOffset(1.2)
ratiosSM.GetYaxis().SetRangeUser(0,2)
ratiosSM.GetXaxis().SetTitleSize(16)
ratiosSM.GetXaxis().SetTitleFont(43)
ratiosSM.GetXaxis().SetTitleOffset(1)
ratiosSM.GetXaxis().SetLabelFont(43)
ratiosSM.GetXaxis().SetLabelSize(15)

output_path = "example_histogram.png"
canvas.SaveAs(output_path)

pad1.Clear()
pad2.Clear()
canvas.Clear("D")