using namespace TMVA::Experimental;

/// function to compile the generated model and the declaration of the SofieFunctor
/// used by RDF.
/// Assume that the model name as in the header file
void CompileModelForRDF(const std::string & headerModelFile, unsigned int ninputs, unsigned int nslots=0) {

   std::string modelName = headerModelFile.substr(0,headerModelFile.find(".hxx"));
   std::string cmd = std::string("#include \"") + headerModelFile + std::string("\"");
   auto ret = gInterpreter->Declare(cmd.c_str());
   if (!ret)
    throw std::runtime_error("Error compiling : " + cmd);
   std::cout << "compiled : " << cmd << std::endl;

   cmd = "auto sofie_functor = TMVA::Experimental::SofieFunctor<" + std::to_string(ninputs) + ",TMVA_SOFIE_" +
    modelName + "::Session>(" + std::to_string(nslots) + ");";
   ret = gInterpreter->Declare(cmd.c_str());
   if (!ret)
    throw std::runtime_error("Error compiling : " + cmd);
   std::cout << "compiled : " << cmd << std::endl;
   std::cout << "Model is ready to be evaluated" << std::endl;
   return;
}

void TMVA_SOFIE_RDataFrame_JIT(std::string modelName = "HiggsModel"){

    // check if the input file exists
    std::string modelHeaderFile = modelName + ".hxx";
    if (gSystem->AccessPathName(modelHeaderFile.c_str())) {
        Info("TMVA_SOFIE_RDataFrame","You need to run TMVA_SOFIE_Keras_Higgs_Model.py to generate the SOFIE header for the Keras trained model");
        return;
    }

    // check that also weigh file exists
    std::string modelWeightFile = modelName + std::string(".dat");
    if (gSystem->AccessPathName(modelWeightFile.c_str())) {
        Error("TMVA_SOFIE_RDataFrame","Generated weight file is missing");
        return;
    }

    // now compile using ROOT JIT trained model (see function above)
    CompileModelForRDF(modelHeaderFile,7);

    std::string inputFileName = "Higgs_data.root";
    std::string inputFile = std::string{gROOT->GetTutorialDir()} + "/machine_learning/data/" + inputFileName;

    ROOT::RDataFrame df1("sig_tree", inputFile);
    auto h1 = df1.Define("DNN_Value", "sofie_functor(rdfslot_,m_jj, m_jjj, m_lv, m_jlv, m_bb, m_wbb, m_wwbb)")
                   .Histo1D({"h_sig", "", 100, 0, 1},"DNN_Value");

    ROOT::RDataFrame df2("bkg_tree", inputFile);
    auto h2 = df2.Define("DNN_Value", "sofie_functor(rdfslot_,m_jj, m_jjj, m_lv, m_jlv, m_bb, m_wbb, m_wwbb)")
                   .Histo1D({"h_bkg", "", 100, 0, 1},"DNN_Value");

    h1->SetLineColor(kRed);
    h2->SetLineColor(kBlue);

    auto c1 = new TCanvas();
    gStyle->SetOptStat(0);

    h2->DrawClone();
    h1->DrawClone("SAME");
    c1->BuildLegend();


}