#!/usr/bin/env python

# ----------------------------------------------------------------------------------- #
#  Python/pyROOT macro for generating a simple dataset.
#  Execute in shell, which produces two output files:
#    - SimpleDataset.root     (the data in ROOT Ntuple format)
#    - SimpleDataset.txt      (the same data in ascii format)
#
#  Authors: Lars Egholm Pedersen & Troels Petersen (NBI)
#  Date: 7th of December 2016
#
# ----------------------------------------------------------------------------------- #

import ROOT

# Change output name to create independent datasets for testing.
# Don't forget to change seed (or set 0) the generator seed.

outprefix = "./SimpleDataset"
npoint = 10000
seed = 1

ROOT.gRandom.SetSeed( seed )

Hist_sig = ROOT.TH2F("Hist_sig", "Signal distribution",     100, -5.0, 5.0, 100, -5.0, 5.0)
Hist_bkg = ROOT.TH2F("Hist_bkg", "Background distribution", 100, -5.0, 5.0, 100, -5.0, 5.0)



# ----------------------------------------------------------------------------------- #
# Generate data files:
# ----------------------------------------------------------------------------------- #

# Create a TNtuple:
ntuple = ROOT.TNtuple("ntuple", "ntuple", "x:y:isSignal")

# Open Ascii output (with open as ... automatically closes file when exiting block)
with open( outprefix+".txt", "w" ) as outtxt : 

    # ROOT/TMVA needs to know the data format of the txt file if it is to be used as input,
    # so this has to be written as a header of the .txt file:
    outtxt.write("x:y:isSignal\n")

    # Generate 'signal' and 'background' distributions:
    for i in xrange( npoint ):
        if (i % 1000 == 0) : print "  Number of generated points: ", i

        # Generate signal
        data = [ROOT.gRandom.Gaus(1,1), # x
                ROOT.gRandom.Gaus(1,1), # y
                1 ]                     # signal

        outtxt.write( "%8.4f \t %8.4f \t %d \n"%(data[0], data[1], data[2]) ) # Write to .txt file 
        ntuple.Fill( data[0], data[1], data[2] )                              # Save in ROOT tuple
        Hist_sig.Fill(data[0], data[1])

        # Generate background
        data = [ROOT.gRandom.Gaus(-1,1), # x
                ROOT.gRandom.Gaus(-1,1), # y
                0 ]                     # signal

        outtxt.write( "%8.4f \t %8.4f \t %d \n"%(data[0], data[1], data[2]) ) # Write to .txt file 
        ntuple.Fill( data[0], data[1], data[2] )                              # Save in ROOT tuple
        Hist_bkg.Fill(data[0], data[1])


# Save output
outroot = ROOT.TFile( outprefix+".root", "RECREATE")
ntuple.Write()
outroot.Close()



# ----------------------------------------------------------------------------------- #
#  Plot generated data:
# ----------------------------------------------------------------------------------- #

canvas = ROOT.TCanvas("canvas", "Signal and Background distributions", 20, 20, 800, 600)

Hist_sig.SetMarkerSize(0.5)
Hist_sig.SetMarkerColor(ROOT.kRed)
Hist_sig.SetMarkerStyle(20)
Hist_sig.Draw("")

Hist_bkg.SetMarkerSize(0.5)
Hist_bkg.SetMarkerColor(ROOT.kBlue)
Hist_bkg.SetMarkerStyle(24)
Hist_bkg.Draw("same")

canvas.Update()



raw_input( ' Press enter to exit ' )