#include <iostream>
#include <fstream>
#include <TMath.h>
using namespace std;
using namespace TMath;
using namespace RooFit;

void Fit(){
  TFile *f = new TFile("Example1.root");
  TTree *t = (TTree*)f->Get("data");
  RooRealVar En("x","x",0,100);
  RooDataSet data("data","data",RooArgSet(En),Import(*t));

  //signal gaussian at 50keV with sigma of 5
  RooRealVar mean("mean","mean",50);
  RooRealVar sig("sig","sig",5);
  RooGaussian gss("gss","gss",En,mean,sig);

  //uniform background
  RooPolynomial p0("px","px",En);

  //sigy and bkgy are yields for each component, only variables allowed to vary
  RooRealVar sigy("sigy","sigy",1,-5,5);
  RooRealVar bkgy("bkgy","bkgy",1,0,9000);

  RooFormulaVar sigNorm("sigNorm","1./(x[0]*TMath::Sqrt(2*TMath::Pi()))",{sig});
  RooFormulaVar bkgNorm("bkgNorm","1./100.",{});
  RooFormulaVar totalYield("totalYield","x[0]+x[1]",{sigy,bkgy});
  RooFormulaVar coeff("coeff","x[0]/TMath::Abs(x[0]+x[1])",{sigy,bkgy});

  RooGenericPdf modelNE("ModelNonExtended",
		     "x[5] * (1-x[3]) + x[4] * x[3] * TMath::Exp(-0.5*((x[0]-x[1])*(x[0]-x[1]))/(x[2] * x[2]))",
		     {En,mean,sig,coeff,sigNorm,bkgNorm});
  RooExtendPdf model("model","model",modelNE,totalYield);
		     
  

  //create likelihood, minimize
  RooAbsReal *nll = model.createNLL(data,Extended(true));
  RooMinimizer min(*nll);

  //run migrad, hesse
  int result_migrad = min.migrad();
  int result_hesse = min.hesse();

  //Plot result
  TCanvas *c1 = new TCanvas("c1","c1",700,500);
  RooPlot *P = En.frame();
  data.plotOn(P,MarkerColor(kBlue));
  model.plotOn(P,LineColor(kRed));
  P->SetTitle("Spectra 816");
  P->SetName("p816"); 
  P->Draw(); 
  
  //check migrad and hesse
  if(result_migrad == 0 && result_hesse == 0){
    cout << "fit succeeded" << endl;
  }
  else{
    cout << "problem" << endl;
  }
  cout << "Signal Yield: " << sigy.getVal() << endl;
  cout << "Total yield: " << totalYield.getVal() << endl;
}