#include <cstdlib>
#include "Stntuple/obj/TGenpBlock.hh"


ClassImp(TGenpBlock)

//______________________________________________________________________________
void TGenpBlock::ReadV1(TBuffer &R__b)
{
  // Read V1 of class TGenpBlock 

  R__b >> fNParticles;
  if (fNParticles > 0) {
    R__b >> fNInteractions;
					// make sure we're not overwriting 
					// memory on read
    fIntData->Set(2*fNInteractions+2);
    R__b.ReadFastArray(fIntData->GetArray(),2*(fNInteractions+1));
    fList->Streamer(R__b);
  }
}

//______________________________________________________________________________
void TGenpBlock::ReadV2(TBuffer &R__b)
{
  ReadParticles(R__b);
}

int TGenpBlock::ReadParticles(TBuffer &R__b) {
//-----------------------------------------------------------------------------
//  read version 2 of the block
//-----------------------------------------------------------------------------
  int           npart, dima, dimb, loca, locb;
  int           idhep, isthep, jm1, jm2, jd1, jd2;

  TGenParticle* p;

  R__b >> npart;
  fList->Clear();
  if (npart > 0) {

    R__b >> fNInteractions;
    fIntData->Streamer(R__b);
    
    dima =  6*npart;
    dimb = 11*npart;

    ushort* a1  = new ushort[dima];
    ushort* a2  = new ushort[dima];
    float* b1   = new float [dimb];

    R__b.ReadFastArray(a1,dima);
    R__b.ReadFastArray(a2,dima);
    R__b.ReadFastArray(b1,dimb);

    for (int i=0; i<npart; i++) {
      loca = 6*i;
      locb = 11*i;
	  
      idhep  = (a1[loca  ] | (a2[loca  ] << 16));
      isthep = (a1[loca+1] | (a2[loca+1] << 16));
      jm1    = (a1[loca+2] | (a2[loca+2] << 16));
      jm2    = (a1[loca+3] | (a2[loca+3] << 16));
      jd1    = (a1[loca+4] | (a2[loca+4] << 16));
      jd2    = (a1[loca+5] | (a2[loca+5] << 16));
      
      p = new ((*fList)[i]) TGenParticle();
      p->Init(i,idhep,isthep,jm1,jm2,jd1,jd2,
	      b1[locb+ 1],b1[locb+ 2],b1[locb+ 3],b1[locb+ 4],
	      b1[locb+ 5],b1[locb+ 6],b1[locb+ 7],b1[locb+ 8]);
      p->SetWeight    (b1[locb   ]);
      p->SetPolarTheta(b1[locb+ 9]);
      p->SetPolarPhi  (b1[locb+10]);
    }

    delete [] a1;
    delete [] a2;
    delete [] b1;
    
    fNParticles = npart;
  }

  return 0;
}

int TGenpBlock::WriteParticles(TBuffer &R__b) {
//-----------------------------------------------------------------------------
//  write version 2 of the block
//-----------------------------------------------------------------------------
  int     loca, locb;

  R__b << fNParticles;
  if (fNParticles > 0) {
    R__b << fNInteractions;
    fIntData->Streamer(R__b);
//-----------------------------------------------------------------------------
// write out list of particles: first integers, then floats, no doubles.
// assume that most integers have only 2 bytes meaningful
// writing 17*4 = 68 bytes per particle
//-----------------------------------------------------------------------------
    int    dima =  6*fNParticles;
    int    dimb = 11*fNParticles;

    ushort* a1  = new ushort[dima];
    ushort* a2  = new ushort[dima];
    float*  b1  = new float [dimb];
//-----------------------------------------------------------------------------
// fill arrays, then write them out
//-----------------------------------------------------------------------------
    for (int i=0; i<fNParticles; i++) {
      TGenParticle* p = Particle(i);
	
      loca = 6*i;
	
      a1[loca  ] = (p->GetPdgCode()          ) & 0xffff;
      a2[loca  ] = (p->GetPdgCode()     >> 16) & 0xffff;
      a1[loca+1] = (p->GetStatusCode()       ) & 0xffff;
      a2[loca+1] = (p->GetStatusCode()  >> 16) & 0xffff;
      a1[loca+2] = (p->GetMother(0)        ) & 0xffff;
      a2[loca+2] = (p->GetMother(0)   >> 16) & 0xffff;
      a1[loca+3] = (p->GetMother(1)        ) & 0xffff;
      a2[loca+3] = (p->GetMother(1)   >> 16) & 0xffff;
      a1[loca+4] = (p->GetDaughter(0)      ) & 0xffff;
      a2[loca+4] = (p->GetDaughter(0) >> 16) & 0xffff;
      a1[loca+5] = (p->GetDaughter(1)      ) & 0xffff;
      a2[loca+5] = (p->GetDaughter(1) >> 16) & 0xffff;
	
      locb = 11*i;
	
      b1[locb   ]  = p->GetWeight();
      b1[locb+ 1]  = p->Px();
      b1[locb+ 2]  = p->Py();
      b1[locb+ 3]  = p->Pz();
      b1[locb+ 4]  = p->Energy();
      b1[locb+ 5]  = p->Vx();
      b1[locb+ 6]  = p->Vy();
      b1[locb+ 7]  = p->Vz();
      b1[locb+ 8]  = p->T();
      b1[locb+ 9]  = p->PolarTheta();
      b1[locb+10]  = p->PolarPhi();
    }

    R__b.WriteFastArray(a1,dima);
    R__b.WriteFastArray(a2,dima);
    R__b.WriteFastArray(b1,dimb);
    
    delete [] a1;
    delete [] a2;
    delete [] b1;
  }
  return 0;
}

//______________________________________________________________________________
void TGenpBlock::Streamer(TBuffer &R__b)
{
   // Stream an object of class TGenpBlock as compact, as possible
  int nwf = &fEventWeight - &fEOR;

  if (R__b.IsReading()) {
    Version_t R__v = R__b.ReadVersion(); 
    if (R__v == 1) { 
      ReadV1(R__b);
    }
    else if (R__v == 2) {
      ReadParticles(R__b);
    }
    else if (R__v == 3) {
      ReadParticles(R__b);
      R__b.ReadFastArray(&fEventWeight,nwf);
    }
  } 
  else {
    R__b.WriteVersion(TGenpBlock::IsA());
    WriteParticles(R__b);
    R__b.WriteFastArray(&fEventWeight,nwf);
  }
}

//_____________________________________________________________________________
TGenpBlock::TGenpBlock() {
  fNParticles    = 0;
  fNInteractions = 0;
  fNCodes        = 0;
  fEtaMin        = -100.;
  fEtaMax        = -100.;
  fPtMin         = 0.;
  fIntData       = new TStnArrayI(20);
  (*fIntData)[0] = 0;
  (*fIntData)[1] = 0;
  fList          = new TClonesArray("TGenParticle",100);
  fList->BypassStreamer(kFALSE);
  fUseCuts       = 0;
  fEventWeight   = 0;
  fPthat         = 0;

}


//_____________________________________________________________________________
TGenpBlock::~TGenpBlock() {
  delete fIntData;
  fList->Delete();
  delete fList;
}


//_____________________________________________________________________________
void TGenpBlock::Clear(const char* opt) {
  fNParticles    = 0;
  fNInteractions = 0;
  fEventWeight   = 0.;
  fPthat         = 0.;

  fIntData->Clear();
					// don't modify cut values at run time
  fList->Clear();
}

//_____________________________________________________________________________
Int_t TGenpBlock::AddPrimaryVertex() {
  // add new primary vertex to the list

  fIntData->Set(2*fNInteractions+3);
  (*fIntData)[2*fNInteractions+2] = First(fNInteractions);
  fNInteractions++;
  return 0;
}


//_____________________________________________________________________________
void TGenpBlock::AddPdgCode(Int_t code) {
  // add new code to the list of codes for the particles to be stored

  if (fNCodes >= fPdgCodeList.fN) {
					// reallocate memory
    fPdgCodeList.Set(2*(fNCodes+1));
  }
  fPdgCodeList.fArray[fNCodes++] = code;
}


//_____________________________________________________________________________
TGenParticle* 
TGenpBlock::NewParticle(Int_t i, Int_t idhep, Int_t istdhep, 
			Int_t m1, Int_t m2, Int_t d1, Int_t d2, 
			Float_t px, Float_t py, Float_t pz, Float_t e,
			Float_t vx, Float_t vy, Float_t vz, Float_t t)
{
  // add new particle to the block. Block is filled sequentially, so 
  // assume that this is the last particle (not necessarily!) and it is 
  // added to the last primary interaction

  TGenParticle* p;

  (*fIntData)[2*fNInteractions]++;

  p = new ((*fList)[i]) TGenParticle();
  p->Init(i,idhep,istdhep,m1,m2,d1,d2,px,py,pz,e,vx,vy,vz,t);

  fNParticles = fList->GetEntriesFast();

  return p;
}


//_____________________________________________________________________________
TGenParticle* TGenpBlock::NewParticle(Int_t Ip, const TParticle* P)
{
  // add new particle to the block. Block is filled sequentially, so this is 
  // the last particle (not necessarily!) and it is added to the last primary 
  // interaction

  TGenParticle* p;

  (*fIntData)[2*fNInteractions]++;
  if (Ip < fNParticles) p = Particle(Ip);
  else                  p = new ((*fList)[fNParticles++]) TGenParticle();

  if (P) p->Init(Ip,P);
  return p;
}


//_____________________________________________________________________________
void TGenpBlock::Print(const char* opt) const {
  int banner_printed(0), i1, i2;

  TGenpBlock* genp = (TGenpBlock*) this;
  for (int ii=0; ii<fNInteractions; ii++) {
    printf(" ------------------------ interaction %i --------------\n",ii);
    i1 = genp->First(ii);
    i2 = i1+genp->NParticles(ii);

    for (int i=i1; i<i2; i++) {
      TGenParticle* p = genp->Particle(i);
      if (! banner_printed) {
	p->Print("banner");
	banner_printed = 1;
      }
      p->Print("data");
    }
  }
}


//_____________________________________________________________________________
void TGenpBlock::GetMissingEt(TVector2* Met) const {
  // calculate GENP-level missing Et in the event using all the neutrinos

  Double_t sum_px(0), sum_py(0);

  TGenpBlock* genp = (TGenpBlock*) this;

  for (int i=0; i<fNParticles; i++) {
    TGenParticle* p = genp->Particle(i);
    if ((p->GetStatusCode() == 1) && p->IsNeutrino()) {
      sum_px += p->Px();
      sum_py += p->Py();
    }
  }
  Met->Set(sum_px,sum_py);
}

//_____________________________________________________________________________
void TGenpBlock::PrintParticleWithDaughters(TGenParticle* P) const {

  TGenParticle* d;
  TGenpBlock* block = (TGenpBlock*) this;

  P->Print("data");

  int nd = P->GetNDaughters();
  if (nd <= 0) return;

  int d1 = P->GetFirstDaughter();
  int d2 = P->GetLastDaughter ();

  for (int i=d1; i<=d2; i++) {
    d = block->Particle(i);
    PrintParticleWithDaughters(d);
  }
} 



//_____________________________________________________________________________
void TGenpBlock::Print(Int_t   PdgCode, 
		       Float_t PtMin  , 
		       Float_t EtaMin, 
		       Float_t EtaMax) const 
{
  int banner_printed;
  banner_printed = 0;
  TGenParticle* p;

  TGenpBlock* block = (TGenpBlock*) this;
  for (int i=0; i<fNParticles; i++) {
    p = block->Particle(i);
    if ((abs(p->GetPdgCode()) == PdgCode) && 
	(p->Pt() > PtMin ) && 
	(p->Eta() >= EtaMin) && 
	(p->Eta() < EtaMax )    ) {

      if (! banner_printed) {
	p->Print("banner");
	banner_printed = 1;
      }

      PrintParticleWithDaughters(p);
    }
  }
}