#include <iostream>
#include <fstream>
#include <vector>
#include <string>
#include <map>

#include "boost/program_options.hpp"
#include "boost/lexical_cast.hpp"

#include "TROOT.h"
#include "TF1.h"
#include "TFile.h"
#include "TMath.h"
#include "TChain.h"
#include "TLegend.h"
#include "TCanvas.h"
#include "RooPlot.h"
#include "RooRealVar.h"
#include "RooWorkspace.h"
#include "RooDataSet.h"
#include "RooAbsData.h"
#include "RooNDKeysPdf.h"
#include "RooAbsPdf.h"
#include "RooGaussian.h"
#include "RooProdPdf.h"
#include "RooArgSet.h"
#include "RooFitResult.h"
#include "RooMinuit.h"
#include "RooMinimizer.h"
#include "RooMsgService.h"
#include "RooDataHist.h"
#include "RooExtendPdf.h"
#include "TRandom3.h"
#include "TLatex.h"
#include "TMacro.h"
#include "TH1F.h"
#include "TH1I.h"
#include "TArrow.h"
#include "TKey.h"

#include "RooCategory.h"

#include <Math/PdfFuncMathCore.h>
#include <Math/ProbFunc.h>
#include <iomanip>

using namespace std;
using namespace RooFit;
using namespace boost;

namespace po = program_options;

bool runFtestCheckWithToys=false;
int nBinsForMass = 320;

TRandom3 *RandomGen = new TRandom3();




void plot(RooRealVar *mass, RooAbsReal *pdf, RooDataSet *data, string name, string xTitle, double xMin, double xMax, int bins, double *prob){
  
  // Chi2 taken from full range fit
  mass->setRange(xMin, xMax);


  /*
  RooPlot *plot_chi2 = mass->frame();
  data->plotOn(plot_chi2,Binning(nBinsForMass));
  pdf->plotOn(plot_chi2);

  int np = pdf->getParameters(*data)->getSize()+1; //Because this pdf has no extend
  double chi2 = plot_chi2->chiSquare(np);

  *prob = getGoodnessOfFit(mass,pdf,data,name);
  */


  *prob = 0;
  RooPlot *plot = mass->frame();
  if(data != NULL)  data->plotOn(plot,Binning(bins));

  TCanvas *canv = new TCanvas();
  pdf->plotOn(plot);//,RooFit::NormRange("fitdata_1,fitdata_2"));
  //pdf->paramOn(plot,RooFit::Layout(0.34,0.96,0.89),RooFit::Format("NEA",AutoPrecision(1)));
  plot->SetMinimum(0.0001);
  plot->SetXTitle(xTitle.c_str());
  plot->SetYTitle("Events");
  plot->Draw();

  TLatex *lat = new TLatex();
  lat->SetNDC();
  lat->SetTextFont(42);
  //lat->DrawLatex(0.1,0.92,Form("#chi^{2} = %.3f, Prob = %.2f",chi2*(nBinsForMass-np),*prob));
  canv->SaveAs(Form("%s.png",name.c_str()));
  
  delete canv;
  delete lat;
}

void plot(string name, TH1 *histo, string xTitle){
  
  TCanvas *canv = new TCanvas();
  histo->GetXaxis()->SetTitle(xTitle.c_str());
  histo->GetYaxis()->SetTitle("Events");
  histo->Draw();

  TLatex *lat = new TLatex();
  lat->SetNDC();
  lat->SetTextFont(42);
  //lat->DrawLatex(0.1,0.92,Form("#chi^{2} = %.3f, Prob = %.2f",chi2*(nBinsForMass-np),*prob));
  canv->SaveAs(Form("%s.png",name.c_str()));
  
  delete canv;
  delete lat;
}

void plot(string name, TGraph *histo, string xTitle){
  
  TCanvas *canv = new TCanvas();

  histo->SetLineColor(1);
  histo->SetMarkerColor(1);
  histo->SetLineWidth(2.5);
  histo->SetMarkerSize(1.);
  histo->SetMarkerStyle(21);
  histo->GetXaxis()->SetTitle(xTitle.c_str());
  histo->GetYaxis()->SetTitle("Fraction");
  histo->Draw("AC");

  TLatex *lat = new TLatex();
  lat->SetNDC();
  lat->SetTextFont(42);
  //lat->DrawLatex(0.1,0.92,Form("#chi^{2} = %.3f, Prob = %.2f",chi2*(nBinsForMass-np),*prob));
  canv->SaveAs(Form("%s.png",name.c_str()));
  
  delete canv;
  delete lat;
}









void plot(RooRealVar *mass, map<string,RooAbsPdf*> pdfs, RooDataSet *data, string name, int cat, int bestFitPdf=-1){
  
  int color[7] = {kBlue,kRed,kMagenta,kGreen+1,kOrange+7,kAzure+10,kBlack};
  TCanvas *canv = new TCanvas();
  TLegend *leg = new TLegend(0.6,0.65,0.89,0.89);
  leg->SetFillColor(0);
  leg->SetLineColor(0);
  RooPlot *plot = mass->frame();

  mass->setRange("unblindReg_1",100,110);
  mass->setRange("unblindReg_2",150,180);
  data->plotOn(plot,Binning(80));

  TObject *datLeg = plot->getObject(int(plot->numItems()-1));
  leg->AddEntry(datLeg,Form("Data - cat%d",cat),"LEP");
  int i=0;
  int style=1;
  for (map<string,RooAbsPdf*>::iterator it=pdfs.begin(); it!=pdfs.end(); it++){
    int col;
    if (i<=6) col=color[i];
    else {col=kBlack; style++;}
    it->second->plotOn(plot,LineColor(col),LineStyle(style));//,RooFit::NormRange("fitdata_1,fitdata_2"));
    TObject *pdfLeg = plot->getObject(int(plot->numItems()-1));
    std::string ext = "";
    if (bestFitPdf==i) ext=" (Best Fit Pdf) ";
    leg->AddEntry(pdfLeg,Form("%s%s",it->first.c_str(),ext.c_str()),"L");
    i++;
  }
  plot->SetTitle(Form("Category %d",cat));
  plot->SetMinimum(0.0001);
  plot->Draw();
  leg->Draw("same");
  canv->SaveAs(Form("%s.pdf",name.c_str()));
  canv->SaveAs(Form("%s.png",name.c_str()));
  delete canv;
}

vector<int>  findXMinMax(TH1D *h_histo){

     vector<int>  MinMax; int Xmin; int Xmax;
     for(int iloop=1; iloop<=h_histo->GetNbinsX(); iloop++){

          if(h_histo->GetBinContent(iloop) == 0) continue;
          Xmin = iloop;
          break;
     }

     for(int iloop=1; iloop<=h_histo->GetNbinsX(); iloop++){

          if(h_histo->GetBinContent(h_histo->GetNbinsX()-iloop+1) == 0) continue;
          Xmax = h_histo->GetNbinsX()-iloop+1;
          break;
     }

     MinMax.push_back(Xmin);
     MinMax.push_back(Xmax);

     return MinMax;

}

void runFit(string pdfNom, RooArgList *varList, RooDataSet *data, string option, int varIndex, string xTitle, double xMin, double xMax, int bins){

        cout << "FAN running fit and plot..." << endl;

        double prob=0;

        RooNDKeysPdf *pdf = new RooNDKeysPdf(pdfNom.c_str(), pdfNom.c_str(), *varList, *data, option.c_str());
       
        RooRealVar *plotVar = (RooRealVar*)varList->at(varIndex);
        plot(plotVar, pdf, data, pdfNom, xTitle, xMin, xMax, bins, &prob); 

}

TGraph *integralHisto(TH1 *histo){
 
      TGraph *graph = new TGraph(); 
      for(int iloop=1; iloop<=histo->GetNbinsX(); iloop++){
          
             graph->SetPoint(iloop-1, histo->GetXaxis()->GetBinUpEdge(iloop), 1.0-histo->Integral(1, iloop));

      }

      
      return graph;

}





















void runPrediction(string pdfNom, string outDir, string ofileName, RooArgList *varListTrain, RooArgList *varListTest, RooDataSet *dataTrain, RooDataSet *dataTest, string option, double high, string outLog, int bins){

        ofstream outfile(outLog.c_str());
        outfile << "FAN running prediction..." << endl;

        TFile *myFile = new TFile(Form("%s/%s", outDir.c_str(), ofileName.c_str()),"RECREATE");
        RooWorkspace *outWS = new RooWorkspace("outWS");
        myFile->cd();

        double prob=0;
        string pdfNomTrain = Form("%sTrain", pdfNom.c_str());
        string pdfNomTest = Form("%sTest", pdfNom.c_str());

        RooNDKeysPdf *pdfTrain = new RooNDKeysPdf(pdfNomTrain.c_str(), pdfNomTrain.c_str(), *varListTrain, *dataTrain, option.c_str());
        outWS->import(*dataTrain, RecycleConflictNodes());
        outWS->import(*pdfTrain, RecycleConflictNodes());
        outfile << "FAN running prediction step 1 finish..." << endl;

        RooNDKeysPdf *pdfTest = new RooNDKeysPdf(pdfNomTest.c_str(), pdfNomTest.c_str(), *varListTest, *dataTest, option.c_str());
        outWS->import(*dataTest, RecycleConflictNodes());
        outWS->import(*pdfTest, RecycleConflictNodes());
        outfile << "FAN running prediction step 2 finish..." << endl;
      
        // product conditional pdf
        RooRealVar *varConditional = (RooRealVar*)varListTrain->at(0); 
        RooProdPdf *pdfConditional = new RooProdPdf("Conditional", "Conditional", *pdfTest, Conditional(*pdfTrain, *varConditional)); 
        outWS->import(*pdfConditional, RecycleConflictNodes());
        //pdfConditional->Print("v");
        outfile << "FAN running prediction step 3 finish..." << endl;
        //RooFit::Precision(1) 
       
        RooArgSet *varSetTest = new RooArgSet(*varListTest);
        RooAbsPdf *pdfProjection = pdfConditional->createProjection(*varSetTest); 
        outWS->import(*pdfProjection, RecycleConflictNodes());
        //pdfProjection->Print();
        outfile << "FAN running prediction step 4 finish..." << endl;

        //TH1 *TH1Conditional = pdfProjection->createHistogram("predict", *varConditional, Binning(bins));
        TH1 *TH1Conditional = pdfProjection->createHistogram("predict", *varConditional);
        outfile << TH1Conditional->Integral(1, 100) << endl; 
        plot(Form("%s/%s_TH1", outDir.c_str(), pdfNom.c_str()), TH1Conditional, "dose/Gy");
        TH1Conditional->Write();
        outfile << "FAN running prediction step 5 finish..." << endl;

        TGraph *predictDVH = (TGraph*)integralHisto(TH1Conditional);
        outfile << "FAN running prediction integration finish.." << endl; 
        plot(Form("%s/%s_DVH", outDir.c_str(), pdfNom.c_str()), predictDVH, "dose/Gy");
        predictDVH->Write();
        outfile << "FAN running prediction step 6 finish..." << endl;
         
        outWS->Print();
        outWS->Write();
        myFile->Close(); 


 
        varConditional->setRange("inte", 20, high) ; 
        RooArgSet *varSetTrain = new RooArgSet(*varListTrain);
        //RooAbsReal *predicTest = pdfConditional->createIntegral(*varSetTrain, NormSet(*varSetTrain), Range("inte"));
        //RooAbsReal *predicTest = pdfConditional->createIntegral(*varSetTrain, NormSet(*varSetTrain), Range(""));
        RooAbsReal *predicTest = pdfProjection->createIntegral(*varSetTrain, NormSet(*varSetTrain), Range(""));

        outfile << predicTest->getVal() << endl;    
               




        /* 
        RooArgSet *varSetTest = new RooArgSet(*varListTest); 
        RooAbsReal *predicTestO = pdfConditional->createIntegral(*varSetTest, NormSet(*varSetTest), Range("")) ;
        cout << predicTestO->getVal() << endl;    
        //predicTestO->Print();      
        RooAbsReal *predicTest = predicTestO->createIntegral(*varConditional, NormSet(*varConditional), Range("inte")) ;
        cout << predicTest->getVal() << endl;
        //predicTest->Print();      
        */






}











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

  gROOT->ProcessLine(".x setTDRStyle.C");
  
 
  string fileName;
  string ofileName;
  int ncats;
  double intevalue;
  string datfile;
  string outDir;
  string inDirTr;
  string inDirTe;
  bool verbose=false;

  po::options_description desc("Allowed options");
  desc.add_options()
    ("help,h",                                                                                  "Show help")
    ("infilename,i", po::value<string>(&fileName),                                              "In file name")
    ("outfilename,o", po::value<string>(&ofileName),                                              "Out file name")
    ("ncats,c", po::value<int>(&ncats)->default_value(5),                                       "Number of categories")
    ("intevalue,b", po::value<double>(&intevalue)->default_value(20.),                                       "Integral high bound")
    ("datfile,d", po::value<string>(&datfile)->default_value("dat/fTest.dat"),                  "Right results to datfile for BiasStudy")
    ("outDir,D", po::value<string>(&outDir)->default_value("plots"),                      "Out directory for plots")
    ("inDirTr,I", po::value<string>(&inDirTr)->default_value("input train"),                      "In directory for train roots")
    ("inDirTe,T", po::value<string>(&inDirTe)->default_value("input Test"),                      "In directory for test roots")
    ("runFtestCheckWithToys", 									"When running the F-test, use toys to calculate pvals (and make plots) ")
    ("verbose,v",                                                                               "Run with more output")
  ;
  po::variables_map vm;
  po::store(po::parse_command_line(argc,argv,desc),vm);
  po::notify(vm);
  if (vm.count("help")) { cout << desc << endl; exit(1); }
  if (vm.count("verbose")) verbose=true;
  if (vm.count("runFtestCheckWithToys")) runFtestCheckWithToys=true;

  if (!verbose) {
    RooMsgService::instance().setGlobalKillBelow(RooFit::ERROR);
    RooMsgService::instance().setSilentMode(true);
    gErrorIgnoreLevel=kWarning;
  }
  

  system(Form("mkdir -p %s",outDir.c_str()));
  TFile *inFile = TFile::Open(fileName.c_str());
  //RooWorkspace *inWS = (RooWorkspace*)inFile->Get("cms_hgg_workspace");

  
  // store results here
  /*
  FILE *resFile ;
  if  (singleCategory >-1) resFile = fopen(Form("%s/fTestResults_cat%d.txt",outDir.c_str(),singleCategory),"w");
  else resFile = fopen(Form("%s/fTestResults.txt",outDir.c_str()),"w"); 
  vector<map<string,int> > choices_vec;
  vector<map<string,std::vector<int> > > choices_envelope_vec;
  vector<map<string,RooAbsPdf*> > pdfs_vec;
  */

  
  double rx, ry, rz, dist, dose;
  TTree *tree = (TTree*)inFile->Get("probDose");
  TTree *newtree = tree->CopyTree("", "", 1000,0);
  TChain *chain = new TChain("probDose"); 
  TChain *chainTr = new TChain("probDose"); 
  TChain *chainTe = new TChain("probDose"); 
  chain->Add(Form("/sps/cms/jfan/CMSSW_7_1_8/src/RadioTherapy/AutoPlan/ProbDose/%s/*.root", inDirTr.c_str()));
  chain->Add(Form("/sps/cms/jfan/CMSSW_7_1_8/src/RadioTherapy/AutoPlan/ProbDose/%s/*.root", inDirTe.c_str()));
  chainTr->Add(Form("/sps/cms/jfan/CMSSW_7_1_8/src/RadioTherapy/AutoPlan/ProbDose/%s/*.root", inDirTr.c_str()));
  chainTe->Add(Form("/sps/cms/jfan/CMSSW_7_1_8/src/RadioTherapy/AutoPlan/ProbDose/%s/*.root", inDirTe.c_str()));

  chain->SetBranchAddress("rx", &rx);
  chain->SetBranchAddress("ry", &ry);
  chain->SetBranchAddress("rz", &rz);
  chain->SetBranchAddress("dist", &dist);
  chain->SetBranchAddress("dose", &dose);

  chainTr->SetBranchAddress("rx", &rx);
  chainTr->SetBranchAddress("ry", &ry);
  chainTr->SetBranchAddress("rz", &rz);
  chainTr->SetBranchAddress("dist", &dist);
  chainTr->SetBranchAddress("dose", &dose);

  chainTe->SetBranchAddress("rx", &rx);
  chainTe->SetBranchAddress("ry", &ry);
  chainTe->SetBranchAddress("rz", &rz);
  chainTe->SetBranchAddress("dist", &dist);
  chainTe->SetBranchAddress("dose", &dose);

  cout << "Total entries train and test chain " << chain->GetEntries() << endl;
  cout << "Total entries train chain " << chainTr->GetEntries() << endl;
  cout << "Total entries test chain " << chainTe->GetEntries() << endl;
  cout << "Total entries tree " << tree->GetEntries() << endl;
  //cout << "Total entries new  " << newtree->GetEntries() << endl;
  /*
  for(int i=0; i<tree->GetEntries(); i++){
 
        tree->GetEntry(i); 

        //cout << rx << " " << ry << " " << rz << " " << dist << " " << dose << endl; 

  }
  */
 
  TH1D *h_temp_dist = new TH1D("temp_dist", "temp_dist", 300, -100, 200);
  TH1D *h_temp_dose = new TH1D("temp_dose", "temp_dose", 100, 0., 100);
  chain->Draw("dist>>temp_dist"); 
  chain->Draw("dose>>temp_dose"); 

  vector<double>  MinMax_dist;
  MinMax_dist.push_back(h_temp_dist->GetXaxis()->GetBinLowEdge(findXMinMax(h_temp_dist)[0]));
  MinMax_dist.push_back(h_temp_dist->GetXaxis()->GetBinUpEdge(findXMinMax(h_temp_dist)[1]));
  cout << "dist Min---> " << MinMax_dist[0] << "  dist Max---> " << MinMax_dist[1] << endl;

  vector<double>  MinMax_dose;
  MinMax_dose.push_back(h_temp_dose->GetXaxis()->GetBinLowEdge(findXMinMax(h_temp_dose)[0]));
  MinMax_dose.push_back(h_temp_dose->GetXaxis()->GetBinUpEdge(findXMinMax(h_temp_dose)[1]));
  cout << "dose Min---> " << MinMax_dose[0] << "  dose Max---> " << MinMax_dose[1] << endl;

  
  RooRealVar *datasetrx = new RooRealVar("rx", "rx", 0); 
  RooRealVar *datasetry = new RooRealVar("ry", "ry", 0); 
  RooRealVar *datasetrz = new RooRealVar("rz", "rz", 0); 
  RooRealVar *datasetdist = new RooRealVar("dist", "dist", 0, int(MinMax_dist[0]*0.9)*1.0, int(MinMax_dist[1]*1.1)*1.0); 
  RooRealVar *datasetdose = new RooRealVar("dose", "dose", 1, int(MinMax_dose[0]*0.9)*1.0, int(MinMax_dose[1]*1.1)*1.0); 

  RooDataSet *data_distTr = new RooDataSet("data_distTr", "data_distTr", chainTr, RooArgSet(*datasetdist));
  RooDataSet *data_distDTr = new RooDataSet("data_distDTr", "data_distDTr", chainTr, RooArgSet(*datasetdose, *datasetdist));
  RooDataSet *data_distTe = new RooDataSet("data_distTe", "data_distTe", chainTe, RooArgSet(*datasetdist));
  RooDataSet *data_doseTe = new RooDataSet("data_doseTe", "data_doseTe", chainTe, RooArgSet(*datasetdose));
  //RooDataSet *data_dist = new RooDataSet("data_dist", "data_dist", newtree, RooArgSet(*datasetdist));
  //RooDataSet *data_distD = new RooDataSet("data_distD", "data_distD", newtree, RooArgSet(*datasetdose, *datasetdist));

  RooArgList *varListdist = new RooArgList(*datasetdist);
  RooArgList *varListdose = new RooArgList(*datasetdose);
  RooArgList *varListdistD = new RooArgList(*datasetdose, *datasetdist);

  // 1D fit plot train
  if(ncats == 0)  runFit(Form("%s/distOnlyTr", outDir.c_str()), varListdist, data_distTr, "am", 0, "dist/mm", int(MinMax_dist[0]*0.9)*1.0, int(MinMax_dist[1]*1.1)*1.0, int(int(MinMax_dist[1]*1.1)*1.0-int(MinMax_dist[0]*0.9)*1.0));

  // 2D fit plot train
  if(ncats == 1)  runFit(Form("%s/distDoseTr_dist", outDir.c_str()), varListdistD, data_distDTr, "am", 1, "dist/mm", int(MinMax_dist[0]*0.9)*1.0, int(MinMax_dist[1]*1.1)*1.0, int(int(MinMax_dist[1]*1.1)*1.0-int(MinMax_dist[0]*0.9)*1.0));
  if(ncats == 2)  runFit(Form("%s/distDoseTr_dose", outDir.c_str()), varListdistD, data_distDTr, "am", 0, "dose/Gy", int(MinMax_dose[0]*0.9)*1.0, int(MinMax_dose[1]*1.1)*1.0, int(int(MinMax_dose[1]*1.1)*1.0-int(MinMax_dose[0]*0.9)*1.0));

  // 1D fit plot test
  if(ncats == 3)  runFit(Form("%s/doseOnlyTe", outDir.c_str()), varListdose, data_doseTe, "am", 0, "dose/Gy", int(MinMax_dose[0]*0.9)*1.0, int(MinMax_dose[1]*1.1)*1.0, int(int(MinMax_dose[1]*1.1)*1.0-int(MinMax_dose[0]*0.9)*1.0));
 
  // prediction
  if(ncats == 4){
           runPrediction("predicDoseT", outDir, ofileName, varListdistD, varListdist, data_distDTr, data_distTe, "am", intevalue, Form("%s/value_%.f.log", outDir.c_str(), intevalue), int(int(MinMax_dose[1]*1.1)*1.0-int(MinMax_dose[0]*0.9)*1.0)*100); 







   }













}