//
#include <TLatex.h>
#include <TCanvas.h>
#include <TPad.h>
#include <TF1.h>
#include <TH1F.h>
#include <TH2F.h>
#include <TMath.h>
#include <TFile.h>
#include <TStyle.h>
#include <TROOT.h>
#include <TTree.h>
#include <string>
#include "B2DPi.C"
#include "B2D3Pi.C"

int nevUpdate = 20;
int ShortPause = 5000000;
int LongPause = 50000000;

int ShowAllVertices = -1;  // Set to 1 to show all PVs


double minFlightDistOfB = 10.0;
double minFlightDistOfD = 10.0;

void setHistProp(TH1F* h)
{
  h->SetLineWidth(2);
  h->SetNdivisions(505,"X");
  h->SetNdivisions(505,"Y");
  h->SetXTitle(xtitle);
  h->SetYTitle(ytitle);
  h->GetXaxis()->SetTitleSize(0.05);
  h->GetYaxis()->SetTitleSize(0.05);
  h->GetXaxis()->SetLabelSize(0.05);
  h->GetYaxis()->SetLabelSize(0.05);
  h->GetXaxis()->SetTitleOffset(1.0);
  h->GetYaxis()->SetTitleOffset(1.0);
}

void pause(int n)
{
  for(int i=0;i<n;i++){continue;}
  return;
}


void run5()
{  

  // tt.Loop();  

  char input;  // a character variable

  gStyle->SetStatW(0.4);
  gStyle->SetStatH(0.2);
  gStyle->SetTitleSize(0.3);
  gStyle->SetPadLeftMargin(0.13);
  gStyle->SetPadBottomMargin(0.13);  
  
  // Allows histograms to stay in scope after leaving this macro.
  TFile *f = new TFile("temp.root", "RECREATE");

  int iProc;
  cout << "Which type of decay mode would you like to look at?" << endl;
  cout << " --> For  B -->D pi,        enter 1" << endl;
  cout << " --> For  B -->D pi pi pi,  enter 2" << endl;
  cout << " --> Please answer (1 or 2) : ";
  
  cin >> iProc;
  if(iProc < 1 || iProc>2) break;

  int nAnim;
  cout << "How many event animations do you wish to see [ after this, just fill histograms ]?" << endl;
  cin >> nAnim;
  
  B2DPi ttPi;
  B2D3Pi tt3Pi;
  
  // Book histograms
  bmass = new TH1F("bmass","B#rightarrowD#pi Mass",60,5000,5600);
  benergy = new TH1F("benergy","B#rightarrowD#pi Energy (GeV)",50,0,500);
  bdecaytime = new TH1F("bdecaytime","B#rightarrowD#pi Decay Time (ps)",50,0,10);
  bflightdist = new TH1F("bflightdist","B#rightarrowD#pi Flight Distance (mm)",50,0,50);
  if(iProc == 2){
    bmass->SetTitle("B#rightarrowD#pi#pi#pi Mass");
    benergy->SetTitle("B#rightarrowD#pi#pi#pi Energy (GeV)");
    bdecaytime->SetTitle("B#rightarrowD#pi#pi#pi Decay Time (ps)");
    bflightdist->SetTitle("B#rightarrowD#pi#pi#pi Flight Distance (mm)");
  }
  setHistProp(bmass, "Mass (GeV/c^{2})", "Number/10 MeV/c^{2}");
  setHistProp(benergy, "Energy (GeV)", "Number/10 GeV");
  setHistProp(bdecaytime, "Decay Time (ps)", "Number/0.2 ps");
  setHistProp(bflightdist, "Distance (mm)", "Number/mm");  

  // Number of Entries
  Long64_t nentries = ttPi.fChain->GetEntriesFast();
  if(iProc==2) nentries = tt3Pi.fChain->GetEntriesFast();

  // Create Canvases
  TCanvas *cp = new TCanvas("cp","Distributions",900,0,500,500);
  cp->SetFillColor(kYellow-9);
  cp->Divide(2,2);

  // Loop over events
  int j = 0;
  for(int i=0; i<nentries;i++){
      if(iProc==1){        
        Long64_t ientry = ttPi.LoadTree(i);
        if (ientry < 0) break;
        ttPi.fChain->GetEntry(i);
        ttPi.LoadToLocalVariables();
      }else if(iProc==2){
        Long64_t ientry = tt3Pi.LoadTree(i);
        if (ientry < 0) break;
        tt3Pi.fChain->GetEntry(i);
        tt3Pi.LoadToLocalVariables();
      }
      
      double dzb = BVERTEX_z - BPV_z;
      double dzd = DVERTEX_z - BVERTEX_z;
    
      // Use flight distance cuts only for the first few events that we're animating
      if(j < nAnim){
        if(dzb<minFlightDistOfB) continue;
        if(dzd<minFlightDistOfD) continue;
      }
      
	  TCanvas *ccxz; 
      TCanvas *ccxy;
      
      if(B_MCMatch >=30) continue;
      if(j < nAnim){
        // Create Event Display
        ccxz = drawPictureXZ_Anim(iProc);
        ccxy = drawPictureXY(iProc);
        ccxz->Draw();
        ccxy->Draw();
        ccxz->Modified();
	    ccxy->Modified();
        ccxz->ForceUpdate();
        ccxy->ForceUpdate();
      }


      // Fill histograms
      bmass->Fill(B_Mass);
      benergy->Fill(B_Energy);
      bdecaytime->Fill(B_DecayTime);
      bflightdist->Fill(B_FlightDist);

      // Draw Histograms
      cp->cd(1)->SetFillColor(kGreen-9); bmass->DrawCopy();
      cp->cd(2)->SetFillColor(kGreen-9); benergy->DrawCopy();
      cp->cd(3)->SetFillColor(kGreen-9); bdecaytime->DrawCopy();
      cp->cd(4)->SetFillColor(kGreen-9); bflightdist->DrawCopy();
      // Only update every N events (N = nevUpdate value), OR for the first N animated events
      if(j%nevUpdate == 0 || j<nAnim) cp->ForceUpdate();   
      
      
      if(j < nAnim){
        cout << " Continue to next event (y or n)";
        cin >> input;
        if (input == 'n' || input=='N') break;
      }
      j++;
  }
}

TCanvas* drawPictureXZ_Anim(int iProc)
{
  TCanvas *c =  new TCanvas("c", "", 50, 0, 800, 600);  


  double zl = BPV_z + zneg;
  double zh = DVERTEX_z + zpos;
  double xl = (BVERTEX_x + BPV_x)/2.0 - xscale;
  double xh = (BVERTEX_x + BPV_x)/2.0 + xscale;      

  c->Range(zl, xl, zh, xh);
  c->SetFillColor(kBlack);


  double zEnd = 250.0;
  cout << " ===================================================" << endl;
  cout << "Number of Primary Vertices = " << NumPV << endl;

  //----------------------------------------------
  // Animate the protons coming into the B vertex
  //----------------------------------------------
  double xsiz = 0.15;
  int nstepb = 20;
  
  static const int nproton = 14;
  
  double zsiz = xsiz*(zh-zl)/(xh-xl);
  TEllipse* e1[nproton];
  TEllipse* e2[nproton];
  for(int j=0; j<nproton; j++){
    e1[j] = new TEllipse(0, 0, zsiz, xsiz, 0, 360,0);          
    e2[j] = new TEllipse(0, 0, zsiz, xsiz, 0, 360,0);          
  }

  // Draw "star" as sort of explosion
  TMarker *m = new TMarker(BPV_z, BPV_x, 29);
  
  double pposz[nproton]={0,10,20,30,5, 15, 25, 5,15,25,10,20,10,20};
  double pposx[nproton]={0,0,0,0,0.25,0.25,0.25,-0.25,-0.25,-0.25,0.5,0.5,-0.5,-0.5};

  double zstepl = abs(zl-BPV_z)/nstepb;
  double zstepr = abs(zh-BPV_z)/nstepb;      

  // Get closest point between left and right beams
  double mindiff = 1000;
  for(int i=0; i<2*nstepb; i++){
    double poszl0 = zl+i*zstepl;    
    double poszr0 = zh-i*zstepr;
    double diff = fabs(poszl0-poszr0);
    if(diff < mindiff) mindiff = diff;
  }
  
  int beampause = 600000;
  bool changepause = false;
  for(int i=0; i<4*nstepb; i++){
    double poszl0 = zl+i*zstepl;    
    double poszr0 = zh-i*zstepr;
    for(int j=0; j<nproton; j++){
      double poszl = poszl0 + pposz[j];
      double poszr = poszr0 + pposz[j];
      double posxl = pposx[j];
      double posxr = pposx[j];
      e1[j]->SetFillStyle(1001);
      e1[j]->SetFillColor(kWhite);
      e2[j]->SetFillStyle(1001);
      e2[j]->SetFillColor(kWhite);
      e1[j]->SetX1(poszl);
      e1[j]->SetY1(posxl);
      e2[j]->SetX1(poszr);
      e2[j]->SetY1(posxr);
      e1[j]->Draw();
      e2[j]->Draw();      
    }
    if(changepause) beampause = 10000;
    pause(beampause);
   
    if(fabs(poszl0-poszr0) < mindiff + 1){
      changepause = true;  
      for(int k=0; k<6; k++){
        m->SetMarkerSize(10*k);    
        m->SetMarkerColor(kRed);
        if(k%2==0) m->SetMarkerColor(kGreen);
        m->Draw();
        c->ForceUpdate();
        pause(1000000);
      }
    }
    //m->SetX(10000);
    m->SetMarkerSize(0);
    //pause(100);
    c->ForceUpdate();
    if(poszl0 > zh && poszr0 < zl) break;
    
  }
  // Now make them go off screen
  for(int j=0; j<nproton; j++){
    e1[j]->SetX1(-10000);e2[j]->SetX1(10000);
    //pause(100000);    
    e1[j]->Draw(); e2[j]->Draw();
  }
  c->ForceUpdate();    
    
  //----------------------------------------------------
  // Now animate the tracks coming out of the B vertex
  //----------------------------------------------------

  TLine* el[200];
  TLine* lb;
  TLine* ld;
  TLine* lbdau[3];
  TLine* lddau[3];

  for(int i=0; i<200; i++){
    el[i] = new TLine(0,0,10,10);
  }
  lb = new TLine(0,0,10,10);  
  ld = new TLine(0,0,10,10);  
  for(int i=0; i<3; i++){
    lbdau[i] = new TLine(0,0,10,10);  
    lddau[i] = new TLine(0,0,10,10);  
  }
  

  
  int nstep = 20;

  int iTrk = nTracksBPV;
  for(int j=0; j<nstep; j++){
    for(int i=0; i<iTrk; i++){
        double xsl = xslopeBPV[i];
        double ysl = yslopeBPV[i];
        double xint = xintcBPV[i];
        double yint = yintcBPV[i];
        double zf = zFinalBPV[i];
        double zstepl = (zf-BPV_z)/nstep;
        double zz = BPV_z + j*zstepl;
        double posz = zz;
        double posx = BPV_x + xsl*(zz - BPV_z);
        double z1 = BPV_z;
        double z2 = zz+zstepl/2.;
        double x1 = BPV_x + xsl*(z1 - BPV_z);
        double x2 = BPV_x + xsl*(z2 - BPV_z);
        el[i]->SetX1(z1);el[i]->SetY1(x1);el[i]->SetX2(z2);el[i]->SetY2(x2);
        el[i]->SetLineColor(14);
        el[i]->Draw();

    }    
    // Now animate the B particle track
    double zz = BPV_z + j*zstepl;
    double posz = zz;
    double posx = BPV_x + xslopeB*(zz - BPV_z);
    double r = sqrt((posx-BPV_x)*(posx-BPV_x)+(posz-BPV_z)*(posz-BPV_z));
    double rb = sqrt((BVERTEX_x-BPV_x)*(BVERTEX_x-BPV_x)+(BVERTEX_z-BPV_z)*(BVERTEX_z-BPV_z));
    double rd = sqrt((DVERTEX_x-BPV_x)*(DVERTEX_x-BPV_x)+(DVERTEX_z-BPV_z)*(DVERTEX_z-BPV_z));
    if(r < rb){
      double zstepl = (zh - BPV_z)/nstep;
      double zz = BPV_z + j*zstepl;
      double z1 = BPV_z;
      double z2 = zz+zstepl/2.;
      if(z2 > BVERTEX_z) z2 = BVERTEX_z;
      double x1 = BPV_x + xslopeB*(z1 - BPV_z);
      double x2 = BPV_x + xslopeB*(z2 - BPV_z);
      lb->SetX1(z1);lb->SetY1(x1);lb->SetX2(z2);lb->SetY2(x2);      
      lb->SetLineColor(kViolet);
      lb->SetLineWidth(2);
      lb->SetLineStyle(7);  
      lb->Draw();      
    }
    // If raidus larger than B vertex radius, and less then D vertex radius, animate the D
    if(r>rb && r<rd){
      double zstepl = (zh - BPV_z)/nstep;
      double zz = BPV_z + j*zstepl;
      double z1 = BVERTEX_z;
      double z2 = zz+zstepl/2.;
      if(z2 > DVERTEX_z) z2 = DVERTEX_z;
      if( (z2 + zstepl/2.) >  DVERTEX_z) z2 = DVERTEX_z;
      double x1 = BVERTEX_x + xslopeD*(z1 - BVERTEX_z);
      double x2 = BVERTEX_x + xslopeD*(z2 - BVERTEX_z);
      ld->SetX1(z1);ld->SetY1(x1);ld->SetX2(z2);ld->SetY2(x2);      
      ld = new TLine(z1,x1,z2,x2);
      ld->SetLineColor(kOrange);
      ld->SetLineWidth(2);  
      ld->SetLineStyle(2);  
      ld->Draw();
    }   
   
    // Draw bachelor tracks associated with the B vertex
    if(r>rb){
      for(int k=0; k<nBachTracks; k++){
        double zstepl = (zh - BPV_z)/nstep;
        double zz = BPV_z + j*zstepl;
        double z1 = BVERTEX_z;
        double z2 = zz+zstepl/2.;
        double x1 = BVERTEX_x + xslopeBdau[k]*(z1 - BVERTEX_z);
        double x2 = BVERTEX_x + xslopeBdau[k]*(z2 - BVERTEX_z);
        lbdau[k]->SetX1(z1);lbdau[k]->SetY1(x1);lbdau[k]->SetX2(z2);lbdau[k]->SetY2(x2);
        lbdau[k]->SetLineColor(kBlue);
        lbdau[k]->SetLineWidth(1);
        lbdau[k]->Draw();
      }
    }
    // Draw D daughter tracks associated with the D vertex
    if(r>rd){
      for(int k=0; k<3; k++){
        double zstepl = (zh - BPV_z)/nstep;
        double zz = BPV_z + j*zstepl;
        double z1 = DVERTEX_z;
        double z2 = zz+zstepl/2.;
        double x1 = DVERTEX_x + xslopeDdau[k]*(z1 - DVERTEX_z);
        double x2 = DVERTEX_x + xslopeDdau[k]*(z2 - DVERTEX_z);
        lddau[k]->SetX1(z1);lddau[k]->SetY1(x1);lddau[k]->SetX2(z2);lddau[k]->SetY2(x2);
        lddau[k]->SetLineColor(kRed);
        lddau[k]->SetLineWidth(1);
        lddau[k]->Draw();
      }
    }
    
    pause(1000000);
    
    c->ForceUpdate();
  }

  if(ShowAllVertices > 0){
    pause(1000000);
    // Now add additional PVs
    for(int j=0; j<NumPV; j++){
      int iTrk = nPVTracks[j];
      if(fabs(PV_z[j]-BPV_z)<0.01) continue;
      
      for(int i=0; i<iTrk; i++){    
        double xsl = xslopePV[i][j];
        double ysl = yslopePV[i][j];
        double xint = xintcPV[i][j];
        double yint = yintcPV[i][j];
        double zf = zFinalPV[i][j];
        double xf = xsl*zf + xint;
        double yf = ysl*zf + yint;        
        l = new TLine(PV_z[j],xint, zf, xf);     
        l->SetLineColor(12);
        l->SetLineWidth(1);
        l->Draw();      
      }    
      c->ForceUpdate();
    }
    
  }
    
    
  double dx = 0.05*(xh-xl);
  double dx1 = 0.035*(xh-xl);
  double dx2 = 0.07*(xh-xl);
  
  TLatex tt;
  tt.SetTextAlign(22);
  tt.SetTextSize(0.04);
  tt.SetTextColor(kBlue);
  tt.DrawLatex(BVERTEX_z+0.5,BVERTEX_x+dx1 ,"B_{vertex}");
  tt.SetTextColor(kRed);
  tt.DrawLatex(DVERTEX_z+ 0.2 , DVERTEX_x+dx2 , "D_{vertex}");

  //TLatex t;
  tt.SetTextSize(0.08); 
  tt.SetTextAlign(12);
  tt.SetTextColor(kRed);
  tt.DrawLatex(zl+15,xl+dx,"LHCb Event Display");

  double z1 = BPV_z-5;
  double z2 = z1 + 10;
  double zz = (z1 + z2)/2.;
  double xx1 = xl+3.5*dx;
  double xx2 = xx1 - dx;
  TLine *lz = new TLine(z1,xx1,z2,xx1);
  lz->SetLineColor(kWhite);
  lz->SetLineWidth(2);
  lz->Draw();
  tt.SetTextSize(0.03); 
  tt.SetTextColor(kWhite);
  tt.SetTextAlign(22);
  tt.DrawLatex((z1+z2)/2.,xx2,"1 cm");
  double xxa = xx1 + 1.0;
  double xxb = xx1 - 1.0;
  TLine *lx = new TLine(zz,xxa,zz,xxb);
  lx->SetLineColor(kWhite);
  lx->SetLineWidth(2);
  lx->Draw();
  tt.SetTextSize(0.03); 
  tt.SetTextColor(kWhite);
  tt.SetTextAlign(22);
  tt.SetTextAngle(90);
  tt.DrawLatex(z2+5.0,xx1,"2 mm");
  

  c->ForceUpdate();
  
  return c;

}

TCanvas* drawPictureXY(int iProc)
{
  TCanvas *cc =  new TCanvas("cc", "", 0, 0, 300, 300);  

  double xl = (BPV_x+DVERTEX_x)/2 - xscaleXY;
  double xh = xl + 2*xscaleXY;
  double yl = (BPV_y + DVERTEX_y)/2.0 - xscaleXY;
  double yh = yl+2*xscaleXY;
  cc->Range(xl, yl, xh, yh);
  cc->SetFillColor(kBlack);
  
  cout << "  B Prim. Vertex Pos (x,y,z): " << BPV_x << " " << BPV_y << " " << BPV_z << endl;
  cout << "  B Vertex Pos (x,y,z):       " << BVERTEX_x << " " << BVERTEX_y << " " << BVERTEX_z << endl;
  cout << "  D Vertex Pos (x,y,z):       " << DVERTEX_x << " " << DVERTEX_y << " " << DVERTEX_z << endl;
  
  
  double xEnd = 20.0;  
  if(B_Px<0) xEnd = -20.0;
  

  TLine * l; 

  int iTrk = nPVTracks;

  cout << "Number of proton-proton Collsions = " << NumPV << endl;
  
  TArrow *a;
  for(int j=0; j<NumPV; j++){
    int iTrk = nPVTracks[j];
    bool bVertex = false;
    if(fabs(PV_z[j]-BPV_z)<0.01) bVertex = true;
    if(!bVertex) cout << " ---> pp Collision # " << j << " has " << iTrk << " tracks" << endl;
    if(bVertex) cout << " ---> pp Collision # " << j << " has " << iTrk << " tracks (B Primary Vertex)" << endl;
    if(!bVertex) continue;    
    
    for(int i=0; i<iTrk; i++){    
      double xsl = 6.28*(gRandom->Rndm()-0.5);
      double ysl = 6.28*(gRandom->Rndm()-0.5);
      xsl = tan(xsl);
      ysl = tan(ysl);      
      double xint = 0.03*gRandom->Gaus();
      double yint = 0.03*gRandom->Gaus();
      double xint = PV_x[j]+xint;      
      double yint = PV_y[j]+yint;      
      double r = gRandom->Rndm();
      double zf = xEnd;
      if(r<0) zf = -xEnd;
      double xf = xsl*zf + xint;
      double yf = ysl*zf + yint;

      l = new TLine(xint, yint, xf, yf);     
      l->SetLineColor(12);
      l->SetLineWidth(1);
      if(bVertex) l->SetLineColor(14);
      l->Draw();
      //cout << i << endl; 
      
    }
    cc->ForceUpdate();
    if(fabs(PV_z[j]-BPV_z)<0.01) bVertex = true;
    
  }
  //Draw B particle track

  l = new TLine(BPV_x,BPV_y,BVERTEX_x,BVERTEX_y);
  l->SetLineStyle(7);
  l->SetLineWidth(3);
  l->SetLineColor(kViolet);  
  l->Draw();
  cc->ForceUpdate();
  
  // Draw line for D trajectory
  l = new TLine(BVERTEX_x,BVERTEX_y,DVERTEX_x,DVERTEX_y);
  l->SetLineStyle(7);
  l->SetLineWidth(3);
  l->SetLineColor(kOrange);  
  l->Draw();
  cc->ForceUpdate();

  if(iProc==1){
    double ysl = PI_Py/PI_Px;
    double yint = PI_POSy-ysl*PI_POSx;
    double yf = ysl*xEnd+yint;
    l = new TLine(BVERTEX_x,BVERTEX_y, xEnd, yf); 
    l->SetLineColor(kBlue); l->SetLineWidth(2);
    l->Draw();
    cc->ForceUpdate();
  }else if(iProc==2){
  // Draw PiPiPi daughters
    for(int i=0; i<3; i++){
      double ysl = BDAU_Py[i]/BDAU_Px[i];
      double yint = BDAU_POSy[i]-ysl*BDAU_POSx[i];
      double xEnd2 = 10.0;  
      if(BDAU_Px[i]<0) xEnd2 = -10.0;
      double yf = ysl*xEnd2+yint;
      l = new TLine(BVERTEX_x, BVERTEX_y, xEnd2, yf); 
      l->SetLineColor(kBlue); l->SetLineWidth(2);
      l->Draw();    
    }
  }
  
    
  

  // Draw D daughters
  for(int i=0; i<3; i++){ 
    double ysl = DDAU_Py[i]/DDAU_Px[i];
    double yint = DDAU_POSy[i]-ysl*DDAU_POSx[i];
    double xEnd2 = 10.0;  
    if(DDAU_Px[i]<0) xEnd2 = -10.0;
    double yf = ysl*xEnd2+yint;
    l = new TLine(DVERTEX_x, DVERTEX_y, xEnd2, yf); 
    l->SetLineColor(kRed); l->SetLineWidth(2);
    l->Draw();    
  }

  TLatex ttt;
  ttt.SetTextSize(0.06); 
  ttt.SetTextAlign(12);
  ttt.SetTextColor(kGreen);
  
  ttt.DrawLatex(xl+0.5,yl+0.4,"LHCb Event Display (XY View)");
   
  cc->ForceUpdate();
  return c;

}