#include <iostream>
#include <vector>
#include "TFile.h"
#include "TTree.h"

#include "TMVA/Reader.h"

using namespace std;

int main()
{

    //-----------------------------------
    // 1 创建 TMVA Reader
    //-----------------------------------
    TFile *f = TFile::Open("tmva_output.root");

    TMVA::Reader *reader = new TMVA::Reader("!Color:!Silent");

    //-----------------------------------
    // 2 定义输入变量
    //-----------------------------------

    float det[10];

    for(int i=0;i<10;i++)
    {
        reader->AddVariable(Form("det%d",i), &det[i]);
    }

    //-----------------------------------
    // 3 加载训练好的模型
    //-----------------------------------

    reader->BookMVA(
        "DNN",
        "dataset/weights/NeutronSpectrum_DNN.weights.xml"
    );

    //-----------------------------------
    // 4 输入数据
    //-----------------------------------

    det[0]=15;
    det[1]=22;
    det[2]=10;
    det[3]=18;
    det[4]=5;
    det[5]=7;
    det[6]=9;
    det[7]=11;
    det[8]=6;
    det[9]=8;

    //-----------------------------------
    // 5 进行预测
    //-----------------------------------

    const std::vector<float> &result =
        reader->EvaluateRegression("DNN");

    //-----------------------------------
    // 6 输出结果
    //-----------------------------------

    cout<<"Prediction:"<<endl;

    for(size_t i=0;i<result.size();i++)
    {
        cout<<"spec"<<i<<" = "<<result[i]<<endl;
    }

}