#include "TCanvas.h"
#include "TMath.h"
#include "TH1.h"
#include "TROOT.h"
#include "TSystem.h"
#include "TMinuit.h"

#include "RooExponential.h"
#include "RooAddPdf.h"
#include "RooExtendPdf.h"
#include "RooPolynomial.h"
#include "RooChebychev.h"
#include "RooGaussModel.h"

#include "RooRealVar.h"
#include "RooDataSet.h"
#include "RooFitResult.h"
#include "RooPlot.h"
#include "RooStats/SPlot.h"
#include "TMatrixT.h"

double massL = 4600.;
double massH = 6000.;
double mBs0 = 5360.;
double sigmaBs = 90.;
double mK0 = 497.614;
double nSIG_Gen = 3000;
double nBKG_Gen = 1100;


void makefit(bool doSplot, bool useMyMatrix, bool debug) {
  
  TString myfilename = "log_splot.txt";
  const int totEntries = nSIG_Gen + nBKG_Gen;
  
  std::cout<<"\n ======> Started to execute code. Total entries to be generated: "<<totEntries<< " <====== \n"<<std::endl;
  TCanvas *canb = new TCanvas("canb", "canb", 1000, 1200); canb->Divide(2,2); canb->cd(1);
  TCanvas *can = new TCanvas("can", "can", 1000, 1200); can->Divide(2,2); can->cd(1);

  
  // // **************************************************** // // 
  // //                   Variables, models etc              // //
  // // **************************************************** // // 

  RooRealVar B_MM ("B_MM","B_MM",massL,massH);
  RooPlot *p = B_MM.frame();  
  
  //#-->model for the invariant mass of Bd: double-side Crystal Ball
  RooRealVar mu ("mu","mu",mBs0, mBs0-100,mBs0+100);
  RooRealVar sigma ("sigma","sigma",sigmaBs, 80,110);
  RooGaussModel modelSignal("modelSignal","modelSignal", B_MM, mu, sigma);

  // combinatorial 
  RooRealVar tau ("tau","tau",-0.00038,-0.1,0.1);
  RooExponential modelCombinatorial ("modelCombinatorial","modelCombinatorial",B_MM,tau);

  RooRealVar aSIGYield ("aSIGYield","aSIGYield", nSIG_Gen, 0,1e6);
  RooExtendPdf modelSignal_Norm ("modelSignal_Norm","modelSignal_Norm",modelSignal,aSIGYield);
  RooRealVar bCombYield ("bCombYield","bCombYield", nBKG_Gen, 0,1e6);
  RooExtendPdf modelCombinatorial_Norm ("modelCombinatorial_Norm","modelCombinatorial_Norm",modelCombinatorial,bCombYield);
  RooAddPdf model ("model","model", RooArgList(modelSignal_Norm, modelCombinatorial_Norm));
  
  // // **************************************************** // // 
  // //              Make and plot toy data                  // //
  // // **************************************************** // // 

  RooDataSet* toydata = model.generate(RooArgSet(B_MM),nSIG_Gen + nBKG_Gen);
  can->cd(1);
  toydata->plotOn(p);
  model.plotOn(p, RooFit::LineColor(kBlue));
  model.plotOn(p, RooFit::Components("modelSignal_Norm"), RooFit::LineColor(kBlue));
  model.plotOn(p, RooFit::Components("modelCombinatorial_Norm"), RooFit::LineColor(kGreen));
  p->Draw();

  // // **************************************************** // // 
  // //              Fit toy data. All floating              // //
  // // **************************************************** // // 
  gSystem->RedirectOutput(myfilename);
  RooFitResult *res_mll =  model.fitTo(*toydata, RooFit::Save(true), RooFit::Extended(kTRUE), RooFit::SumW2Error(kTRUE) );
  gSystem->RedirectOutput(0);
  double sigYieldFitted = aSIGYield.getVal();
  double combYieldFitted = bCombYield.getVal();
  std::cout<<" \n All parameters floating... the yields are (S) "<<sigYieldFitted
	   <<"  (B) "<< combYieldFitted <<" and total entries = "<<totEntries<<"\n"<<std::endl;

  // // **************************************************** // // 
  // //              Fit toy data. yields floating           // //
  // // **************************************************** // // 

  mu.setConstant();
  sigma.setConstant();
  tau.setConstant();

  gSystem->RedirectOutput(myfilename);
  RooFitResult *res_mll_fix =  model.fitTo(*toydata, RooFit::Save(true), RooFit::Extended());
  gSystem->RedirectOutput(0);
  sigYieldFitted = aSIGYield.getVal();
  combYieldFitted = bCombYield.getVal();
  std::cout<<"\n Only yields floating.... the yields are (S) "<<sigYieldFitted
	   <<" (B) "<< combYieldFitted <<" and total entries = "<<totEntries<<"\n" <<std::endl;

  model.plotOn(p, RooFit::LineColor(kBlue), RooFit::LineStyle(kDashed) );
  model.plotOn(p, RooFit::Components("modelSignal_Norm"), RooFit::LineStyle(kDashed), RooFit::LineColor(kMagenta));
  model.plotOn(p, RooFit::Components("modelCombinatorial_Norm"), RooFit::LineStyle(kDashed), RooFit::LineColor(kRed));
  model.paramOn(p);
  can->cd(2);
  p->Draw();

  //  TMinuit *gMinuit = new TMinuit(2);
  //  gMinuit->SetFCN(model);

  // // **************************************************** // // 
  // //                Print fit results                     // //
  // // **************************************************** // // 

  std::cout<<"\n  ==================================================    \n result of initial fit "<<std::endl;
  res_mll->Print();
  const TMatrixDSym &res_cov = res_mll->covarianceMatrix();
  std::cout<<"Printing covariance matrix of first fit .... "<<std::endl;
  res_cov.Print();
  std::cout<<" ===== done printing ===== \n"<<std::endl;

  std::cout<<"\n ==================================================    \n result of the seocond fit "<<std::endl;
  res_mll_fix->Print();
  const TMatrixDSym &res_cov_fix = res_mll_fix->covarianceMatrix();
  std::cout<<"Printing covariance matrix of the second fit .... "<<std::endl;
  res_cov_fix.Print();
  std::cout<<" ===== done printing ===== \n"<<std::endl;

  double cov_sig_sig = res_cov_fix[0][0]; 
  double cov_sig_comb = res_cov_fix[0][1];
  double cov_comb_comb = res_cov_fix[1][1]; 

  std::cout<<"The covariances terms from the fit matrix are [0][0] = "<<cov_sig_sig
	   <<" [0][1] = "<<cov_sig_comb<<"  [1][1] = "<<cov_comb_comb<<std::endl;

  // // **************************************************** // // 
  // //  Store mass and pdf values for each data point       // //
  // // **************************************************** // // 

  std::vector<double> sigpdf_val;
  std::vector<double> combpdf_val;
  std::vector<double>  model_val;
  std::vector<double>  mass_val;
  double Bmass = 0.;
  
  for (Int_t ievt = 0; ievt < totEntries; ++ievt)   {
    const RooArgSet *argset = toydata->get(ievt);
    Bmass = argset->getRealValue("B_MM"); 
    B_MM.setVal(Bmass);

    sigpdf_val.push_back(modelSignal.evaluate());
    combpdf_val.push_back(modelCombinatorial.evaluate());
    model_val.push_back(model.evaluate());
    mass_val.push_back(Bmass);
   
    if(debug) { 
      std::cout<<ievt<<". sanity check "<<modelSignal.evaluate()<<" == "<<modelSignal_Norm.evaluate()
	       << " and "<< modelCombinatorial.evaluate()<<" == "<<modelCombinatorial_Norm.evaluate()<<std::endl;
    }
  }
    
  std::cout<<"\n // // **************************************************** // // "<<std::endl;
  std::cout<<" //                          SOME tests                       // //"<<std::endl;
  std::cout<<"// // **************************************************** // // \n"<<std::endl;
  
  
  B_MM.setRange( "integralRange", massL, massH);
  RooAbsReal *integral = model.createIntegral(B_MM, RooFit::NormSet(B_MM), RooFit::Range("integralRange") ); 
  std::cout<<"norm of model "<<integral->getVal() <<std::endl;
  integral = modelSignal_Norm.createIntegral(B_MM, RooFit::NormSet(B_MM), RooFit::Range("integralRange") ); 
  std::cout<<"norm of signal "<<integral->getVal() <<std::endl;
  integral = modelCombinatorial_Norm.createIntegral(B_MM, RooFit::NormSet(B_MM), RooFit::Range("integralRange") ); 
  std::cout<<"norm of comb.  "<<integral->getVal() <<std::endl;
  //double pdfnormErr = integral->getPropagatedError(*res_mll_fix);
  
  RooArgSet vars(aSIGYield, bCombYield, mu, sigma, tau);
  B_MM.setVal(5484.27);
  std::cout<<model.getVal(&vars)<<" v/s evaluate "<<model.evaluate()<<std::endl;
  

  std::cout<<"\n // // **************************************************** // // "<<std::endl;
  std::cout<<"//        Calculate the inverse covariance matrix          // //"<<std::endl;
  std::cout<<"// // **************************************************** // // \n"<<std::endl;
  
  const int nspec = 2;
  TMatrixD covInv(nspec, nspec);
  for (Int_t i = 0; i < nspec; i++) for (Int_t j = 0; j < nspec; j++) covInv(i,j) = 0;
   
  std::cout<< "\n Calculating inverse covariance matrix... \n"<<std::endl;
  
  for (Int_t ievt = 0; ievt < totEntries; ++ievt)   {

    double myVals[nspec] = { sigpdf_val.at(ievt), combpdf_val.at(ievt)};
    // Sum for the denominator
    Double_t dsum = 0.;
    dsum += myVals[0]*sigYieldFitted;
    dsum += myVals[1]*combYieldFitted;
    
    for(Int_t n=0; n<nspec; ++n)
      for(Int_t j=0; j<nspec; ++j) 
	{
	  covInv(n,j) +=  myVals[n]*myVals[j]/(dsum*dsum) ;
	}
  } // for ievt

  std::cout<<"Printing the inverse covariance matrix "<<std::endl;   covInv.Print();

  // // **************************************************** // // 
  //           Invert to get the covariance matrix           // //
  // // **************************************************** // // 

  if (covInv.Determinant() <=0) {
    std::cout<< " Error: covariance matrix is singular; I can't invert it!" << std::endl;
  }
  else {
    std::cout<<"\n ===== > Matrix was inverted... \n "<<std::endl;
  }

  TMatrixD covMatrix(TMatrixD::kInverted,covInv);
  std::cout<<"Printing the covariance matrix "<<std::endl;  covMatrix.Print();

  // // set them to my new vals:
  if(useMyMatrix){
    cov_sig_sig = covMatrix(0,0);
    cov_comb_comb = covMatrix(1,1);
    cov_sig_comb = covMatrix(0,1);
    std::cout<<"The [NEW] covariances are [0][0] "<<cov_sig_sig
	     <<" [0][1] = "<<cov_sig_comb<<"  [1][1] "<<cov_comb_comb<<std::endl;
  }
  
  // // do some checks on the matrix
  if(debug){
    std::cout<< "\n Yield of specie | Sum of Row in Matrix |  Norm \n"<<
      " ---------------------------------------------- " << std::endl; 
    
    double yieldvalues[nspec] = { sigYieldFitted, combYieldFitted};
    for(Int_t k=0; k<nspec; ++k) 
      {
	Double_t covnorm(0) ;
	for(Int_t m=0; m<nspec; ++m) covnorm += covInv[k][m]*yieldvalues[m];
	Double_t sumrow(0) ;
	for(Int_t m = 0; m < nspec; ++m) sumrow += covMatrix[k][m] ;
	std::cout <<"   "<< yieldvalues[k] << "       |           " << sumrow << " | " << covnorm << endl ;
    }
  }


  std::cout<<"// // **************************************************** // // "<<std::endl;
  std::cout<<"//                Compute sweights  (Fatima)               // //"<<std::endl;
  std::cout<<"// // **************************************************** // // "<<std::endl;

  std::cout<<"\n Computing sweights. The covariance matrix elements I use are [0][0] "<<cov_sig_sig
	   <<" [0][1] = "<<cov_sig_comb<<"  [1][1] "<<cov_comb_comb<<"\n"<<std::endl;

  double sigwt[totEntries];
  double combwt[totEntries];
  double myyield_sig  = 0;
  double myyield_comb = 0.;

  for(int iEv = 0; iEv < totEntries; iEv++){

    double combpdf = combpdf_val.at(iEv);
    double sigpdf  = sigpdf_val.at(iEv);
    if(debug) std::cout<< iEv<<". check... "<<(sigYieldFitted + combYieldFitted)*model_val.at(iEv)
		       <<" == "<<(combpdf*combYieldFitted +  sigpdf*sigYieldFitted) <<" ?? "<< std::endl;
    double denominator = combpdf*combYieldFitted + sigpdf*sigYieldFitted;
    double comb_ev = ( cov_comb_comb*combpdf + cov_sig_comb*sigpdf) / denominator; 
    double signal_ev = ( cov_sig_sig*sigpdf +  cov_sig_comb*combpdf) / denominator; 
    
    if(iEv%1000==0) std::cout<<iEv<<". Bmass = "<<mass_val.at(iEv)<<" W_s = "<<signal_ev
			     <<" W_comb = "<<comb_ev<<" their sum "<<signal_ev+comb_ev<<std::endl;
    
    sigwt[iEv] = signal_ev;
    combwt[iEv] = comb_ev;
    myyield_sig +=signal_ev;
    myyield_comb += comb_ev;
  } // for iEv

  std::cout<<"\n Sum of weights  sig "<<myyield_sig<<" comb "<<myyield_comb
	   <<" (should be == the fitted yields: "<<sigYieldFitted<<" and "<<combYieldFitted<<")\n"<<std::endl;

  // //  ***** Make some plots from my weights ***** // //
  canb->cd(1);
  TH1F *h1 = new TH1F("h1", "h1", 100, massL, massH); h1->Sumw2();
  TH1F *h2 = new TH1F("h2", "h2", 100, massL, massH); h2->Sumw2();
  for(int i=0 ; i<totEntries; i++){
    double m_tmp = mass_val.at(i);
    h1->Fill(m_tmp, sigwt[i]);
    h2->Fill(m_tmp, combwt[i]);
  }
  h2->SetLineColor(kRed);   h2->SetMarkerColor(kRed);
  h1->DrawCopy(); 
  canb->cd(2); h2->DrawCopy();
  // // end of plotting 
  

  if(doSplot){
  
    std::cout<<"// // **************************************************** // // "<<std::endl;
    std::cout<<"// //                 Using sPlot in roostats              // // "<<std::endl;
    std::cout<<"// // **************************************************** // //"<<std::endl;
    
    
    mu.setConstant();
    sigma.setConstant();
    tau.setConstant();
    aSIGYield.setVal(sigYieldFitted*0.9);
    bCombYield.setVal(combYieldFitted*0.9);

    RooStats::SPlot* sData = new RooStats::SPlot("sData","An SPlot", *toydata, &model, RooArgList(aSIGYield, bCombYield) );
    
    std::cout << "\n Check SWeights:" << std::endl;
    std::cout << "Yield of Sig is " 
	      << aSIGYield.getVal() << ".  From sWeights it is "
	      << sData->GetYieldFromSWeight("aSIGYield") << std::endl;
    std::cout << "Yield of BKG is " 
	      << bCombYield.getVal() << ".  From sWeights it is "
	      << sData->GetYieldFromSWeight("bCombYield") << std::endl;
    std::cout << std::endl;
    
    
    for(Int_t i=0; i < totEntries; i+=1000)
      {
	std::cout << i<<". Bmass = "<< mass_val.at(i)
		  << " sig Weight   " << sData->GetSWeight(i,"aSIGYield") 
		  << " comb Weight   " << sData->GetSWeight(i,"bCombYield") 
		  << " Total Weight   " << sData->GetSumOfEventSWeight(i) 
		  << std::endl;
      }
    
    
    B_MM.setBins(100);
    RooDataSet * data_sig = new RooDataSet(toydata->GetName(),toydata->GetTitle(), toydata, *toydata->get(),0,"aSIGYield_sw"); 
    RooPlot *p2[2];
    p2[0]= B_MM.frame();
    data_sig->plotOn(p2[0]);
    can->cd(3); p2[0]->Draw();
    
    RooDataSet * data_bkg = new RooDataSet(toydata->GetName(),toydata->GetTitle(), toydata, *toydata->get(),0,"bCombYield_sw"); 
    p2[1]= B_MM.frame();
    data_bkg->plotOn(p2[0], RooFit::LineColor(kRed), RooFit::MarkerColor(kRed));
    can->cd(4); p2[0]->Draw();
    
  } // if  doSplot
  
  std::cout<<"\n Fit output saved in file "<<myfilename<< "\n" <<std::endl;
  
}