#include <math.h>

void plottif(char* estim="harm2")//MTCC
{
  gROOT->Reset();
  gStyle->SetOptStat(000000);
  gStyle->SetCanvasBorderMode(-1);
  gStyle->SetCanvasBorderSize(1);
  gStyle->SetCanvasColor(10);

  allhits("dedx",false);
  //  allhits("N");
  /*
  together(1,estim);
  together(2,estim);
  final(estim);
  together(3,estim);
  together(4,estim);
  plotSimple("p");
  plotSimple("pt");
  plotSimple("eta");
  plotSimple("n");
  */
}

void allhits(char* var="dedx", bool etacut=false) // MTCC
{
  gROOT->Reset();
  gStyle->SetOptStat(000000);
  gStyle->SetCanvasBorderMode(-1);
  gStyle->SetCanvasBorderSize(1);
  gStyle->SetCanvasColor(10);

  TString dt("tif");
  TString mc("mtcc_data_tickmark");

  TCanvas* myCanvas = new TCanvas("myCanvas","all hits",200,10,900,800);
  myCanvas->SetFillColor(10);
  myCanvas->Divide(1,2);
  myCanvas->cd(1);
  plotHistoHits(var,"tib",dt,true,1,"",true);
  plotHistoHits(var,"tib",mc,false,2,"same");
  TPaveLabel *legend1 = new TPaveLabel(0.75,0.80,0.95,0.95,"Data TIF   ","NDC");
  TPaveLabel *legend2 = new TPaveLabel(0.75,0.65,0.95,0.80,"Data MTCC","NDC");
  legend1->Draw();
  legend2->SetTextColor(2);
  legend2->Draw();
  myCanvas->Update();
  myCanvas->cd(2);
  plotHistoHits(var,"tob",dt,true,1,"",true);
  plotHistoHits(var,"tob",mc,false,2,"same");

  TString gif("allhits_");
  gif+=var;
  gif+="_tif";
  gif+=".png";

  myCanvas->Print(gif);

}


void together(int mode=1,char* var="mean") // MTCC
{
  gROOT->Reset();
  //gStyle->SetOptStat(0000);
  gStyle->SetCanvasBorderMode(-1);
  gStyle->SetCanvasBorderSize(1);
  gStyle->SetCanvasColor(10);

  TString name1 = filename("cosmic");
  TString name2 = filename("simu");
  
  TFile f1(name1);
  TFile f2(name2);
  TTree *tree1 = (TTree*)f1.Get("T");
  TTree *tree2 = (TTree*)f2.Get("T");

  bool log=false;
  if (mode == 2) log=true;
  bool norma=false;
  if (mode == 1) norma=true;

  Int_t nbinsx=25;
  Double_t xlow,xup;
  if (log) {xlow=0.5;xup=2.1;} else {xlow=2.;xup=7.;}
  TH1F *h1 = new TH1F("hdata", "dE/dx [MeV/cm]", nbinsx, xlow, xup);
  TH1F *h2 = new TH1F("hsimu", "dE/dx [MeV/cm]", nbinsx, xlow, xup);

  TString gif(var);
  gif+="_mtcc";

  TCanvas* myCanvas = new TCanvas("myCanvas",var,200,10,900,800);
  myCanvas->SetFillColor(18);

  if (mode == 1) {
    // lin scale
    myCanvas->cd();
    fillHisto(h1,var,tree1,norma,log,"data");
    fillHisto(h2,var,tree2,norma,log,"simu",2,"same");
    gif+="_lin.png";
  } else if (mode == 2) {
    // log scale
    myCanvas->Divide(1,2);
    myCanvas->cd(1);
    fillHisto(h1,var,tree1,norma,log,"data");
    myCanvas->cd(2);
    fillHisto(h2,var,tree2,norma,log,"simu");
    gif+="_log.png";
  } else if (mode == 3) {
    // scatter plot dE/dx vs P
    myCanvas->cd();
    gif+="_vsP";
    gif+=".png";
    plotScatter(var,tree1,1);
    plotScatter(var,tree2,2,"same");
  } else if (mode == 4 || mode==5 || mode==6) {
    // profile histo dE/dx vs P
    Int_t nbinsx=30;
    Double_t xlow;
    Double_t xup;
    char* title;
    char* abscissa;
    if (mode == 4) {title="dE/dx [MeV/cm] vs P [GeV]"; abscissa="p"; xlow=0.8; xup=9.8;}
    if (mode == 5) {title="dE/dx [MeV/cm] vs Pt [GeV]"; abscissa="pt"; xlow=0.8; xup=4.8;}
    if (mode == 6) {title="dE/dx [MeV/cm] vs eta"; abscissa="eta"; xlow=0.; xup=1.2;}
    TProfile *prof1 = new TProfile("pdata", title, nbinsx, xlow, xup);
    TProfile *prof2 = new TProfile("psimu", title, nbinsx, xlow, xup);
    myCanvas->cd();
    fillProfile(prof1,var,abscissa,tree1,0,25,0,"data",1);
    fillProfile(prof2,var,abscissa,tree2,0,25,0,"simu",2,"same");
    // legend:
    Double_t xmin=xlow+(xup-xlow)*0.65;
    Double_t xmax=xlow+(xup-xlow)*0.95;
    Double_t ymax=findmax(prof1,nbinsx);
    Double_t ymin=findmin(prof1,nbinsx);
    //    TLegend *legend = new TLegend(xmin,ymin+0.75*(ymax-ymin),xmax,ymin+0.99*(ymax-ymin),"particles:","");
    TLegend *legend = new TLegend(0.65,0.75,1.00,0.95,"particles:","NDC");
    legend->AddEntry(prof1,"cosmics","l");
    legend->AddEntry(prof2,"sim.cosmics","l");
    legend->Draw();
    myCanvas->Update();
    if (abscissa=="p") {gif+="_prof.png";} else if (abscissa=="pt") {gif+="_profPt.png";} else if (abscissa=="eta") {gif+="_profEta.png";} else {gif+="else";}
    delete prof1;
    delete prof2;
  } else {
    cout << "Mode " << mode << " does not exist." << endl;
  }
  
   
  myCanvas->Print(gif);

}

void plotSimple(char* var="eta") // MTCC
{
  gROOT->Reset();

  TString name1 = filename("data");
  TString name2 = filename("simu");
  
  TFile f1(name1);
  TFile f2(name2);
  TTree *tree1 = (TTree*)f1.Get("T");
  TTree *tree2 = (TTree*)f2.Get("T");

  double variable;
  int ivariable;
  if (var=="n") {
    tree1->SetBranchAddress(var,&ivariable);
    tree2->SetBranchAddress(var,&ivariable);
  } else {
    tree1->SetBranchAddress(var,&variable);
    tree2->SetBranchAddress(var,&variable);
  }

  Int_t nbinsx=50;
  if (var=="n") nbinsx=7;
  Double_t xlow,xup;
  if (var=="p") {xlow=0.; xup=10.; }
  if (var=="pt") {xlow=0.; xup=10.; }
  if (var=="eta") {xlow=-2.5; xup=2.5; }
  if (var=="n") {xlow=0; xup=7; }
  TH1F *h1 = new TH1F("hdata", var, nbinsx, xlow, xup);
  TH1F *h2 = new TH1F("hsimu", var, nbinsx, xlow, xup);

  for (int i = 0; i < tree1->GetEntries(); i++) {
    tree1->GetEntry(i);
    if (var=="n") { h1->Fill(ivariable); } else { h1->Fill(variable); }
  }
  for (int i = 0; i < tree2->GetEntries(); i++) {
    tree2->GetEntry(i);
    if (var=="n") { h2->Fill(ivariable); } else { h2->Fill(variable); }
  }

  normalize(h1);
  normalize(h2);

  h1->SetLineColor(1);
  h2->SetLineColor(2);

  h1->SetLineWidth(3);
  h2->SetLineWidth(3);

  TCanvas* myCanvas = new TCanvas("myCanvas",var,200,10,900,800);
  myCanvas->cd();

  h1->DrawCopy();
  h2->DrawCopy("same");

  TString gif(var);
  gif+="_mtcc";
  gif+=".png";
  myCanvas->Print(gif);

  delete h1;
  delete h2;
}

TString filename(char* job="tif"))//MTCC
{
  TString pippo("trackhisto_");
  pippo+=job;
  pippo+=".root";
  return pippo;
}


Double_t findmax(TProfile* prof, Int_t nbins)
{
  Double_t max=0;
  for (unsigned int i=0; i<nbins; i++) {
    Double_t temp1=prof->GetBinContent(i);
    Double_t temp2=prof->GetBinError(i);
    if (temp1+temp2 > max) max=temp1+temp2;
  }
  return max;
}

Double_t findmin(TProfile* prof, Int_t nbins)
{
  Double_t min=99999.;
  for (unsigned int i=0; i<nbins; i++) {
    Double_t temp1=prof->GetBinContent(i);
    Double_t temp2=prof->GetBinError(i);
    if (temp1-temp2 < min) min=temp1-temp2;
  }
  return min;
}

void separation(TGraph* gr, TProfile* prof1, TProfile* prof2, Int_t nbins, Double_t xlow, Double_t xup, char* label="proton/pion separation", char* abscissa="p",int marker=21)
{
  Double_t binwidth=(xup-xlow)/nbins;
  Double_t* x = NULL; 
  Double_t* y = NULL; 
  x = new Double_t[nbins];
  y = new Double_t[nbins];
  for (unsigned int i=0; i<nbins; i++) {
    x[i]=0.;
    y[i]=0.;
  }
  gr = new TGraph(nbins,x,y);
  for (unsigned int i=0; i<nbins; i++) {
    x[i]=xlow+binwidth*i;
    Double_t y1=prof1->GetBinContent(i+1);
    Double_t y2=prof2->GetBinContent(i+1);
    Double_t e1=prof1->GetBinError(i+1);
    Double_t e2=prof2->GetBinError(i+1);
    if (e1>0. && e2>0.) {
      y[i]= (y1-y2)/sqrt(e1**2+e2**2);
      //    cout << "i = "<< i << ", x = " << x[i] << ", y = " << y[i]<< endl;
      gr->SetPoint(i,x[i],y[i]);
    }
  }
  //  TGraph* gr = new TGraph(nbins,x,y);
  //  gr = new TGraph(nbins,x,y);
  //  gr->SetMarkerColor(color);
  gr->SetMarkerStyle(marker);
  gr->SetTitle(label);
  if (abscissa=="p") {gr->GetXaxis()->SetTitle("P [GeV/c]");} else if (abscissa=="pt") {gr->GetXaxis()->SetTitle("P_T [GeV/c]");} else if (abscissa=="\eta") {gr->GetXaxis()->SetTitle("\eta");} else if (abscissa=="n") {gr->GetXaxis()->SetTitle("n.hits");} else {continue;}
  //  gr->GetXaxis()->SetTitle("P [GeV]");
  gr->GetYaxis()->SetTitle("sigma");
  gr->Draw("AP");
  //gr->PaintGraph(nbins,x,y,"A");

  delete [] x;
  delete [] y;
  //delete gr;
}


void fillHisto(TH1F* h, char* var="mean", TTree *tree, double B, bool norma, bool log, char* label="data", int color=1, char* superimpose="")//MTCC
{
  gROOT->Reset();
  gStyle->SetOptStat(0000);
  gStyle->SetCanvasBorderMode(-1);
  gStyle->SetCanvasBorderSize(1);
  gStyle->SetCanvasColor(10);

  double variable;
  double p,eta;
  tree->SetBranchAddress(var,&variable);
  tree->SetBranchAddress("p",&p);
  tree->SetBranchAddress("eta",&eta);
  
  for (int i = 0; i < tree->GetEntries(); i++) {
    tree->GetEntry(i);
    if (log) {
      if (variable>0) h->Fill(log(variable));
    } else {
      h->Fill(variable);
    }
  }

  // normalization:
  if (norma) normalize(h);

  h->SetLineColor(color);
  h->SetLineWidth(3);
  h->DrawCopy(superimpose);
  if (log) {
    // perform a gaussian fit:
    h->Fit("gaus");
    TF1 *fit = h->GetFunction("gaus");
    Double_t chi2 = fround(fit->GetChisquare(),1);
    unsigned ndecimals=3;
    Double_t p1 = fround(fit->GetParameter(1),ndecimals);
    Double_t e1 = fround(fit->GetParError(1),ndecimals);
    Double_t p2 = fround(fit->GetParameter(2),ndecimals);
    Double_t e2 = fround(fit->GetParError(2),ndecimals);
    cout << "--------------------" << endl;
    cout << "Result of the fit:" << endl;
    cout << "--------------------" << endl;
    cout << "chi2 = " << chi2 << endl;
    cout << "mean = " << p1 << "+-" << e1 << endl;
    cout << "sigma = " << p2 << "+-" << e2 << endl;
    cout << "--------------------" << endl;
    // draw a legend with the fit results:
    Double_t xmin=1.5;
    Double_t xmax=2.0;
    Double_t ymax=h->GetMaximum();
    //    TPaveText *text = new TPaveText(xmin,0.65*ymax,xmax,0.95*ymax,"br");
    TPaveText *text = new TPaveText(0.55,0.55,1.0,0.95,"NDC");
    TString t0("chi2 = ");
    t0+=chi2;
    TString t1("mean = ");
    t1+=p1;
    t1+="+-";
    t1+=e1;
    TString t2("sigma = ");
    t2+=p2;
    t2+="+-";
    t2+=e2;
    text->AddText(xmin,0.65*ymax,label);
    text->AddText(xmin,0.65*ymax,t0);
    text->AddText(xmin,0.65*ymax,t1);
    text->AddText(xmin,0.65*ymax,t2);
    text->Draw();
    // write the fit results in an ascii file
    TString ascii("txt/");
    ascii+=var;
    ascii+="_";
    ascii+=label;
    ascii+=".txt";   
    fstream f1;
    f1.open(ascii,ios::out);
    f1 << chi2 << endl;
    f1 << p1 << endl;
    f1 << p2 << endl;
    f1.close();
    cout << "Fit results written into: " << ascii << endl;
  }
  
  //  delete h;
}


void final(char* var="mean") //MTCC
{
  const int numfiles = 2;
  Double_t chi2[numfiles], p1[numfiles], p2[numfiles];
  TString label[numfiles]={"data","simu"};
  Double_t separation[numfiles];

  cout << "***************************" << endl;
  cout << "Estimator: " << var << endl;
  cout << "***************************" << endl;
  for (unsigned int i=0; i<numfiles; i++) {
    TString ascii("txt/");
    ascii+=var;
    ascii+="_";
    ascii+=label[i];
    ascii+=".txt";
    cout << "Reading file: " << ascii << endl;
    ifstream f1;
    f1.open(ascii);
    if (f1.good()) {
      f1 >> chi2[i];
      f1 >> p1[i];
      f1 >> p2[i];
      f1.close();
      if (i>0) separation[i]=(p1[i]-p1[0])/sqrt(p2[i]**2 + p2[0]**2);
    } else {
      cout << ascii << " doesn't exist!" << endl;
    }
  }
  cout << "data chi2: " << chi2[0] << endl;
  cout << "data sigma: " << 100.*p2[0]/p1[0] << " %" << endl;
  cout << "MC/data discrepancy: " << separation[1] << " sigmas" << endl;
  cout << "***************************" << endl;

}

void plotScatter(char* var="mean", TTree *tree, int color=1, char* superimpose="")
{
  gROOT->Reset();
  gStyle->SetOptStat(0000);
  gStyle->SetCanvasBorderMode(-1);
  gStyle->SetCanvasBorderSize(1);
  gStyle->SetCanvasColor(10);
  
  double variable;
  double p,eta;
  tree->SetBranchAddress(var,&variable);
  tree->SetBranchAddress("p",&p);
  tree->SetBranchAddress("eta",&eta);

  TH2F *h = new TH2F("h", "dE/dx [MeV/cm] vs P (GeV)",100,0.4,7.0,100,1.5,6.0);

  for (int i = 0; i < tree->GetEntries(); i++) {
    tree->GetEntry(i);
    h->Fill(p,variable);
  }
  
  h->SetMarkerColor(color);
  h->SetMarkerStyle(23); // 21: squares, 20, 22: asymmetric stars, 23: triangles (pointing downwards)
  h->SetMarkerSize(0.5);
  h->DrawCopy(superimpose);

  delete h;
}

/* round number n to d decimal points */
double fround(double n, unsigned d)
{
      return floor(n * pow(10., d) + .5) / pow(10., d);
}

void normalize(TH1F* h)
{
  Double_t scale;
  if (h->Integral()!=0) scale = 1/h->Integral();
  h->Scale(scale);
}

void plotHistoHits(char* var="dedx", char* type="tib", char* job="tif", bool dofit=false, int color=1, char* superimpose="", bool marker=false)//MTCC
{
  gROOT->Reset();
  gStyle->SetOptStat(000000);
  gStyle->SetCanvasBorderMode(-1);
  gStyle->SetCanvasBorderSize(1);
  gStyle->SetCanvasColor(10); // no more grey bkg

  TString name1 = filename(job);
  TFile* f1 = new TFile(name1);

  TString histo("h");
  histo+=var;
  histo+="_";
  histo+=type;

  TH1F *h = (TH1F*)f1->Get(histo);
  if (marker) {
    h->SetMarkerStyle(21);
    h->SetMarkerSize(0.7);
    h->Sumw2();   // this makes sure that the sum of squares of weights will be stored
  }
  h->SetLineColor(color);
  h->SetLineWidth(3);
  h->Rebin(2); // temp
  h->SetMinimum(0);
  normalize(h);
  if (dofit) {
    h->Fit("landau","m","",1.8,6);
    TF1 *fit = h->GetFunction("landau");
    Double_t chi2 = fround(fit->GetChisquare(),2);
    unsigned ndecimals=3;
    Double_t p1 = fround(fit->GetParameter(1),ndecimals);
    Double_t e1 = fround(fit->GetParError(1),ndecimals);
    Double_t p2 = fround(fit->GetParameter(2),ndecimals);
    Double_t e2 = fround(fit->GetParError(2),ndecimals);
    cout << "--------------------" << endl;
    cout << "Result of the fit:" << endl;
    cout << "--------------------" << endl;
    cout << "chi2  = " << chi2 << endl;
    cout << "peak  = " << p1 << "+-" << e1 << endl;
    cout << "width = " << p2 << "+-" << e2 << endl;
    cout << "--------------------" << endl;
  }

  // axis labels:
  char* xlabel = "dE/dx [MeV/cm]";
  char* ylabel = "Arbitrary units";
  h->GetXaxis()->SetTitle(xlabel);
  h->GetYaxis()->SetTitle(ylabel);

  if (marker) {
    h->DrawCopy("e1");
  } else {
    h->DrawCopy(superimpose);
  }
  
  delete h;
}


vector<double> regularize(vector<double> v)
{
  vector<double> w;
  w.clear();
  double n=v.at(0);
  double eta=v.at(1);
  double p=v.at(2);
  double dpdr=v.at(3);
  double dp2=v.at(4);
  double dedx=v.at(5);
  w.push_back((1.*n-15.)/8.);
  w.push_back(eta/2.4);
  w.push_back((atan(p)-1.2)/.4);
  w.push_back((dpdr-0.007)/0.005);
  w.push_back(dp2/0.1);
  w.push_back(log(dedx)-1.);
  return w;
}

vector<double> invert(vector<double> v)
{
  vector<double> w;
  w.clear();
  double n=v.at(0);
  double eta=v.at(1);
  double p=v.at(2);
  double dpdr=v.at(3);
  double dp2=v.at(4);
  double dedx=v.at(5);
  w.push_back(8.*n+15.);
  w.push_back(eta*2.4);
  w.push_back(tan(0.4*p+1.2));
  w.push_back(dpdr*0.005+0.007);
  w.push_back(dp2*0.1);
  w.push_back(exp(dedx+1.));
  return w;
}