#include <iostream.h>
#include <fstream.h>
#include <stdlib.h>
#include <time.h>

#include <TROOT.h>
#include <TH1.h>
#include <TH2.h>
#include <TFile.h>
#include <TObject.h>
#include <TPythia6.h>
#include <TMCParticle.h>
#include <TTree.h>
#include <TRandom.h>
#include <TVector3.h>
#include <TLorentzVector.h>
#include <TCanvas.h>
#include <TArrow.h>
#include <TLatex.h>
#include <TArc.h>

#include "header.h"
#include "TRunCard.h"
#include "TSiDetector.h"
#include "TVertex.h"
#include "TTrackFinder.h"
#include "TMCTrack.h"
#include "TSiTrack.h"

double pi = TMath::Pi();
 

Float_t DeltaPhi(Float_t phi1, Float_t phi2);
int chargedHadron(TMCParticle* );
 

int main(int argc, char **argv) {

  TROOT root("readtree","reads a pythia ttree file");

  Int_t   numberEvents;   // number of events to generate
  if ( argc==2 ) {
    numberEvents = atoi(argv[1]);
  } else if ( argc==3 ){
    numberEvents = atoi(argv[1]);
    if ( argv[2][0]=='k' || argv[2][0]=='K' ) numberEvents *= 1000;
    if ( argv[2][0]=='m' || argv[2][0]=='M' ) numberEvents *= 1000000;
  }

 
  header *hdr = new header;
  TMCParticle *particle = new TMCParticle;
  TClonesArray *part = new TClonesArray("TMCParticle",500);
 

  TFile *fout = new TFile("output.root","RECREATE","pythia analysis output");

  //====================================================
  // histogramming
  //====================================================
  TH1F * hnpart  = new TH1F("hnpart","hnpart",100,0.,200.);
  TH1F * HPt     = new TH1F("HPt","pt  ",100, 0, 15);
  TH1F * HEta    = new TH1F("HEta","eta  ",100, -3, 3);
  TH1F * HPhi    = new TH1F("HPhi","phi inclus ",100,-pi, pi);
  TH1F * hprocid = new TH1F("hprocid","procid  ",100, -0.5, 99.5);

  TH1F * htracks = new TH1F("htracks","no Si tracks  ",100, -0.5, 99.5);
  TH1F * hdispx  = new TH1F("hdispx","disp x   ",200, -0.5, 0.5);
  TH1F * hdispy  = new TH1F("hdispy","disp y   ",200, -0.5, 0.5);
  TH1F * hdisprxy= new TH1F("hdisprxy","disp rxy",200,  0.0, 1.5);
  TH1F * hdispr  = new TH1F("hdispr","disp r   ",200,  0.0, 1.5);
  TH1F * hvr     = new TH1F("hvr","true r   ",200,  0.0, 1.5);

 
  //====================================================
  // detector definition and single vertex and track instances
  //====================================================
  TRunCard runCard("config/siveta.cfg");
  TSiDetector *sidet[12];
  int isi, noPlanes=12;
  for(isi=0;isi<noPlanes;isi+=2){
    sidet[isi]   = new TSiDetector("config/sidet1.setup", isi*30, runCard.GetVerbose());
    sidet[isi+1] = new TSiDetector("config/sidet2.setup", isi*30, runCard.GetVerbose());
  }

  TMCTrack     *MCTrack     = new TMCTrack();
  TClonesArray *MCtrackList = new TClonesArray("TMCTrack",500);
  TClonesArray *trackList   = new TClonesArray("TSiTrack",500);

  TVertex *vertex   = new TVertex();
  TTrackFinder *trackFinder = new TTrackFinder();
  TRandom ranGen(21341233);
 
  //====================================================
  // graphics
  //====================================================
  TCanvas *c1;
  if(runCard.GetVerbose()==2) {
    c1 = new TCanvas("c1", "c1",506,8,598,598);
    c1->SetFillColor(0);
    c1->SetBorderMode(0);
    c1->Range(-10, -10, 10, 10);
    TArrow *axX = new TArrow(-10,0,10,0,0.02,">"); axX->SetLineStyle(2); axX->Draw();
    TArrow *axY = new TArrow(0,-10,0,10,0.02,">"); axY->SetLineStyle(2); axY->Draw();
    c1->cd();

    for(isi=0;isi<noPlanes;isi++) sidet[isi]->Draw();
    TLatex *tex; TArrow *arrow;
    tex = new TLatex(7.44511,-1.02725,"x (p_{x})");
    tex->SetTextSize(0.04); tex->Draw();
    tex = new TLatex(0.658683,9.32914,"y (p_{y})");
    tex->SetTextSize(0.04); tex->Draw();
 
    tex = new TLatex(-9.,8.,"charm tracks");
    tex->SetTextColor(2); tex->SetTextSize(0.03); tex->Draw();
    arrow = new TArrow(-4.8, 8.2, -3.8, 8.2, 0.01,">");
    arrow->SetFillColor(1);arrow->SetLineColor(2); arrow->Draw();

    tex = new TLatex(-9.,7.5,"tracks");
    tex->SetTextColor(1); tex->SetTextSize(0.03); tex->Draw();
    arrow = new TArrow(-4.8, 7.7, -3.8, 7.7, 0.01,">");
    arrow->SetFillColor(1);arrow->SetLineColor(1); arrow->Draw();
 
    tex = new TLatex(-9.,7.0,"hits");
    tex->SetTextColor(1); tex->SetTextSize(0.03); tex->Draw();
    TArc *dhit = new TArc(-4.8, 7.2, 0.1);
    dhit->SetLineColor(6);dhit->Draw();

    c1->Update();
  }

 
  //====================================================
  // Access the event header branch and the particle information array branch
  //====================================================
  for (int ifile =0; ifile<runCard.GetNoFiles(); ifile++) {
 
    cout <<endl<<"Reading  " << runCard.GetFileName(ifile) <<" file, ";
    TFile fin(runCard.GetFileName(ifile),"READ");
    TTree *t = (TTree *)fin.Get("pyttree");  // access the ttre

    TBranch *b[2];
    b[0] = t->GetBranch("event");
    b[0]->SetAddress(&hdr);
    b[1] = t->GetBranch("part");
    b[1]->SetAddress(&part);
 
    Stat_t nevents = t->GetEntries();
    printf("there are %d events.\n", (int)nevents);

    if ( argc<2 ) numberEvents = (int)nevents;
    int ieout = numberEvents/10;
    if (ieout<1) ieout=1;
    cout << "Looping over "<<numberEvents<<" events"<<endl;
 

    //====================================================
    // Loop over all events in the TTree
    //====================================================
    for (Int_t evt=0; evt<numberEvents; evt++) {
 
      t->GetEvent(evt);
      Int_t npart = hdr->GetNpart();
      Int_t  procID = hdr->GetProcessid();
 
      hnpart->Fill(npart);
      hprocid->Fill(procID);
 
      // Create the vertex;  100um beamprofile and  +-30cm in z flat
      // vertex->SetRPhiZ(ranGen.Gaus(0,0.01), ranGen.Rndm()*2*pi-pi, ranGen.Rndm()*60-30 );
      vertex->SetRPhiZ(ranGen.Gaus(0,1.21), ranGen.Rndm()*2*pi-pi, ranGen.Rndm()*60-30 );
      //vertex->SetRPhiZ(0.9, 50*pi/180.0, ranGen.Rndm()*60-30 );

      //====================================================
      // Loop over all particles in event
      //====================================================
      int nChargedTracks=0;
      for (Int_t i=0; i<npart; i++) {
        particle = (TMCParticle *)part->At(i);
 
 Float_t Pt,Eta,Phi;
  //Int_t   type       = particle->GetKF();
 Pt  = sqrt(pow(particle->GetPx(),2)+pow(particle->GetPy(),2));
 Eta = -log(tan(atan2(Pt,particle->GetPz())/2));
 Phi = atan2(particle->GetPy(),particle->GetPx());
 Int_t charge = chargedHadron(particle);

 if(charge!=0){
   //TMCTrack     *MCTrack     = new TMCTrack();
   TLorentzVector mom(particle->GetPx(),
        particle->GetPy(),
        particle->GetPz(),
        particle->GetEnergy());
   MCTrack->Set4P(mom);
   TVector3 sumVertex(particle->GetVx(),particle->GetVy(),particle->GetVz());
   //TVector3 sumVertex(0,0,0);
   hvr->Fill(sumVertex.Mag()/10.0);
   if(sumVertex.Mag()>0) MCTrack->SetCharm(1); else  MCTrack->SetCharm(0);
   sumVertex = (sumVertex*0.1) + vertex->GetVector();   //Pythia vertex is in mm!!!!!
   MCTrack->SetV(sumVertex);
   new ((*MCtrackList)[nChargedTracks++]) TMCTrack(*MCTrack);
 }
 HPt->Fill(Pt);
 HEta->Fill(Eta);
 HPhi->Fill(Phi);
      }

      //make hits----------------------------------------------
      for(isi=0;isi<noPlanes;isi++) sidet[isi]->makeHit(MCtrackList);
      //-------------------------------------------------------
 

      //make tracks--------------------------------------------
      for(isi=0;isi<noPlanes;isi+=2)          trackFinder->makeSimpleTracks(sidet[isi], sidet[isi+1], trackList);
      //trackFinder->makeSimpleTracks(&sidet1, &sidet2, trackList);
       //-------------------------------------------------------
 

      //analyse tracks-----------------------------------------
      TSiTrack  *track, *track2;
      htracks->Fill(trackList->GetEntries());
 
      for (int itr=0; itr<trackList->GetEntries(); itr++) {
 track = (TSiTrack*)trackList->At(itr);
 if(runCard.GetVerbose()==2) {
   track->Draw();
   c1->Update();
 }
 //cout<<track->GetMomentumVector().Mag()<<endl;
 TVector3 disp = vertex->GetDisplacementVector( track );
 hdispx->Fill(disp.X());
 hdispy->Fill(disp.Y());
 hdisprxy->Fill( sqrt(disp.X()*disp.X()+disp.Y()*disp.Y()) );
 hdispr->Fill(disp.Mag());
 
 // Attempt to determine Displaced Vertex
 if(disp.Mag()>0.001){
   for (int itr2=0; itr2<trackList->GetEntries(); itr2++) {
     if(itr != itr2){
        track2 = (TSiTrack*)trackList->At(itr2);
  TVector3 dispt = vertex->GetDisplacementVector( track, track2 );
  cout<<dispt.Mag()<<endl;
     }
   }
 }
 //--------------------------------------

      }
      //--------------------------------------------------------
 

      MCtrackList->Clear();
      trackList->Clear();
      for(isi=0;isi<noPlanes;isi++) sidet[isi]->ClearHitList();

      if (evt % ieout == 0) cout<<evt<<" "<<float(evt)/numberEvents*100<<"%"<<endl;
    }
    fin.Close();
  }
 
  if(runCard.GetVerbose()==2) {
    cout<<"Save siveta.eps ? y/(n or RET)"<<endl;
    if(getchar()=='y')
      c1->Print("siveta.eps");
      c1->Print("siveta.root");
  }
 
  fout->Write();
  fout->Close();

}
 
 

Float_t DeltaPhi(Float_t phi1, Float_t phi2)
{
 
  Float_t res = phi1 - phi2;
  if (fabs(res) > pi) res = atan2(sin(res), cos(res));
  return res;
}

int chargedHadron(TMCParticle* particle) {
  int charged = 0;
  const char* name = particle->GetName();
  TString nameSt(name);
  if (nameSt.Contains("pi+")) {
    charged=+1;
  } else if (nameSt.Contains("pi-")) {
    charged=-1;
  } else if (nameSt.Contains("K-")) {
    charged=-1;
  } else if (nameSt.Contains("K+")) {
    charged=+1;
  }  else if (nameSt.Contains("pbar")) {
    charged=-1;
  } else if (nameSt.Contains("p+")) {
    charged=+1;
  }
 
  return charged;
}