#include <cstdlib>
#include <iostream>
#include <map>
#include <string>

#include "TChain.h"
#include "TFile.h"
#include "TTree.h"
#include "TString.h"
#include "TObjString.h"
#include "TSystem.h"
#include "TROOT.h"

#include "TMVA/Tools.h"
#include "TMVA/Factory.h"
#include "TMVA/DataLoader.h"
#include "TMVA/TMVARegGui.h"


using namespace TMVA;

void DNNRegression(){
//定义输出文件
    TFile *outfile=new TFile("EnergyConstruction.root","RECREATE");
//初始化多变量分析器
    TMVA::Factory *factory = new TMVA::Factory( "TMVARegression", outfile,
                                               "!V:!Silent:Color:DrawProgressBar:AnalysisType=Regression" );


//利用DataLoader加载训练与测试数据
    TMVA::DataLoader *dataloader=new TMVA::DataLoader("dataset");

    for(int i=0;i<64;i++){
        TString PE,PMTNUM;
        PE.Form("PE[%d]",i);
        PMTNUM.Form("PMT%d",i);
        dataloader->AddVariable( PE, PMTNUM, "units", 'F' );
    }
    dataloader->AddTarget( "Energy" );

    TFile *input=TFile::Open("NeuralNetworksData.root","READ");

    TTree *regTree=(TTree *)input->Get("event");
    dataloader->AddRegressionTree( regTree, 1.0 );

    TCut mycut = "Energy >20 && Energy<50";

    dataloader->PrepareTrainingAndTestTree( mycut,
                                         "nTrain_Regression=40000:nTest_Regression=5000:SplitMode=Random:NormMode=NumEvents:!V" );
//定义神经网络架构，隐含层、单元数、激活函数、学习策略
    TString training0("LearningRate=1e-5,Momentum=0.5,Repetitions=1,ConvergenceSteps=100,BatchSize=40,"
                        "TestRepetitions=10,WeightDecay=0.01,Regularization=L1,DropConfig=0");
    TString training1("LearningRate=1e-5,Momentum=0.5,Repetitions=1,ConvergenceSteps=100,BatchSize=40,"
                        "TestRepetitions=5,WeightDecay=0.01,Regularization=L1,DropConfig=0");
    TString training2("LearningRate=1e-5,Momentum=0.5,Repetitions=1,ConvergenceSteps=100,BatchSize=40,"
                        "TestRepetitions=5,WeightDecay=0.01,Regularization=L1");


    TString trainingStrategyString("TrainingStrategy=");

   // trainingStrategyString+="LearningRate=1e-1, BatchSize=256| LearningRate=1e-2, BatchSize=256| LearningRate=1e-3, BatchSize=256";

    trainingStrategyString += training0 + "|" + training1+ "|" + training2+"|" + training3+"|" + training4+"|" + training5;
    TString layoutString("Layout=SIGMOID|60,Layout=SIGMOID|60,Layout=RELU|60,RELU");

    TString nnOptions(
         "!H:V:ErrorStrategy=SUMOFSQUARES:VarTransform=G:WeightInitialization=XAVIER:Architecture=CPU");
      nnOptions.Append(":");
      nnOptions.Append(layoutString);
      nnOptions.Append(":");
      nnOptions.Append(trainingStrategyString);

    factory->BookMethod(dataloader, TMVA::Types::kDNN, "DNN_CPU", nnOptions);

    // factory->BookMethod( dataloader,  TMVA::Types::kBDT, "BDT",
    //                        "!H:!V:NTrees=200:MinNodeSize=1.0%:BoostType=AdaBoostR2:SeparationType=RegressionVariance:nCuts=20:PruneMethod=CostComplexity:PruneStrength=30" );

    factory->TrainAllMethods();

    factory->TestAllMethods();

    factory->EvaluateAllMethods();

    outfile->Close();

   std::cout << "==> Wrote root file: " << outfile->GetName() << std::endl;
   std::cout << "==> TMVARegression is done!" << std::endl;

   delete factory;
   delete dataloader;

}

int main(){
    DNNRegression();
    return 0;

}