// ----------------------------------------------------------------------------
// $Id: LCDJetFinder.cxx,v 1.2 2001/05/11 21:34:50 toshi Exp $
// ----------------------------------------------------------------------------
//
// $Log: LCDJetFinder.cxx,v $
// Revision 1.2  2001/05/11 21:34:50  toshi
// Minar changes. MS VC++ 6 picked up some multi-difinition error of
// Int_t i .
//
// Revision 1.1  2001/05/10 19:32:06  toshi
// C++ file name convention has been changed from .C to .cxx .
//
//
// Jet finder base class routines.
//
// V0.0 Mar 01 99 : R. Shanks  Derived from Java routines written by G.Bower. 
// V0.1 Jul 02 99 : M.Iwasaki  Make some Modification on masscut and ymass 
//                             calculation (in JadeEJetfonder.cxx 
//                             or DurhamJetFinder.cxx) on doFindJets().
//                             Adding Jade cluster (yclus).
// V0.2 Jul 07 99 : M.Iwasaki  Change Mod to Mag function in TVector3
//                             so as to use root 2.22.x
// V0.3 Jul 21 99 : M.Iwasaki  Make temporary 4-vectors copied from input 
//                             particle 4-vectors in doFindJets() 
//                             to protect the initial 4-vectors.
// V0.4 Sep 23 99 : M.Iwasaki  Fix setEvent memory leak.
//
// V1.0 May 15 00 : M.Iwasaki  Make necessary modifications, and change classes
//                             Merge JadeEJetFinder, JadeJetFinder, 
//                             DurhamJetfinder, and JetFinder into one JetFinder.
// V1.1 Aug 15 00 : M.Iwasaki  Change calcinvmass. Remove DurhamJetFinder, 
//                             JadeJetFinder, and JadeEJetFinder classes.
//                             Add SetDURHAM, SetJADE, or SetJADEE. 
//                             Now we can change algorithms by one of them.
//                             Change Class name from JetFinder to LCDJetFinder.

#include "LCDJetFinder.h"

//_____________________________________________________________________
//  ------------------
//   JetFinder Class
//  ------------------
//
ClassImp(LCDJetFinder)

const Int_t LCDJetFinder::UNASSOC = -999;

const Int_t LCDJetFinder::DURHAM = 1;
const Int_t LCDJetFinder::JADE   = 2;
const Int_t LCDJetFinder::JADEE  = 3;

 LCDJetFinder::LCDJetFinder(): 
    m_ymin(1e20),m_evis(0.),m_dycut(0.1),
    m_algorithm(LCDJetFinder::DURHAM){ // Default constructor  uses Durham
    m_jet  = new TClonesArray("TLorentzVector", 20);
    m_part = new TClonesArray("TLorentzVector",200);
}

 LCDJetFinder::LCDJetFinder(Double_t ycut): 
    m_ymin(0.0),m_evis(0.),m_dycut(ycut),
    m_algorithm(LCDJetFinder::DURHAM){ // Default constructor  uses Durham
    m_jet  = new TClonesArray("TLorentzVector", 20);
    m_part = new TClonesArray("TLorentzVector",200);
}
//______________________________________________________

 LCDJetFinder::~LCDJetFinder() {
  m_jet->Delete();  delete m_jet;
  m_part->Delete(); delete m_part;
}
//______________________________________________________

 TLorentzVector* LCDJetFinder::jet4vec(Int_t index) {
  return (TLorentzVector*)m_jet->At(index);
}
//_______________________________________________________

 Int_t LCDJetFinder::nParticlesPerJet(Int_t index) {
  return m_inparts_per_jet[index];
}
//_______________________________________________________

 void LCDJetFinder::SetYCut(Double_t ycut) {
  m_dycut = ycut;
}
//_______________________________________________________

// Input the particle 4(3)-vector list
// e: 4-vector  TLorentzVector ..(px,py,pz,E) or
//    3-vector  TVector3       ..(px,py,pz) 
// If you input TVector3, the energy of particle
// will be E = sqrt(px**2 + py**2 + pz**2)
 void LCDJetFinder::SetPartList(TObjArray* e) {

  m_evis = 0;
  m_4vec = e;

  Int_t ne = e->GetEntries();
  TObject* o;
  for(Int_t i=0;i<ne;i++) {
    o = m_4vec->At(i);
    TString nam(o->IsA()->GetName());
    if (nam.Contains("TLorentzVector")) {
      m_evis += ((TLorentzVector *) o)->T();
    } else if (nam.Contains("TVector3")) {
      m_evis += ((TVector3 *) o)->Mag();
    } else {
      printf("LCDJetFinder::SetEvent input is not a TVector3 or a TLorentzVectorn");
    }
  }

}
//_______________________________________________________

 void LCDJetFinder::doFindJets(Double_t ycut){
  SetYCut(ycut);
  doFindJets();
}
//______________________________________________________

 void LCDJetFinder::doFindJets(){

  Int_t np = m_4vec->GetEntries(); 	
  if (np < 2) return;

  TObject* o;
  Int_t ipart;
  for(ipart=0; ipart < np ; ipart++) {
    o = m_4vec->At(ipart);
    TString nam(o->IsA()->GetName());
    if (nam.Contains("TLorentzVector")) {
      new((*m_part)[ipart]) TLorentzVector(*((TLorentzVector*)o));
    } else if (nam.Contains("TVector3")) {
      new((*m_part)[ipart]) TLorentzVector(((TVector3 *) o)->X(),
					   ((TVector3 *) o)->Y(),
					   ((TVector3 *) o)->Z(),
					   ((TVector3 *) o)->Mag());
    } else {
      printf("LCDJetFinder::SetEvent input is not a TVector3 or a TLorentzVectorn");
    }
  }

  m_ipart_jet_assoc.Reset();
  m_ipart_jet_assoc.Set(np);
  for (Int_t m=0; m<np; m++) m_ipart_jet_assoc[m] = UNASSOC;

  m_njets = 0;

  //
  // create invariant mass pair array.
  //
  TMatrix ymass = TMatrix(np,np);
  
  for (Int_t i1 = 0; i1 < np - 1; i1++ ) {
    for (Int_t i2 = i1 + 1 ; i2 < np ; i2++ ) {
      ymass(i1,i2) = calcinvmass(*(TLorentzVector*)m_part->At(i1),
				 *(TLorentzVector*)m_part->At(i2));
    }
  }
  
  Double_t masscut = m_dycut * m_evis * m_evis ;

  Int_t im,jm;
  Double_t minmass;
  Int_t i,j;
  Int_t imin;
  Int_t imax;
  for (;;) {
    im = -1;
    jm = -1;
    minmass = 100000000.;
    if ( minmass <= masscut ) minmass = masscut * 10.;
    //
    // find least invariant mass pair.
    //
    for(i = 0 ; i < np - 1 ; i++ ) {
      for(j = i + 1 ; j < np ; j++ ) {
	if (m_ipart_jet_assoc[i] != LCDJetFinder::UNASSOC) continue;
	if (m_ipart_jet_assoc[j] != LCDJetFinder::UNASSOC) continue;
	if (ymass(i,j) > minmass) continue;
	  
	minmass = ymass(i,j);
	im = i;  jm = j;
      }
    }

    m_ymin=minmass/m_evis/m_evis;

    if (minmass > masscut) break;
    //
    // combine particles im and jm.
    //
    *((TLorentzVector*)m_part->At(im)) += *((TLorentzVector*)m_part->At(jm));

    for(ipart = 0; ipart < np; ipart++ ){
      if(m_ipart_jet_assoc[ipart] == jm) m_ipart_jet_assoc[ipart] = im;
    }
    //
    // Recalculate invariant masses for newly combined particle
    //
    m_ipart_jet_assoc[jm] = im;
    for (ipart = 0; ipart < np ; ipart++) {
      if (ipart == im) continue;
      if (m_ipart_jet_assoc[ipart] != UNASSOC ) continue;
      
      imin = TMath::Min(ipart,im);
      imax = TMath::Max(ipart,im);
      
      ymass(imin,imax) = calcinvmass(*((TLorentzVector*)m_part->At(imin)),
				     *((TLorentzVector*)m_part->At(imax)));
    }
    
  }
  //
  // finish up by filling jet array.
  //
  
  for (ipart = 0 ; ipart < np ; ipart++) {
    if (m_ipart_jet_assoc[ipart] == UNASSOC) m_njets++;			
  }
  
  m_jet->Delete();
  m_inparts_per_jet.Reset();
  m_inparts_per_jet.Set(m_njets);

  Int_t nj = 0;
  Int_t npart;
  m_ifewest_parts = 100000; // Starting min value
  for(i = 0 ; i < np ; i++ ){
    if (m_ipart_jet_assoc[i] != UNASSOC) continue;

    new((*m_jet)[nj]) TLorentzVector(*((TLorentzVector*)m_part->At(i)));

    npart = 1;
    for (j = 0 ; j < np ; j++) {
      if(m_ipart_jet_assoc[j] == i) {
	m_ipart_jet_assoc[j] = nj;
	npart++;
      }
    }
    m_ipart_jet_assoc[i] = nj;
    m_inparts_per_jet[nj] = npart;
    if( npart < m_ifewest_parts) m_ifewest_parts = npart;
    nj++;
  }  
  m_part->Delete();

}
//______________________________________________________

 Int_t LCDJetFinder::doFindJets(Int_t njet){

  Int_t np = m_4vec->GetEntries(); 	
  if (np < njet) return -1;

  TObject* o;
  Int_t ipart;
  for(ipart=0; ipart < np ; ipart++) {
    o = m_4vec->At(ipart);
    TString nam(o->IsA()->GetName());
    if (nam.Contains("TLorentzVector")) {
      new((*m_part)[ipart]) TLorentzVector(*((TLorentzVector*)o));
    } else if (nam.Contains("TVector3")) {
      new((*m_part)[ipart]) TLorentzVector(((TVector3 *) o)->X(),
					   ((TVector3 *) o)->Y(),
					   ((TVector3 *) o)->Z(),
					   ((TVector3 *) o)->Mag());
    } else {
      printf("LCDJetFinder::SetEvent input is not a TVector3 or a TLorentzVectorn");
    }
  }

  m_ipart_jet_assoc.Reset();
  m_ipart_jet_assoc.Set(np);
  for (Int_t m=0; m<np; m++) m_ipart_jet_assoc[m] = UNASSOC;

  m_njets = 0;

  //
  // create invariant mass pair array.
  //
  TMatrix ymass = TMatrix(np,np);
  
  for (Int_t i1 = 0; i1 < np - 1; i1++ ) {
    for (Int_t i2 = i1 + 1 ; i2 < np ; i2++ ) {
      ymass(i1,i2) = calcinvmass(*(TLorentzVector*)m_part->At(i1),
				 *(TLorentzVector*)m_part->At(i2));
    }
  }
  
  Int_t im,jm;
  Int_t i,j;
  Int_t imin;
  Int_t imax;
  Int_t nj=np;
  for (;;) {
    im = -1;
    jm = -1;
    m_ymin=1E20;
    //
    // find least invariant mass pair.
    //
    for(i = 0 ; i < np - 1 ; i++ ) {
      for(j = i + 1 ; j < np ; j++ ) {
	if (m_ipart_jet_assoc[i] != LCDJetFinder::UNASSOC) continue;
	if (m_ipart_jet_assoc[j] != LCDJetFinder::UNASSOC) continue;
	if (ymass(i,j) > m_ymin) continue;
	  
	m_ymin = ymass(i,j);
	im = i;  jm = j;
      }
    }

    //
    // combine particles im and jm.
    //
    *((TLorentzVector*)m_part->At(im)) += *((TLorentzVector*)m_part->At(jm));

    for(ipart = 0; ipart < np; ipart++ ){
      if(m_ipart_jet_assoc[ipart] == jm) m_ipart_jet_assoc[ipart] = im;
    }
    m_ipart_jet_assoc[jm] = im;
    m_ymin/= m_evis;
    m_ymin/= m_evis;
    if (--nj == njet) break;

    //
    // Recalculate invariant masses for newly combined particle
    //
    for (ipart = 0; ipart < np ; ipart++) {
      if (ipart == im) continue;
      if (m_ipart_jet_assoc[ipart] != UNASSOC ) continue;
      
      imin = TMath::Min(ipart,im);
      imax = TMath::Max(ipart,im);
      
      ymass(imin,imax) = calcinvmass(*((TLorentzVector*)m_part->At(imin)),
				     *((TLorentzVector*)m_part->At(imax)));
    }
    
  }
  //
  // finish up by filling jet array.
  //
  
  for (ipart = 0 ; ipart < np ; ipart++) {
    if (m_ipart_jet_assoc[ipart] == UNASSOC) m_njets++;			
  }

  //For debug...
  if (m_njets != njet) {
    fprintf(stderr,"LCDJetFinder::doFindJets error m_njets=%d njet=%dn",
	    m_njets,njet);
    return -1;
  }
  
  m_jet->Delete();
  m_inparts_per_jet.Reset();
  m_inparts_per_jet.Set(m_njets);

  nj = 0;
  Int_t npart;
  m_ifewest_parts = 100000; // Starting min value
  for(i = 0 ; i < np ; i++ ){
    if (m_ipart_jet_assoc[i] != UNASSOC) continue;

    new((*m_jet)[nj]) TLorentzVector(*((TLorentzVector*)m_part->At(i)));

    npart = 1;
    for (j = 0 ; j < np ; j++) {
      if(m_ipart_jet_assoc[j] == i) {
	m_ipart_jet_assoc[j] = nj;
	npart++;
      }
    }
    m_ipart_jet_assoc[i] = nj;
    m_inparts_per_jet[nj] = npart;
    if( npart < m_ifewest_parts) m_ifewest_parts = npart;
    nj++;
  }
  m_part->Delete();

  return 0;
}
//______________________________________________________

 void LCDJetFinder::SetDURHAM(){
  m_algorithm = LCDJetFinder::DURHAM; 
}
 void LCDJetFinder::SetJADE(){
  m_algorithm = LCDJetFinder::JADE;
}
 void LCDJetFinder::SetJADEE(){
  m_algorithm = LCDJetFinder::JADEE;
}

 Double_t LCDJetFinder::calcinvmass(const TLorentzVector& jet1,
				   const TLorentzVector& jet2){
  TVector3 P_jet1 = jet1.Vect();
  TVector3 P_jet2 = jet2.Vect();
  Double_t costh = (P_jet1 * P_jet2)/P_jet1.Mag()/P_jet2.Mag();

  if     (m_algorithm == LCDJetFinder::DURHAM) {  // DURHAM
    Double_t minT = TMath::Min(jet1.E(),jet2.E());
    return 2. * minT*minT * (1.-costh);
  } else if (m_algorithm == LCDJetFinder::JADE)   {  // JADE    
    return 2. * (jet1.E())*(jet2.E()) * (1.-costh) ; 
  } else if (m_algorithm == LCDJetFinder::JADEE)  {  // JADE E
    return (jet1 + jet2).M2();
  }

  printf(" Strange Algorithm!!!! n");
  return 0.;

};