#!/usr/bin/env python
# coding: utf-8

# <img src="http://oproject.org/img/ROOT.png" height="30%" width="30%">
# <img src="http://oproject.org/img/tmvalogo.png" height="30%" width="30%">

# <hr style="border-top-width: 4px; border-top-color: #34609b;">

# In[ ]:


from ROOT import TMVA, TFile, TTree, TCut, TString, TH1F, TCanvas, gStyle
gStyle.SetOptStat(0)
get_ipython().run_line_magic('jsroot', 'on')
from subprocess import call
from os.path import isfile
from array import array

from keras.models import Sequential
from keras.layers import Dense


# In[ ]:


TMVA.Tools.Instance()
TMVA.PyMethodBase.PyInitialize()

output = TFile.Open('TMVA.root', 'RECREATE')
factory = TMVA.Factory('TMVARegression', output,
        '!V:!Silent:Color:!DrawProgressBar:Transformations=D,G:AnalysisType=Regression')


# In[ ]:


if not isfile('tmva_reg_example.root'):
    call(['curl', '-O', 'http://root.cern.ch/files/tmva_reg_example.root'])

data = TFile.Open('tmva_reg_example.root')
tree = data.Get('TreeR')

dataloader = TMVA.DataLoader('dataset')
for branch in tree.GetListOfBranches():
    name = branch.GetName()
    if name != 'fvalue':
        dataloader.AddVariable(name)
dataloader.AddTarget('fvalue')

dataloader.AddRegressionTree(tree, 1.0)
dataloader.PrepareTrainingAndTestTree(TCut(''),
        'nTrain_Regression=4000:SplitMode=Random:NormMode=NumEvents:!V')


# In[ ]:


model = Sequential()
model.add(Dense(64, activation='tanh', input_dim=2))
model.add(Dense(1, activation='linear'))

model.summary()


# In[ ]:


model.compile(loss='mean_squared_error', optimizer="sgd")

model.save("model.h5")


# In[ ]:


factory.BookMethod(dataloader, TMVA.Types.kPyKeras, 'PyKeras',
        'H:!V:VarTransform=D,G:FilenameModel=model.h5:NumEpochs=20:BatchSize=32')
factory.BookMethod(dataloader, TMVA.Types.kBDT, 'BDTG',
        '!H:!V:NTrees=1000:nCuts=20:MaxDepth=3')


# In[ ]:


factory.TrainAllMethods()
factory.TestAllMethods()
factory.EvaluateAllMethods()


# In[ ]:


reader = TMVA.Reader("Color:!Silent")


# In[ ]:


branches = {}
for branch in tree.GetListOfBranches():
    branchName = branch.GetName()
    branches[branchName] = array('f', [-999])
    tree.SetBranchAddress(branchName, branches[branchName])
    if branchName != 'fvalue':
        reader.AddVariable(branchName, branches[branchName])


# In[ ]:


reader.BookMVA('PyKeras', TString('dataset/weights/TMVARegression_PyKeras.weights.xml'))
reader.BookMVA('BDTG', TString('dataset/weights/TMVARegression_BDTG.weights.xml'))


# In[ ]:


c = TCanvas()
h_keras = TH1F("PyKeras", "PyKeras", 30, 0, 400)
h_bdtg = TH1F("BDTG", "BDTG", 30, 0, 400)
h_truth = TH1F("Ground truth", "Ground truth", 30, 0, 400)
for i in range(tree.GetEntries()):
    tree.GetEntry(i)
    h_keras.Fill(reader.EvaluateMVA("PyKeras"))
    h_bdtg.Fill(reader.EvaluateMVA("BDTG"))
    h_truth.Fill(branches["fvalue"][0])


# In[ ]:


h_bdtg.SetTitle("BDTG vs PyKeras")
h_bdtg.GetXaxis().SetTitle("Target")
h_bdtg.Draw()
h_keras.Draw("SAME")
h_truth.Draw("SAME")
c.Draw()