using namespace std;
using namespace RooFit;
using namespace RooStats;

void GetObjectCount(TString fileName)
{
    TFile* inFile = new TFile(fileName,"update");
    RooWorkspace* wksp = (RooWorkspace*)inFile->Get("wksp");
    cout << "Getting objectCount from RooRealVar: " << (int)wksp->var("objectCount")->getVal() << endl;
    TProcessID::SetObjectCount((int)wksp->var("objectCount")->getVal());
    wksp->Write();
    inFile->Close();
}

void SetObjectCount(TString fileName)
{
    TFile* inFile = new TFile(fileName,"update");
    RooWorkspace* wksp = (RooWorkspace*)inFile->Get("wksp");
    cout << "Setting objectCount RooRealVar to " << TProcessID::GetObjectCount() << endl;
    wksp->var("objectCount")->setVal(TProcessID::GetObjectCount());
    wksp->Write();
    inFile->Close();
}

void CreateWorkspace(TString fileName)
{
    RooWorkspace* wksp = new RooWorkspace("wksp");
    wksp->factory("x[0.,1.]");
    wksp->factory("objectCount[0,0,1e6]");

    TFile* outFile = new TFile(fileName,"recreate");
    wksp->Write();
    outFile->Close();

    SetObjectCount(fileName);
}

void CreatePdf(TString fileName)
{
    GetObjectCount(fileName);

    TFile* inFile = new TFile(fileName,"read");
    RooWorkspace* wksp = (RooWorkspace*)inFile->Get("wksp")->Clone();
    inFile->Close();

    wksp->factory("SUM::bkgPlusSigModel(prod::bkgSize(bu[1.,0.,2.],bkgN[100000])*Exponential::bkgModel(x,k[-8,-14,-3]), prod::sigSize(mu[0.2,-5.,5.],sigN[1000])*BreitWigner::sigModel(x,m[0.25,0.1,0.9],w[0.03]))");

    TFile* outFile = new TFile(fileName,"recreate");
    wksp->Write();
    outFile->Close();

    SetObjectCount(fileName);
}

void CreateModel(TString fileName)
{
    GetObjectCount(fileName);

    TFile* inFile = new TFile(fileName,"read");
    RooWorkspace* wksp = (RooWorkspace*)inFile->Get("wksp")->Clone();
    inFile->Close();

    ModelConfig* model = new ModelConfig("model",wksp);
    if(wksp->var("x"))               model->SetObservables("x");
    if(wksp->pdf("bkgPlusSigModel")) model->SetPdf("bkgPlusSigModel");
    wksp->import(*model);

    TFile* outFile = new TFile(fileName,"recreate");
    outFile->cd();
    wksp->Write();
    outFile->Close();

    SetObjectCount(fileName);
}

void PrintModel(TString fileName)
{
    TFile* inFile = new TFile(fileName,"read");
    RooWorkspace* wksp = (RooWorkspace*)inFile->Get("wksp")->Clone();
    if(wksp->obj("model")) wksp->obj("model")->Print();
    inFile->Close();
}

void createThings_fix(std::string arg)
{
    TString mode       = TString(arg); //wksp, pdf, model, any combination
    TString fileName = "workspace.root";

    if(mode.Contains("workspace")) CreateWorkspace(fileName);
    if(mode.Contains("pdf"))       CreatePdf(fileName);
    if(mode.Contains("model"))     CreateModel(fileName);

    PrintModel(fileName);
    exit(0);
}