import random

import numpy

import ROOT

import os

ROOT_OBJECTS = []

def set_root_verbosity(verbose_level):

    ROOT.gErrorIgnoreLevel = 6000

def setup_gstyle():

    stops = [0.0000, 0.1250, 0.2500, 0.3750, 0.5000, 0.6250, 0.7500, 0.8750, 1.0000]

    red   = [0.2082, 0.0592, 0.0780, 0.0232, 0.1802, 0.5301, 0.8186, 0.9956, 0.9764]

    green = [0.1664, 0.3599, 0.5041, 0.6419, 0.7178, 0.7492, 0.7328, 0.7862, 0.9832]

    blue  = [0.5293, 0.8684, 0.8385, 0.7914, 0.6425, 0.4662, 0.3499, 0.1968, 0.0539]

    s = numpy.array(stops)

    r = numpy.array(red)

    g = numpy.array(green)

    b = numpy.array(blue)

    ncontours = 255

    npoints = len(s)

    ROOT.TColor.CreateGradientColorTable(npoints, s, r, g, b, ncontours)

    ROOT.gStyle.SetNumberContours(ncontours)

    # axes and labels

    ROOT.gStyle.SetPadBottomMargin(0.15)

    ROOT.gStyle.SetPadLeftMargin(0.15)

    ROOT.gStyle.SetPadRightMargin(0.15)

    for axis in "X", "Y":

        ROOT.gStyle.SetNdivisions(505, axis)

        ROOT.gStyle.SetLabelSize(0.05, axis)

        ROOT.gStyle.SetTitleSize(0.06, axis)

        ROOT.gStyle.SetTitleOffset(1.10, axis)

def get_frame_fill():

    level = 0.9

    return ROOT.TColor.GetColor(0.2082*level, 0.1664*level, 0.5293*level)

def text_box():

    text_box = ROOT.TPaveText(0.6, 0.8, 0.85, 0.899, "NDC")

    text_box.SetBorderSize(0)

    text_box.SetFillColor(0)

    text_box.SetTextSize(0.04)

    text_box.SetTextAlign(12)

    text_box.AddText("MICE preliminary")

    text_box.Draw()

    ROOT_OBJECTS.append(text_box)

    text_box = ROOT.TPaveText(0.6, 0.7, 0.85, 0.8, "NDC")

    text_box.SetBorderSize(0)

    text_box.SetFillColor(0)

    text_box.SetTextSize(0.03)

    text_box.SetTextAlign(12)

    text_box.AddText("ISIS Cycle 2017/03")

    text_box.AddText("Run 7469 MAUS v1.2.3")

    text_box.Draw()

    ROOT_OBJECTS.append(text_box)

def test_hist_2d():

    canvas = ROOT.TCanvas("test 2d", "test 2d")

    hist = ROOT.TH2D("test 2d", ";x axis [units];y axis [units]", 50, -1, 1, 70, -1, 1.8)

    hist.SetStats(False)

    for i in range(100000):

        x, y = random.gauss(0, 0.3), random.gauss(0, 0.3)

        hist.Fill(x, y)

    canvas.Draw()

    # if you want to emphasise the tails, leave the frame white

    hist.Draw("COLZ")

    text_box()

    canvas.Update()

    canvas.Print("test_hist_2d_no_fill.eps")

    canvas.Print("test_hist_2d_no_fill.png")

    # if you want to de-emphasise the tails, use a frame fill

    canvas.SetFrameFillColor(get_frame_fill())

    hist.Draw("COLZ")

    text_box()

    canvas.Update()

    canvas.Print("test_hist_2d_fill.eps")

    canvas.Print("test_hist_2d_fill.png")

    # these two lines force python to keep the ROOT stuff in memory

    ROOT_OBJECTS.append(canvas)

    ROOT_OBJECTS.append(hist)

def test_hist_1d():

    canvas = ROOT.TCanvas("test 1d", "test 1d")

    hist_data = ROOT.TH1D("test 1d data", ";x axis [units];y axis [units]", 100, -1, 1)

    hist_mc = ROOT.TH1D("test 1d mc", ";x axis [units];y axis [units]", 100, -1, 1)

    hist_data.SetStats(False)

    hist_mc.SetStats(False)

    for i in range(100000):

        x_data = random.gauss(0, 0.3)

        hist_data.Fill(x_data)

        x_mc = random.gauss(0.01, 0.27)

        hist_mc.Fill(x_mc)

    canvas.Draw()

    hist_mc.SetFillColor(ROOT.kOrange-2)

    hist_mc.Draw()

    hist_data.SetMarkerStyle(20)

    hist_data.Draw("e1 p same")

    text_box()

    canvas.Update()

    canvas.Print("test_hist_1d.eps")

    canvas.Print("test_hist_1d.png")

    # these two lines force python to keep the ROOT stuff in memory

    ROOT_OBJECTS.append(canvas)

    ROOT_OBJECTS.append(hist_data)

    ROOT_OBJECTS.append(hist_mc)

    # if you are comparing MC graph with data graph, then following style is recommended:

    hist_mc.SetMarkerStyle(26)

    hist_mc.SetMarkerColor(ROOT.kRed)

    hist_mc.SetLineColor(ROOT.kRed)

    hist_mc.Draw("e1 p")

    hist_data.SetMarkerStyle(20)

    hist_data.Draw("e1 p same")

    text_box()

    canvas.Update()

    canvas.Print("test_graph_1d.eps")

    canvas.Print("test_graph_1d.png")

if __name__ == "__main__":

    test_dir = "style_test"

    try:

        os.mkdir(test_dir)

    except OSError:

        pass

    os.chdir(test_dir)

    setup_gstyle()

    test_hist_2d()

    test_hist_1d()

    raw_input("Finished - output in "+test_dir)