//
// call TMinut directly to fit a power law
// using uncorrelated uncertainties
//

//-------------------------------------------------------
// Header files from C++ and ROOT
//-------------------------------------------------------

// c++ headers
#include <iostream>
#include <fstream>

// root headers
#include <TMath.h>
#include <TH2D.h>
#include <TFile.h>
#include <TF2.h>
#include <TF1.h>
#include <TRandom.h>
#include <TGraphErrors.h>
#include <TPad.h>
#include <TCanvas.h>
#include <TMinuit.h>
#include <TStyle.h>
#include <TGaxis.h>
#include <TLatex.h>
#include <TLegend.h>

// -------------------------------------------------------------------------------
// ranges for fits and plots
// -------------------------------------------------------------------------------
double gxmin = 17., gymin = 15., gxmax = 1400., gymax = 1000.; // limits for plot

// -------------------------------------------------------------------------------
// keep track of chi2
// -------------------------------------------------------------------------------
double gChi2 = -1;

// -------------------------------------------------------------------------------
// parameters for ideal power law
// -------------------------------------------------------------------------------
double lambdaPL = 0.7;
double nPL = 70; // nb
double wPL = 90; // GeV

// -------------------------------------------------------------------------------
// info to generate power law
// -------------------------------------------------------------------------------
const int nPoints = 9; // number of measurements to be simulated
const double wPoints[nPoints] = {24, 30, 40, 50,  70,  130,  190,  390,  700}; // GeV
const double staUnc[nPoints] = {0.05, 0.03, 0.05, 0.07, 0.07, 0.03, 0.03, 0.07,0.04}; // rel stat unc
const double uncUnc[nPoints] = {0.08, 0.08, 0.08, 0.07,
				0.05, 0.07, 0.07, 0.07, 0.10}; // uncorrelated unceratainty in percent
const int nCor = 4; // number of correlated sources of uncertainty 
const double normUnc = 0.03; // global normalization correlated error
const double fwdUnc = 0.05; // correlated uncertainty for fwd points
const double sfUnc = 0.04; // correlated uncertainty for semi-fwd points
const double cbUnc = 0.04; // correlated uncertainty for central barrel points
const double corUnc[nCor][nPoints] = // relative correlated uncertainty
  {
    normUnc, normUnc, normUnc, normUnc, normUnc, normUnc, normUnc, normUnc, normUnc, 
    fwdUnc, fwdUnc, fwdUnc, 0.0, 0.0, 0.0, 0.0, 0.0, fwdUnc, // only fwd
     0.0, 0.0, 0.0, sfUnc, sfUnc,  0.0, 0.0, sfUnc, 0.0,  // only semifwd
     0.0, 0.0, 0.0,  0.0,  0.0, cbUnc, cbUnc,  0.0, 0.0  // only semifwd    
  };

// -------------------------------------------------------------------------------
//  generated power law
// -------------------------------------------------------------------------------
double csGen[nPoints]; // nb
double csPL[nPoints]; // nb
double suUnc[nPoints]; // stat+unc uncertaity in nb
void generateCS()
{
  // initialise random number generator
  gRandom->SetSeed(0);
  cout << endl << " Random seed: " << gRandom->GetSeed() << endl;
  // generate cross section  
  for(int i=0; i<nPoints;i++) {
    csPL[i] = nPL*TMath::Power(wPoints[i]/wPL,lambdaPL);
    csGen[i] = csPL[i]*gRandom->Gaus(1,staUnc[i]);
    csGen[i] *= gRandom->Gaus(1,uncUnc[i]);
    suUnc[i] = csGen[i]*TMath::Sqrt(staUnc[i]*staUnc[i]+uncUnc[i]*uncUnc[i]);
  }

  // shift full cross section
  double normCS = gRandom->Gaus(1,normUnc);
  for(int i=0; i<nPoints;i++) csGen[i] *= normCS;
  cout << " normalisation shift: " << normCS << endl;

  // shift fwd cross sections
  double normFW =  gRandom->Gaus(1,fwdUnc);
  csGen[0] *= normFW; csGen[1] *= normFW; csGen[2] *= normFW; csGen[8] *= normFW;
  cout << " fwd shift: " << normFW << endl;
  
  // shift sf cross sections
  double normSF =  gRandom->Gaus(1,sfUnc);
  csGen[3] *= normSF; csGen[4] *= normSF; csGen[7] *= normSF;
  cout << "  semi-fwd shift: " << normSF << endl;

  // shift cb cross sections
  double normCB =  gRandom->Gaus(1,cbUnc);
  csGen[5] *= normCB; csGen[6] *= normCB;
  cout << " central barrel shift: " << normCB << endl;
  cout << endl;
}


// -------------------------------------------------------------------------------
// Define graph with error band from results of the fit
// -------------------------------------------------------------------------------

TGraph* GetShade(double N, double d, double cov00, double cov11, double cov01)
{
  Float_t wmin=gxmin;
  Float_t wmax=gxmax;
  int nw = 200;
  Color_t color = 29;
  double dw = (wmax-wmin)/nw;
  TGraph *gshade = new TGraph(2*nw+2);
  for (int i=0;i<=nw;i++){
    double w = wmin+i*dw;
    double y = N*pow(w/90,d);
    double dydN = y/N;
    double dydd = y*log(w/90);
    double dy = sqrt(dydN*dydN*cov00 + dydd*dydd*cov11 + 2*dydN*dydd*cov01);
    gshade->SetPoint(i       ,w,y+dy);
    gshade->SetPoint(2*nw-i+1,w,y-dy);
  }
  gshade->SetFillColor(color);
  gshade->SetLineColor(color);
  return gshade;
}

// -------------------------------------------------------------------------------
// chi2 function to fit data according to prescription in
// Eur.Phys.J. C63 (2009) 625-678, section 9, eq (31)
// -------------------------------------------------------------------------------

void fcnChi2Model(int &npar, double *gin, double &f, double *par, int iflag)

{
  // use unused parguments to avoid compiler message
  double tmp = npar; tmp = gin[0]; tmp = iflag;

  // define model according to:
  // ch2 = sum_i [m_i-mu_i-Sij]^2/D + Sbj
  // Sij = sum_j g_ij*m_i*b_j
  // Sbj = sum_j (b_j)^2
  // D = d_i_stat^2*mu_i*(m_i-Sij)+(d_i,unc*m_i)^2
  // mu_i is the measurement
  // g_ij relative normalization uncertainty at point i from source j
  // d_i = relative uncertainty (either stat or uncorr)

  double bj[nCor] = {par[2],par[3],par[4],par[5]};
  double Sbj = 0;
  for(int j=0;j<nCor;j++) Sbj += (bj[j]*bj[j]);
  double chi2 = 0;
  for (int i=0; i<nPoints;i++) {
    double mu_i = csGen[i];
    double m_i = par[0]*TMath::Power(wPoints[i]/wPL,par[1]);
    double Sij = 0;
    for(int j=0;j<nCor;j++) Sij += m_i*bj[j]*corUnc[j][i];
    double d_stat_i = staUnc[i];
    double d_unc_i = uncUnc[i];
    double D = (d_stat_i*d_stat_i*mu_i*(m_i-Sij))+((d_unc_i*m_i)*(d_unc_i*m_i));
    chi2 += ((m_i-mu_i-Sij)*(m_i-mu_i-Sij)/D);
  }
  chi2 += Sbj;
  f = chi2;
  gChi2 = f;
}



// -------------------------------------------------------------------------------
// Main function
// -------------------------------------------------------------------------------

void powerLawFit()
{
  ///////////////////////////////////////////////////////////////////////////////
  // generate data set
  ///////////////////////////////////////////////////////////////////////////////
  generateCS();
  
  ///////////////////////////////////////////////////////////////////////////////  
  // now do TMinut
  ///////////////////////////////////////////////////////////////////////////////  
  cout << endl<< endl<< endl;
  cout << " - o - o - o - o - o - o - o - o -  Minuit   - o - o - o - o - o - o - o - o -" << endl;
  cout << endl<< endl<< endl;
  // initialize minuit with a maximum of parameters
  const int nPar = 2+nCor;
  TMinuit *myMinuit = new TMinuit(nPar);
    
  // set the function with the minimization process
  myMinuit->SetFCN(fcnChi2Model);   
  
  // define parameters
  double minBj = -10;
  double maxBj = -minBj;  
  myMinuit->DefineParameter(0,"N",70,1,0.0,1000.);
  myMinuit->DefineParameter(1,"#delta",0.7,0.1,0.1,10.);
  myMinuit->DefineParameter(2,"global norm",0.001,0.0001,minBj,maxBj);
  myMinuit->DefineParameter(3,"fwd norm",0.001,0.0001,minBj,maxBj);
  myMinuit->DefineParameter(4,"sf norm",0.001,0.0001,minBj,maxBj);
  myMinuit->DefineParameter(5,"cb norm",0.001,0.0001,minBj,maxBj);      

  // migrad
  myMinuit->SetMaxIterations(500);
  myMinuit->Migrad();

  ///////////////////////////////////////////////////////////////////////////////  
  // get results
  ///////////////////////////////////////////////////////////////////////////////  
  double Cov[nPar*nPar];
  myMinuit->mnemat(Cov,nPar);
  double N = 0;  // normalization
  double Nerr = 0;  // normalization error
  myMinuit->GetParameter(0,N,Nerr);
  double D = 0;  // exponent
  double Derr = 0;  // exponent error
  myMinuit->GetParameter(1,D,Derr);
  double covNN = Cov[0];
  double covDD = Cov[nPar+1];
  double covND = Cov[1];
  double rho = covND/TMath::Sqrt(covDD*covNN);
  cout << " N = " << N << " +/- " << Nerr << endl;
  cout << " D = " << D << " +/- " << Derr << endl;
  cout << " rho = " << rho << endl;
  cout << " gChi2 " << gChi2 << endl;      

  ///////////////////////////////////////////////////////////////////////////////  
  // plot results
  ///////////////////////////////////////////////////////////////////////////////
  // data
  TGraphErrors *grGen = new TGraphErrors(nPoints, wPoints, csGen, NULL,suUnc);
  grGen->SetMarkerColor(2);
  grGen->SetMarkerSize(0.8);
  grGen->SetMarkerStyle(20);
  grGen->SetLineColor(2);
  grGen->SetLineWidth(3);

  // fit
  TGraph* gShade = GetShade(N,D,covNN,covDD,covND);
  TF1* power_law_fit = new TF1("power_law_fit","[0]*pow(x/90,[1])",gxmin,gxmax);
  power_law_fit->SetParameter(0,N);
  power_law_fit->SetParameter(1,D);
  power_law_fit->SetLineColor(4);
  power_law_fit->SetLineWidth(1);


  // Canvas
  TCanvas* c1 = new TCanvas("c1","c1",1000,700);
  TH1F* frame1 = gPad->DrawFrame(gxmin,gymin,gxmax,gymax);
  frame1->GetXaxis()->SetMoreLogLabels();
  gPad->SetLeftMargin(0.09);
  gPad->SetRightMargin(0.01);
  gPad->SetTopMargin(0.11);//0.025
  gPad->SetBottomMargin(0.13);
  gPad->SetLogx();
  gPad->SetLogy();
  gStyle->SetOptFit(111111);
  gStyle->SetOptStat("");
  gStyle->SetOptTitle(0);  
  gStyle->SetPalette(1);
  gStyle->SetLineWidth(2);      //axis line
  gStyle->SetFrameLineWidth(2); //frame line
  TGaxis::SetMaxDigits(3);
  frame1->SetTitle(";W_{#gammap} (GeV);#sigma(#gamma+p #rightarrow J/#psi+p) (nb)");
  Float_t siz = 0.035;
  frame1->SetTitleSize(siz);       frame1->SetLabelSize(siz);
  frame1->SetTitleSize(siz, "Y");  frame1->SetLabelSize(siz, "Y");
  frame1->GetXaxis()->SetTitleOffset(1.4);
  frame1->GetYaxis()->SetTitleOffset(1.);

  // draw
  gShade->Draw("fsame");
  power_law_fit->Draw("csame");
  grGen->Draw("p,e,same");

  // new axis
  double mjpsi = 3.0969;
  double bxmin = TMath::Power((mjpsi/gxmax),2.);
  double bxmax = TMath::Power((mjpsi/gxmin),2.);
  TF1 *fbx = new TF1("fbx","TMath::Power(([0]/x),2.)", bxmin, bxmax);
  fbx->SetParameter(0, mjpsi);

  TGaxis *axis = new TGaxis(gxmax, gymax, gxmin, gymax, "fbx", 510, "G+");
  axis->SetTextFont(42);
  axis->SetLabelFont(42);
  axis->SetTitleSize(siz); axis->SetLabelSize(siz);
  axis->SetLabelOffset(-0.035);
  axis->SetTitleOffset();
  axis->Draw("same");

  TLatex *bxtit = new TLatex();
  bxtit->SetTextFont(42);
  bxtit->SetTextSize(siz);
  bxtit->SetTextAlign(31);
  bxtit->DrawLatex(gxmax, gymax+450, "Bjorken-#it{x}");

  // legends
  //upper legend
  double xl = 0.1, dxl = 0.4;
  double yl = 0.6, dyl = 0.2;
  TLegend* leg1 = new TLegend(xl, yl, xl+dxl, yl+dyl);
  leg1->SetFillStyle(0);
  leg1->SetBorderSize(0);
  leg1->SetTextSize(siz);
  leg1->AddEntry(grGen,"Gen","pe");
  leg1->AddEntry(gShade,"#sigma(W_{#gammap}) = #sigma_{0}(W_{#gammap}/W_{0})^{#delta}","lf");
  leg1->Draw("same");

  //lower legend, alignment to (=) sign
  xl = 480; yl = 80.; dyl = 0.35;
  int item = 0;
  TLatex* lat = new TLatex();
  lat->SetTextFont(42);
  lat->SetTextSize(siz);
  lat->DrawLatex(xl, exp(log(yl)-dyl*item), Form("= %.2f #pm %.2f",D,Derr)); item++;
  lat->DrawLatex(xl, exp(log(yl)-dyl*item), Form("= %.1f #pm %.1f nb",N,Nerr)); item++;
  lat->DrawLatex(xl, exp(log(yl)-dyl*item), "= 90 GeV, fixed"); item++;
  lat->DrawLatex(xl, exp(log(yl)-dyl*item), Form("= %.2f",rho)); item++;
  int ndf = nPoints-nPar;
  lat->DrawLatex(xl, exp(log(yl)-dyl*item), Form("= %1.2f/%d",gChi2,ndf)); item++;

  
  xl = xl - 20.;
  item = 0;
  TLatex* lat1 = new TLatex();
  lat1->SetTextFont(42);
  lat1->SetTextSize(siz);
  lat1->SetTextAlign(31);
  lat1->DrawLatex(xl, exp(log(yl)-dyl*item), "#delta"); item++;
  lat1->DrawLatex(xl, exp(log(yl)-dyl*item), "#sigma_{0}"); item++;
  lat1->DrawLatex(xl, exp(log(yl)-dyl*item), "W_{0}"); item++;
  lat1->DrawLatex(xl, exp(log(yl)-dyl*item), "#rho(#sigma_{0}, #delta)"); item++;
  lat1->DrawLatex(xl, exp(log(yl)-dyl*item), "#chi^{2}/NDF"); item++;

  // plot correlation matrix
  TH2F *corrH = new TH2F("corrH","Correlation matrix",nPar,-0.5,nPar-0.5,nPar,-0.5,nPar-0.5);
  const char *names[nPar] = {"N","#lambda","glo","fwd","sf","cb"};
  cout << " here " << flush <<endl;
  for(int i=0;i<nPar;i++) {
    corrH->GetXaxis()->SetBinLabel(i+1,names[i]);
    corrH->GetYaxis()->SetBinLabel(i+1,names[i]);
    for(int j=0;j<nPar;j++) {
      float corrII = Cov[i+nPar*i];
      float corrJJ = Cov[j+nPar*j];
      float corrIJ = Cov[i+nPar*j];
      float corr = corrIJ/TMath::Sqrt(corrII*corrJJ);
      corrH->SetBinContent(i+1,j+1,corr);
    }
  };
  gStyle->SetPaintTextFormat("4.3f");
  TCanvas* cCorr = new TCanvas("cCorr","cCorr",100,100,800,800);
  cCorr->cd();
  corrH->Draw("zcol,text");

}