///////////////////////////////////////////////////////////
//
//  These functions are used to investigate how correlated 
//  momenta affect delta(p) and sigma(p)
//
///////////////////////////////////////////////////////////

///////////////////////////////////////////////////////////////////
//  etcorrel: this function is used to plot the correlation between 
//  line number and energy of particles.
//  It also makes a line of best fit.
///////////////////////////////////////////////////////////////////
void etcorrel(TString fName, const Int_t stepSize=200){

  cout << fName << endl;

  TNtuple *lnt = makeGpNtuple(fName.Data());
  Float_t e1,e2,x,y,z;
  lnt->SetBranchAddress("e1",&e1);
  lnt->SetBranchAddress("e2",&e2);
  lnt->SetBranchAddress("x",&x);
  lnt->SetBranchAddress("y",&y);
  lnt->SetBranchAddress("z",&z);

  TCanvas *c = new TCanvas("c","c",900,600); c->Divide(1,3);
  TCanvas *c1 = new TCanvas("c1","c1",900,600);//c1->Divide(1,2);

  Int_t ncols;
  Int_t nlines = 0;
  //  const Int_t stepSize = 200;

  e1array = TArrayF(stepSize);
  e2array = TArrayF(stepSize);
  zarray = TArrayF(stepSize);

  Int_t aIndex = 0;
  Float_t e1av;
  Float_t e2av;
  Float_t zav;

  Int_t nPlotPoints = 0;

  const Int_t nMax = 10000;
  Float_t numPlotArray[nMax], e1PlotArray[nMax],e2PlotArray[nMax],zPlotArray[nMax];


  for(Int_t j=0;j<lnt->GetEntries();j++) {
    lnt->GetEntry(j);

    //every 10 steps take averages and add to plotpoints array
    if (((j % stepSize) == 0) && (j != 0) ){

      aIndex = 0;
      e1av = 0;
      e2av = 0;
      zav = 0;
      //make averages: adding
      for(Int_t i=0; i < stepSize; i++){
   e1av += e1array[i];
      }
      for(Int_t i=0; i < stepSize; i++){
   e2av += e2array[i];
      }
      for(Int_t i=0; i < stepSize; i++){
   zav += zarray[i];
      }
      //make averages: normalising
      e1av = e1av / stepSize;
      e2av = e2av / stepSize;
      zav = zav / stepSize;
      //write values to arrays that will be plotted
      e1PlotArray[nPlotPoints]= e1av;
      e2PlotArray[nPlotPoints]= e2av;
      zPlotArray[nPlotPoints]= zav;

      numPlotArray[nPlotPoints]= nPlotPoints;
      nPlotPoints++;

    }

    //put in stepsize-array to make averages later
    e1array[aIndex]=e1;
    e2array[aIndex]=e2;
    zarray[aIndex]=z;

    //cout << "got to line" << nlines << endl;
    aIndex++;
    nlines++;
  }

  printf(" found %d points \n",nlines);


  //make 1% error bars
  const Int_t n1 = nPlotPoints;
  Float_t ere1[n1],ere2[n1];
  for (Int_t i=0; i<nPlotPoints;i++){
    ere1[i]= 0.01 * e1PlotArray[i];
    ere2[i]= 0.01 * e2PlotArray[i];
  }

  e1graph = new TGraphErrors( nPlotPoints, numPlotArray, e1PlotArray,0,ere1);
  e2graph = new TGraphErrors( nPlotPoints, numPlotArray, e2PlotArray,0,ere2);
  //  e2graph = new TGraphErrors( nPlotPoints, numPlotArray, e2PlotArray,0,ere2);
  zgraph = new TGraph( nPlotPoints, numPlotArray, zPlotArray);

  c->cd(1);
  e1graph->SetTitle("beam 1 energy vs line number");
  e1graph->Draw("ALP");

  c->cd(2);
  e2graph->SetTitle("beam 2 energy vs line number");
  e2graph->Draw("ALP");

  c->cd(3);
  zgraph->SetTitle("z distance vs line number");
  zgraph->Draw("ALP");

  c1->cd();
  e1graph->Draw("A*");
  e1graph->Fit("pol1");
  e1graph->SetTitle("beam 1 energy vs line number");
  e1graph->GetXaxis()->SetTitle("T");
  e1graph->GetXaxis()->SetTitleOffset(0.65);
  e1graph->GetXaxis()->SetTitleSize(0.065);
  e1graph->GetYaxis()->SetTitle("Energy of e+ lumi pairs");
  e1graph->GetYaxis()->SetTitleOffset(0.65);
  e1graph->GetYaxis()->SetTitleSize(0.065);
  e1graph->Draw("A*");

  cout << "error on 5th is " << e1graph->GetErrorX(5) << endl;

  c1->Print("../tmp-results/fitcorrel.eps");
  c->Print("../tmp-results/correlplots.ps");
}


////////////////////////////////////////////////
//   useful little function to reset the range
////////////////////////////////////////////////
void histrange(histname){
  histname->SetMinimum(histname->GetMinimum()-2);
  histname->SetMinimum(histname->GetMinimum()+1);
}

///////////////////////////////////////////////////////////////
//
//  ecorrel: this function makes a 2-D surface plot of E1 vs E2
//           and also of delta(p) vs sigma(p).
//           First arguement is filepath, but useless
//           The second argument drawArgs can be any combination
//           of 1,2,3, and decides which canvasses get displayed:
//           1: 2D histogram output of e1 vs e2 and del vs sig
//           2: profiles of e2 vs e1, del vs e1, sigma vs e1
//           3: profile of delta p  vs sigma p
///////////////////////////////////////////////////////////////

void ecorrel(TString filePath, TString drawArgs){


  filePath = "/home/fz/gp/run/nlcb-500-nooffset/";
  TString fileName = "lumi.dat";
  TString fName = filePath.Append(fileName);
  cout << fName << endl;
  FILE *fp = fopen( fName.Data() ,"r");

  if(drawArgs.Contains("1")){
    TCanvas *c = new TCanvas("c","2d hist output",900,600);
    c->Divide(2,1);
  }
  if(drawArgs.Contains("2")){
    TCanvas *c1 = new TCanvas("c1","profiles of e1 vs e2",900,600);
    c1->Divide(1,4);
  }
  if(drawArgs.Contains("3")){
    TCanvas *c2 = new TCanvas("c2","profiles of delta p  vs sigma p",900,600);
  }

  Float_t e1,e2,sigma_p,delta_p,x,y,z;
  Int_t ncols;
  Int_t nlines = 0;
  const Int_t nMax = 30000;
  const Int_t avSize = 1;

  Float_t e1tmparr[avSize], e2tmparr[avSize];
  Float_t e1avarr[nMax],e2avarr[nMax];
  Float_t e1arr[nMax], e2arr[nMax];
  Float_t e1av, e2av, rstrue, rsr1;
  Int_t aIndex = 0;

  //do a first run to get the number of data points
  while (1) {
     ncols = fscanf(fp,"%f %f %f %f %f", &e1, &e2, &x, &y, &z);
     if (ncols < 0) break;
     //if (nlines <10) cout << e1 << " " << e1 << endl;
     nlines++;
  }
  printf(" found %d points \n",nlines);
  fclose(fp);

  //create a histogram 
  TH2D *histen2d = new TH2D("histen2d","e1 vs e1 2d histogram",nlines/100,215,246,nlines/100,215,246);
  TH2D *histp2d = new TH2D("histp2d","sigma_e vs delta_e 2d histogram",nlines/100,-90,-40,nlines/100,450,490 ); //360,450 for tail study
  TProfile *e2prof = new TProfile("e2","Profile of e2 versus e1",16,230,246,230,246);
  TProfile *del_pprof = new TProfile("delta_p","Profile of delta_p versus e1",16,230,246,-40,40);
  TProfile *sig_pprof = new TProfile("sigma_p","Profile of sigma_p versus e1",16,230,246,460,490);
  TProfile *delsigprof = new TProfile("delsig","Profile of sigma_p versus delta_p",20,-40,40,450,490);
  TProfile *dldroots = new TProfile("dldroots", "Profile of Dl_D(root_S)",16,230,246);

 //now fill the histogram and do the work
  FILE *fp = fopen( fName.Data() ,"r");
  nlines = 0;
  ncols = 0;

  while(1){
    ncols = fscanf(fp,"%f %f %f %f %f", &e1, &e2, &x, &y, &z);
    e1avarr[nlines] = e1;
    e2avarr[nlines] = e2;
    if (ncols < 0) break;

    //we take averages for avSize number of points
    //avsize should be =1 usually, but can be set to anything.
    if ( ((nlines % avSize) == 0) && (nlines != 0) ){
      e1av = 0;
      e2av = 0;

      for (Int_t i=0; i < avSize ; i++){
   e1av += e1tmparr[i];
   e2av += e2tmparr[i];
      }
      //if ((nlines % 2000)== 0) printf(" line %d e1avbig %d \n",nlines,e1av);
      e1av = e1av / avSize;
      e2av = e2av / avSize;
      //if ((nlines % 2000)== 0) printf(" line %d e1av %d \n",nlines,e1av);

      sigma_p = e1av + e2av;//2*pow(e1av*e2av,0.5); //e1av + e2av;
      delta_p = e1av - e2av;
      rstrue = 2*pow(e1av*e2av,0.5);
      rsr1 = 2*pow(e1av*e2av,0.5);

      //fill 2D histogram
      histen2d->Fill(e1av,e2av);
      histp2d->Fill(delta_p,sigma_p);
      e2prof->Fill(e2av,e1av);
      del_pprof->Fill(e1av,delta_p);
      sig_pprof->Fill(e1av,sigma_p);
      delsigprof->Fill(delta_p,sigma_p);

      //reset averages counter
      aIndex = 0;
    }

    e1tmparr[aIndex]= e1;
    e2tmparr[aIndex]= e2;
    aIndex++;
    nlines++;
  }
  fclose(fp);
  printf("finished filling histogram \n");

  //setting the ranges right
  histrange(sig_pprof);
  histrange(del_pprof);
  histrange(delsigprof);
  histrange(e2prof);

  if(drawArgs.Contains("1")){
    //  draw the contents of c
    c->cd(1);
    histen2d->Draw();
    c->cd(2);
    histp2d->Draw();
    c->Print("../tmp-results/ecorrel.ps");
  };

  if(drawArgs.Contains("2")){
    //draw the contents of c1
    c1->cd(1);
    e2prof->SetMarkerColor(4);
    e2prof->SetMarkerStyle(22);
    e2prof->Draw("ALPE1");
    //gStyle->SetEndErrorSize(10);
    c1->cd(2);
    del_pprof->SetMarkerColor(4);
    del_pprof->SetMarkerStyle(22);
    del_pprof->Draw("ALPE1");
    c1->cd(3);
    sig_pprof->SetMarkerColor(4);
    sig_pprof->SetMarkerStyle(22);
    sig_pprof->Draw("ALPE1");
    c1->cd(4);
    delsigprof->SetMarkerColor(4);
    delsigprof->SetMarkerStyle(22);
    delsigprof->Draw("ALPE1");
    c1->Print("../tmp-results/e1-vs-e2_del-vs-sig.ps");
  };

  if(drawArgs.Contains("3")){
    //draw the contents of c2
    c2->cd(0);
    c2->Print("../tmp-results/del-vs-sig_profile.ps");
  };


}