#include "RooSNLikelihood.hh"

using namespace RooFit;

int main(int argc, char* argv[]){

  // Control the number of files we are looping over.
  int lowerBound = atoi(argv[1]);
  int upperBound = atoi(argv[2]);
  // Setup the input filenames (for the mean and data sets) and 
  // the output filename.
  char* inputMeanFilename = argv[3];
  char* dataFilenamePattern = argv[4];
  char* outputFilename = argv[5];

  // Before we loop over the data files, let us first setup the output ROOT
  // file and its TTree.
  TFile* outputTFile = new TFile(outputFilename, "RECREATE");
  TTree* outputTTree = new TTree("likelihood", "likelihood");
  
  // Create a TCanvas for later.
  TCanvas* c = new TCanvas("c", "c", 800, 800);

  // Let us create a PDF so we can generate toy data.
  //RooRealVar t("t", "t", -0.02, 10);
  //RooSNPdf* pdf = new RooSNPdf("pdf", "pdf", t);
  //pdf->SetMeanFilename(inputMeanFilename);
  //pdf->CreatePdf();

  // Setup the vectors for this TTree. The last five have not been implemented yet.
  std::vector<Double_t> min_vec;
  std::vector<Double_t> err_l_vec;
  std::vector<Double_t> err_h_vec;
  std::vector<Double_t> norm_vec;
  std::vector<Double_t> like_vec;
  std::vector<Double_t> test_stat_vec;
  std::vector<Int_t> evts_vec;
  std::vector<Int_t> burst_num_vec; 

  // Assign these vectors to TBranchs.
  outputTTree->Branch("min_t0", &min_vec);
  outputTTree->Branch("err_l", &err_l_vec);
  outputTTree->Branch("err_h", &err_h_vec);
  outputTTree->Branch("norm", &norm_vec);
  outputTTree->Branch("likelihood", &like_vec);
  outputTTree->Branch("test_statistic", &test_stat_vec);
  outputTTree->Branch("events", &evts_vec);
  outputTTree->Branch("burst_num", &burst_num_vec);


  // Okay now we can begin out for loop. 
  for (int i = lowerBound; i < upperBound; i++){ 

    // This is the observable (t_0, time of core bounce).
    RooRealVar t0("t0", "t0", 0, -0.01, 0.01);
    t0.hasAsymError();
  
    // We need to create out likelihood function here, will use a custom RooAbsReal
    // that I wrote called RooSNLikelihood.
    RooSNLikelihood* customFunc = new RooSNLikelihood("likelihood", "likelihood", t0);

    // This requires four function calls to be used correctly, first must set the data file 
    // name that we are on. 
    std::string dataFilename = dataFilenamePattern;
    dataFilename.append(std::to_string(i));
    dataFilename.append(".root");

    // Now we can pass this to our function and call all other components.
    customFunc->SetDataFilename(dataFilename.c_str());
    customFunc->SetMeanFilename(inputMeanFilename);
    customFunc->ReadMeanEventEvolution();
    customFunc->ReadSimulationData();

    // Create the minimizer for our function.
    RooMinimizer m(*customFunc);

    m.setMaxFunctionCalls(40000);
    m.setMaxIterations(4000);
    m.setOffsetting(kTRUE);

    // Disable any verbose.
    //m.setVerbose(kFALSE);

    // Use migrad to minimize.
    m.migrad();

    // Use minos to get the asym. errors on t_0
    m.minos(t0);

    // Get the minimum and asym. errors at that point.
    Double_t min_t0 = t0.getValV();
    Double_t err_l = t0.getAsymErrorLo();
    Double_t err_h = t0.getAsymErrorHi();
    Double_t like_val = customFunc->getVal(t0);
    // Get some other parameters.
    Double_t norm = customFunc->GetNormalization();
    Double_t KS = customFunc->GetKSTest();
    Int_t numEvents = customFunc->GetNumberOfEvents();
    // Push back there parameteres to our output vectors.
    min_vec.push_back(min_t0);
    err_l_vec.push_back(err_l);
    err_h_vec.push_back(err_h);
    norm_vec.push_back(norm);
    evts_vec.push_back(numEvents);
    burst_num_vec.push_back(i);
    like_vec.push_back(like_val);
    test_stat_vec.push_back(KS);

    // Rename the canvas so it can be written to the file.
    
    std::string canvasName = "canvas-";
    canvasName.append(std::to_string(i));
    c->SetName(canvasName.c_str());

    // Make a frame and plot the above on it.
    RooPlot* frame = t0.frame(Range(min_t0+err_l, min_t0+err_h));
    customFunc->plotOn(frame, ShiftToZero());
    frame->SetMinimum(0);
    frame->SetMaximum(1);
    frame->Draw();

    c->Print("test.png");
    
    // Delete this object. 
    delete customFunc;

  }

  // Okay now we can fill our TTree, write it to a file, clear out vectors
  // and then close the file.
  
  outputTTree->Fill();
  outputTFile->Write();

  min_vec.clear();
  err_l_vec.clear();
  err_h_vec.clear();
  norm_vec.clear();
  like_vec.clear();
  test_stat_vec.clear();
  evts_vec.clear();
  burst_num_vec.clear();

  outputTFile->Close();

  return 0;
     
}