#include <cstdlib>
#include "TDatabasePDG.h"
#include "Stntuple/obj/TObspBlock.hh"


ClassImp(TObspBlock)

//______________________________________________________________________________
void TObspBlock::ReadV1(TBuffer &R__b)
{
  // read version 1 of TObspBlock

  R__b >> fNParticles;
  if (fNParticles > 0) {
    fListOfParticles->Streamer(R__b);
    R__b >> fNVertices;
    if (fNVertices > 0) {
      fListOfVertices->Streamer(R__b);
    }
    R__b >> fNPrimaryVertices;
    if (fNPrimaryVertices > 0) {
      fListOfPrimaryVertices->Streamer(R__b);
    }
  }
}

//______________________________________________________________________________
int TObspBlock::ReadPrimaryVertices(TBuffer &R__b) {
  int npv;
  R__b >> npv;
  if (npv <= 0) return 0;
					// provision for future
  fNPrimaryVertices = npv;

  return 0;
}

//______________________________________________________________________________
int TObspBlock::WritePrimaryVertices(TBuffer &R__b) {
  R__b << fNPrimaryVertices;
  if (fNPrimaryVertices <= 0) return 0;
					// provision for future
  return 0;
}


//______________________________________________________________________________
int TObspBlock::ReadVertices(TBuffer &R__b) {
//-----------------------------------------------------------------------------
// read list of vertices - fNVertices > 0 
// first integers, then floats
// assume that most integers have only 2 bytes meaningful, so split each 
// integer into 2 halves; write 4*8 = 32 bytes per vertex
//-----------------------------------------------------------------------------
  int    nv, im, iv, id1, id2, loca, locb;

  R__b >>  nv;
  if (nv <= 0) return 0;

  int    dima =  4*nv;
  int    dimb =  4*nv;

  short* a1   = new short[dima];
  short* a2   = new short[dima];
  float* b1   = new float[dimb];
//-----------------------------------------------------------------------------
// read the data into arrays
//-----------------------------------------------------------------------------
  R__b.ReadFastArray(a1,dima);
  R__b.ReadFastArray(a2,dima);
  R__b.ReadFastArray(b1,dimb);
//-----------------------------------------------------------------------------
// unpack the data
//-----------------------------------------------------------------------------
  TObsvVertex* v;
  for (int i=0; i<nv; i++) {
    loca = 4*i;
	
    iv    = a1[loca  ] | (a2[loca]   << 16);
    im    = a1[loca+1] | (a2[loca+1] << 16);
    id1   = a1[loca+2] | (a2[loca+2] << 16);
    id2   = a1[loca+3] | (a2[loca+3] << 16);
	
    locb = 4*i;
    v = NewVertex(iv,b1[locb],b1[locb+1],b1[locb+2],b1[locb+3],im,id1,id2);
  }

  delete [] a1;
  delete [] a2;
  delete [] b1;

  return 0;
}

//______________________________________________________________________________
int TObspBlock::WriteVertices(TBuffer &R__b)
{
//-----------------------------------------------------------------------------
// write out list of vertices: first integers, then floats
// assume that most integers have only 2 bytes meaningful, so split each 
// integer into 2 halves; write 4*8 = 32 bytes per vertex
//-----------------------------------------------------------------------------
  //  int nv;

  R__b << fNVertices;

  if (fNVertices <= 0) return 0;

  int    dima =  4*fNVertices;
  int    dimb =  4*fNVertices;

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

  int loca, locb;
//-----------------------------------------------------------------------------
// fill arrays, then write them out
//-----------------------------------------------------------------------------
  for (int i=0; i<fNVertices; i++) {
    TObsvVertex* v = Vertex(i);

    loca = 4*i;
	
    a1[loca  ] = (v->Number()             ) & 0xffff;
    a2[loca  ] = (v->Number()        >> 16) & 0xffff;
    a1[loca+1] = (v->Mother()             ) & 0xffff;
    a2[loca+1] = (v->Mother()        >> 16) & 0xffff;
    a1[loca+2] = (v->FirstDaughter()      ) & 0xffff;
    a2[loca+2] = (v->FirstDaughter() >> 16) & 0xffff;
    a1[loca+3] = (v->LastDaughter ()      ) & 0xffff;
    a2[loca+3] = (v->LastDaughter () >> 16) & 0xffff;
	
    locb = 4*i;
	
    b1[locb   ]  = v->X();
    b1[locb+ 1]  = v->Y();
    b1[locb+ 2]  = v->Z();
    b1[locb+ 3]  = v->T();
  }

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

  delete [] a1;
  delete [] a2;
  delete [] b1;

  return 0;
}


//______________________________________________________________________________
int TObspBlock::ReadParticles(TBuffer &R__b) {

  int    np, pdg_code, iv1, iv2, genp, dc, loca1, loca2, locb;
  float  charge;

  R__b >> np;
  if (np <= 0) return 0;

  int    dima1 =  6*np;
  int    dima2 =  2*np;
  int    dimb  =  5*np;

  short* a1   = new short[dima1];
  short* a2   = new short[dima2];
  float* b1   = new float[dimb ];

  R__b.ReadFastArray(a1,dima1);
  R__b.ReadFastArray(a2,dima2);
  R__b.ReadFastArray(b1,dimb );
//-----------------------------------------------------------------------------
//  read array of particles
//-----------------------------------------------------------------------------
  TObspParticle* p;
  for (int i=0; i<np; i++) {

    loca1 = 6*i;
    loca2 = 2*i;
				// assume that vertex #'s are 2-byte long

    pdg_code    = a1[loca1  ] | (a2[loca2  ] << 16);
    iv1         = a1[loca1+1];
    iv2         = a1[loca1+2];
    genp        = a1[loca1+3] | (a2[loca2+1] << 16);
    dc          = a1[loca1+4];
    charge      = a1[loca1+5]/3.;

    locb = 5*i;
	
    p = NewParticle(i,pdg_code,b1[locb],b1[locb+1],b1[locb+2],b1[locb+3],
		    b1[locb+4],iv1,iv2,genp,dc,charge);
  }

  delete [] a1;
  delete [] a2;
  delete [] b1;

  return 0;
}


//______________________________________________________________________________
int TObspBlock::WriteParticles(TBuffer &R__b) {
//-----------------------------------------------------------------------------
// write out list of OBSP particles: first integers, then floats
// assume that most integers have only 2 bytes meaningful, so split each 
// integer into 2 halves; write 4*8 = 32 bytes per vertex
//-----------------------------------------------------------------------------
  R__b << fNParticles;
  if (fNParticles <= 0) return 0;

  int    dima1 =  6*fNParticles;
  int    dima2 =  2*fNParticles;
  int    dimb  =  5*fNParticles;

  short* a1   = new short[dima1];
  short* a2   = new short[dima2];
  float* b1   = new float[dimb ];

  int   loca1, loca2, locb;
//-----------------------------------------------------------------------------
// fill arrays, then write them out
//-----------------------------------------------------------------------------
  TLorentzVector* mom;
  for (int i=0; i<fNParticles; i++) {
    TObspParticle* p = Particle(i);
    mom      = p->Momentum();

    loca1 = 6*i;
    loca2 = 2*i;
	
    a1[loca1  ] = (p->PdgCode()                ) & 0xffff;

				// assume that vertex #'s are 2-byte long

    a1[loca1+1] = (p->VertexNumber()           ) & 0xffff;
    a1[loca1+2] = (p->DecayVertexNumber()      ) & 0xffff;
    a1[loca1+3] = (p->GenpNumber()             ) & 0xffff;
    a1[loca1+4] = (p->DecayCode()              ) & 0xffff;
    a1[loca1+5] =  int(p->Charge()*3);


    a2[loca2  ] = (p->PdgCode()           >> 16) & 0xffff;
    a2[loca2+1] = (p->GenpNumber()        >> 16) & 0xffff;


    locb = 5*i;
	
    b1[locb   ]  = mom->Px();
    b1[locb+ 1]  = mom->Py();
    b1[locb+ 2]  = mom->Pz();
    b1[locb+ 3]  = mom->M();
    b1[locb+ 4]  = p->Path();
  }

  R__b.WriteFastArray(a1,dima1);
  R__b.WriteFastArray(a2,dima2);
  R__b.WriteFastArray(b1,dimb);

  delete [] a1;
  delete [] a2;
  delete [] b1;

  return 0;
}



//______________________________________________________________________________
void TObspBlock::Streamer(TBuffer &R__b)
{
  // Stream an object of class TObspBlock as fast as possible, dont' write
  // primary vertices in this version, will be added later 

  if (R__b.IsReading()) {
//-----------------------------------------------------------------------------
//  read section
//-----------------------------------------------------------------------------
    Version_t R__v = R__b.ReadVersion(); 
    if      (R__v == 1) ReadV1(R__b);
    else if (R__v == 2) {
      ReadParticles(R__b);
      if (fNParticles > 0) {
	ReadVertices(R__b);
	if (fNVertices > 0) {
	  ReadPrimaryVertices(R__b);
	}
      }
    }
  } 
  else {
//-----------------------------------------------------------------------------
//  write section
//-----------------------------------------------------------------------------
    R__b.WriteVersion(TObspBlock::IsA());
    WriteParticles(R__b);
    if (fNParticles > 0) {
      WriteVertices(R__b);
      if (fNVertices > 0) {
	WritePrimaryVertices(R__b);
      }
    }
  }

}

//_____________________________________________________________________________
TObspBlock::TObspBlock() {
  fNParticles            = 0;
  fNVertices             = 0;
  fNPrimaryVertices      = 0;
  fListOfParticles       = new TClonesArray("TObspParticle",100);
  fListOfParticles->BypassStreamer(kFALSE);
  fListOfVertices        = new TClonesArray("TObsvVertex",50);
  fListOfVertices->BypassStreamer(kFALSE);
  fListOfPrimaryVertices = new TClonesArray("TObsvVertex",10);
  fListOfPrimaryVertices->BypassStreamer(kFALSE);
}


//_____________________________________________________________________________
TObspBlock::~TObspBlock() {
  fListOfParticles->Delete();
  delete fListOfParticles;
  fListOfVertices->Delete();
  delete fListOfVertices;
  fListOfPrimaryVertices->Delete();
  delete fListOfPrimaryVertices;
}


//_____________________________________________________________________________
void TObspBlock::Clear(const char* opt) {
  fNParticles       = 0;
  fNVertices        = 0;
  fNPrimaryVertices = 0;
					// don't modify cut values at run time
  fListOfParticles->Clear();
  fListOfVertices->Clear();
  fListOfPrimaryVertices->Clear();
}

//_____________________________________________________________________________
TObspParticle* TObspBlock::NewParticle(Int_t   IPart, 
				       Int_t   PdgCode, 
				       Float_t Px, Float_t Py, Float_t Pz, 
				       Float_t M,
				       Float_t Path, 
				       Int_t   VertexNumber, 
				       Int_t   DecayVertexNumber, 
				       Int_t   GenpNumber,
				       Int_t   DecayCode,
				       Float_t Charge)
{
  // add new particle to the block. Block is filled sequentially, so this is 
  // the last particle and it is added to the last primary interaction

  TObspParticle* p;
  p = new ((*fListOfParticles)[fNParticles++]) 
    TObspParticle(IPart,PdgCode,Px,Py,Pz,M, Path, 
		  VertexNumber,DecayVertexNumber,GenpNumber, 
		  DecayCode, Charge);
  return p;
}


//_____________________________________________________________________________
TObsvVertex* TObspBlock::NewVertex(Int_t   IVert, Float_t Vx, Float_t Vy, 
				  Float_t Vz, Float_t Time,
				  Int_t   Mother, 
				  Int_t   FirstDaughter,
				  Int_t   LastDaughter)
{
  // add new particle to the block. Block is filled sequentially, so this is 
  // the last particle and it is added to the last primary interaction

  return new ((*fListOfVertices)[fNVertices++]) TObsvVertex(IVert,Vx,Vy,Vz,Time,
							   Mother,
							   FirstDaughter,
							   LastDaughter);
}


//_____________________________________________________________________________
void TObspBlock::Print(const char* opt) const {
  printf(" ********** TObspBlock::Print ***************** \n");
  int banner_printed;
  banner_printed = 0;
  for (int i=0; i<fNParticles; i++) {
    TObspParticle* p = Particle(i);
    if (! banner_printed) {
      p->Print("banner");
      banner_printed = 1;
    }
    p->Print("");
  }

  banner_printed = 0;
  for (int i=0; i<fNVertices; i++) {
    TObsvVertex* v = Vertex(i);
    if (! banner_printed) {
      v->Print("banner");
      banner_printed = 1;
    }
    v->Print("");
  }
}


//_____________________________________________________________________________
Int_t TObspBlock::GetParticles(Int_t       PdgCode, 
			       TObjArray*  List   , 
			       Float_t     PtMin  ,
			       Float_t     EtaMin ,
			       Float_t     EtaMax ) const
{
  // return 'List' of particles with given 'PdgCode' and their antiparticles

  float  pt, eta;

  List->Clear();

  for (int i=0; i<fNParticles; i++) {
    TObspParticle* p = Particle(i);
    if (abs(p->PdgCode()) == PdgCode) {
      pt = p->Momentum()->Pt();
      if (pt >= PtMin) {
	eta = p->Momentum()->Eta();
	if ((eta >= EtaMin) && (eta < EtaMax)) {
	  List->Add(p);
	}
      }
    }
  }
  return List->GetEntriesFast();
}


//_____________________________________________________________________________
Int_t TObspBlock::GetParticles(const char* Name   , 
			       TObjArray*  List   ,
			       Float_t     PtMin  ,
			       Float_t     EtaMin ,
			       Float_t     EtaMax ) const
{
  // return 'List' of particles with given 'Name' and their antiparticles

  List->Clear();

  TParticlePDG* part = TDatabasePDG::Instance()->GetParticle(Name);
  if (!part) return 0;

  Int_t pdg_code = part->PdgCode();
  return GetParticles(pdg_code,List,PtMin,EtaMin,EtaMax);
}


//_____________________________________________________________________________
void TObspBlock::PrintDaughters(TObspParticle* P) const {
  // remember: indices stored in OPSP are the FORTRAN ones

  TObspParticle  *d;
  TObsvVertex    *v1, *v2;
  int            iv1, iv2, i1, i2;

  iv1 = P->DecayVertexNumber()-1;
  if (iv1 < 0) return;

  v1 = Vertex(iv1);
    
  i1 = v1->FirstDaughter()-1;
  i2 = v1->LastDaughter ()-1;

  for (int i=i1; i<=i2; i++) {
    d = Particle(i);
    d->Print ("/data/nolf");
    v1->Print("/data/nolf");
    iv2 = d->DecayVertexNumber()-1;
    if (iv2 > 0) {
      v2 = Vertex(iv2);
      v2->Print("/data");
      PrintDaughters(d);
    }
    else {
      printf("\n");
    }
  }

}

//_____________________________________________________________________________
void TObspBlock::Print(Int_t   PdgCode, 
		       Float_t PtMin  , 
		       Float_t EtaMin , 
		       Float_t EtaMax ) const 
{
  TArrayI        vind(fNVertices);
  TObspParticle* p;

  int banner_printed, iv1,iv2;
  banner_printed = 0;
  for (int i=0; i<fNParticles; i++) {
    p = Particle(i);
    if ((abs(p->PdgCode()) == PdgCode) && 
	(p->Momentum()->Pt() > PtMin ) && 
	(p->Momentum()->Eta() >= EtaMin) && 
	(p->Momentum()->Eta() < EtaMax )    ) {
      
      iv1 = p->VertexNumber     ()-1;
      iv2 = p->DecayVertexNumber()-1;

      vind[iv1] = 1;

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

      p->Print("/data/nolf"); 
      Vertex(iv1)->Print("/data/nolf");
      Vertex(iv2)->Print("/data");
				// and print the daughters
      PrintDaughters(p);
    }
  }
}