///////////////////////////////////////////////////////////
//
// File:     conejet.cpp
//
// Purpose:  make simple cone jets, similar to pxcone jet algorithm
//
// Created:  Shaohua Fu
//
// History:
//
////////////////////////////////////////////////////////////

#include "mctest/conejet.hpp"
#include "kinem_util/AnglesUtil.hpp"

#include <iostream>
#include <iomanip>

using namespace std;

const float EPSILON = 0.001;
const int MAXITER = 30;


// Constructor
ConeJet::ConeJet(list<vector<double> >& intracks, float coneR, float jetPTmin, float overlaplim, bool debug):
  _intracks(intracks), _coneR(coneR), _jetPTmin(jetPTmin), _overlaplim(overlaplim), _debug(debug) {

  _ptracks.clear();
  _pjets.clear();

  //////
  _debug = debug;

  MakeConeJets();
}

void ConeJet::MakeConeJets(){
  if (_debug) {
    cout<<"-------------------- Make Cone Jets --------------------"<<endl;
    cout<<"     Cone Size R = "<<_coneR<<" Radians"<<endl;
    cout<<"     Min Jet PT = "<<_jetPTmin<<" GeV"<<endl;
    cout<<"     Overlap fraction = "<<_overlaplim<<endl;
    cout<<"--------------------------------------------------------"<<endl;
  }

  if (_intracks.size()==0) {
    cout<<"ConeJet WARNING: NO input tracks... EXIT"<<endl;
    return;
  }

  if (_debug) {
    cout<<" INPUT TRACK LISTS: "<<_intracks.size()<<endl;
    for (list<vector<double> >::iterator it = _intracks.begin(); it != _intracks.end(); ++it) {
      for (vector<double>::iterator itt = (*it).begin(); itt != (*it).end(); ++itt) cout<<setprecision(4)<<setw(8)<<(*itt)<<"  ";
      cout<<endl;
    }
  }

  InitSeeds();

  if (_debug) {
    cout<<" PARTICLE TRACK LISTS: "<<endl;
    PrintParts(_ptracks);
  }

  // Look for jets using particle directions as seed axes
  vector<Particle> jets_pre;
  for (vector<Particle>::iterator itrk = _ptracks.begin(); itrk != _ptracks.end(); ++itrk) {
    vector<double> axis;
    axis.push_back((*itrk).eta);
    axis.push_back((*itrk).phi);
    Particle tjet;
    if (JetFromSeed(axis, tjet)) {
      if (IsNewJet(tjet, jets_pre)) {
	// if the jet is a new one, add it to our array
	jets_pre.push_back(tjet);
      }
    }
  }

  if (_debug) {
    cout<<" FIRST LIST OF JETS: "<<endl;
    PrintParts(jets_pre);
  }

  // Now look between all pairs of jets as seed axes
  vector<Particle> jets_post = jets_pre;
  for (int i=0; i<jets_pre.size()-1; ++i) {
    for (int j=i+1; j<jets_pre.size(); ++j) {
      // only bother if their distance is between 1 and 2 cone radii apart
      double eta1 = jets_pre[i].eta;
      double phi1 = jets_pre[i].phi;
      double eta2 = jets_pre[j].eta;
      double phi2 = jets_pre[j].phi;
      double dist = kinem::delta_R(eta1, phi1, eta2, phi2);
      if (dist>_coneR && dist<_coneR*2) {
	vector<double> axis;
	axis.push_back((eta1+eta2)*0.5);
	axis.push_back((phi1+phi2)*0.5);
	Particle tjet;
	if (JetFromSeed(axis, tjet)) {
	  if (IsNewJet(tjet, jets_post)) {
	    jets_post.push_back(tjet);
	  }
	}
      }
    } // j
  } // i

  if (_debug) {
    cout<<" SECOND LIST OF JETS: "<<endl;
    PrintParts(jets_post);
  }

  // now put the jet list into order by PT, eliminating jets below threshold
  SortAndClean(jets_post, _jetPTmin);

  if (_debug) {
    cout<<" JETS AFTER SORT: "<<endl;
    PrintParts(jets_post);
  }

  // take care of jet overlaps
  OverlappedJets(jets_post);

  if (_debug) {
    cout<<" JETS AFTER REMOVING OVERLAP: "<<endl;
    PrintParts(jets_post);
  }

  // order jets again as some have been eliminated, or lost energy
  SortAndClean(jets_post, _jetPTmin);

  if (_debug) {
    cout<<" FINAL JETS: "<<endl;
    PrintParts(jets_post);
  }

  // All done! copy output into output arrays
  _pjets = jets_post;

  return;
}


void ConeJet::GetConeJets(list<vector<double> >& outjets){
  _outjets.clear();
  for (vector<Particle>::iterator it = _pjets.begin(); it != _pjets.end(); ++it) {
    vector<double> jet;
    jet.push_back((*it).Px);
    jet.push_back((*it).Py);
    jet.push_back((*it).Pz);
    jet.push_back((*it).E);
    jet.push_back((*it).eta);
    jet.push_back((*it).phi);
    jet.push_back((*it).PT);
    _outjets.push_back(jet);
  }
  outjets = _outjets;
  return;
}


// (Px,Py,Pz,E) -> (Px,Py,Pz,E,eta,phi,PT)
void ConeJet::InitSeeds(){
  list<vector<double> >::iterator itrk;
  for (itrk = _intracks.begin(); itrk != _intracks.end(); ++itrk) {
    if ((*itrk).size() != 4) {
      cout<<"WARNING: this input track 4-momenta wrong -- removed"<<endl;
      continue;
    }
    Particle track;
    track.Px = (*itrk)[0];
    track.Py = (*itrk)[1];
    track.Pz = (*itrk)[2];
    track.E = (*itrk)[3];
    track.eta = kinem::eta(track.Px, track.Py, track.Pz);
    track.phi = kinem::phi(track.Px, track.Py);
    track.PT = sqrt(track.Px*track.Px + track.Py*track.Py);
    _ptracks.push_back(track);
  }
  SortAndClean(_ptracks, EPSILON);
  return;
}


// sort the list in descending PT, remove PT < threshold
void ConeJet::SortAndClean(vector<Particle>& partlist, double threshold){
  if (partlist.size() == 0) return;
  vector<Particle> oldlist = partlist;
  vector<Particle> newlist;
  for (vector<Particle>::iterator it = oldlist.begin(); it != oldlist.end(); ++it) {
    if ((*it).PT < threshold) continue;
    newlist.push_back(*it);
  }
  sort(newlist.begin(), newlist.end(), compareParticlePT());
  partlist = newlist;

  // added check for duplicates
  if (partlist.size() == 0) return;
  oldlist = newlist;
  newlist.clear();
  newlist.push_back(oldlist[0]);
  for (int i=1; i<oldlist.size(); ++i) {
    Particle previous = newlist[newlist.size()-1];
    if (oldlist[i].Px==previous.Px && oldlist[i].Py==previous.Py && oldlist[i].Pz==previous.Pz && oldlist[i].E==previous.E)
      {cout<<"Duplicate: same Px,Py,Pz,E!"<<endl; continue;}
    if (oldlist[i].eta==previous.eta && oldlist[i].phi==previous.phi && oldlist[i].PT==previous.PT)
      {cout<<"Duplicate: same eta,phi,PT!"<<endl; continue;}
    if (oldlist[i].itrack.size()>0 && oldlist[i].itrack==previous.itrack)
      {cout<<"Duplicate: same non-zero itrack!"<<endl; continue;}
    newlist.push_back(oldlist[i]);
  }
  partlist = newlist;
  return;
}


// make a cone jet from a seed direction
bool ConeJet::JetFromSeed(vector<double>& axis, Particle& pjet){
  // use seed as a trial axis, look for a stable jet
  // check stable jet against those already found and add to list
  // will try up to MAXITER iterations to get a stable set of particles in the cone
  for (int iter = 0; iter < MAXITER; ++iter) {
    Particle tmpjet;
    TryJet(axis, tmpjet);
    // return immediately if there were no particles in the cone
    if (tmpjet.itrack.size()==0) return false;
    if (iter==0) {pjet = tmpjet; continue;}
    if (tmpjet.itrack == pjet.itrack) {
      // we have a stable jet
      return true;
    }
    pjet = tmpjet;
  }
  cout<<"WARNING: this jet is still unstable after "<<MAXITER<<" iterations."<<endl;
  return true;
}


void ConeJet::TryJet(vector<double>& axis, Particle& tmpjet){
  for (int i = 0; i < _ptracks.size(); ++i) {
    double dR = kinem::delta_R(axis[0], axis[1], _ptracks[i].eta, _ptracks[i].phi);
    // check if particle is inside cone
    if (dR < _coneR) {
      tmpjet.Px += _ptracks[i].Px;
      tmpjet.Py += _ptracks[i].Py;
      tmpjet.Pz += _ptracks[i].Pz;
      tmpjet.E += _ptracks[i].E;
      tmpjet.itrack.push_back(i);
    }
  }
  // if there are particles in the cone, calculate new jet axis
  if (tmpjet.itrack.size() > 0) {
    tmpjet.eta = kinem::eta(tmpjet.Px, tmpjet.Py, tmpjet.Pz);
    tmpjet.phi = kinem::phi(tmpjet.Px, tmpjet.Py);
    tmpjet.PT = sqrt(tmpjet.Px*tmpjet.Px + tmpjet.Py*tmpjet.Py);
  }
  return;
}


bool ConeJet::IsNewJet(Particle& tmpjet, vector<Particle>& pjets){
  bool isnew = true;
  for (vector<Particle>::iterator it = pjets.begin(); it != pjets.end() && isnew; ++it) {
    if (tmpjet.itrack == (*it).itrack) isnew = false;
  }
  return isnew;
}


// Look for particles assigned to more than 1 jet, and reassigns them
// If more than a fraction _overlaplim of a jet's PT is contained in
// higher energy jets, that jet is neglected
// Particles assigned to the jet closest in angle
void ConeJet::OverlappedJets(vector<Particle>& pjets){
  if (pjets.size()<1) return;

  // look for jets with large overlaps with higher energy jets
  // from now on only reply on jet.itrack info
  for (int i=1; i<pjets.size(); ++i) {
    if (pjets[i].itrack.size()==0) continue;
    double ptover = 0;
    for (int j=0; j<i; ++j) {
      if (pjets[j].itrack.size()==0) continue;
      vector<int> itrkover;
      for (vector<int>::iterator it=pjets[i].itrack.begin(); it!=pjets[i].itrack.end(); ++it)
	for (vector<int>::iterator jt=pjets[j].itrack.begin(); jt!=pjets[j].itrack.end(); ++jt)
	  if ((*it)==(*jt)) itrkover.push_back(*it);
      for (vector<int>::iterator it=itrkover.begin(); it!=itrkover.end(); ++it)
	ptover += _ptracks[(*it)].PT;
    }
    // is the fraction of energy shared larger than _overlaplim?
    if (ptover > _overlaplim*pjets[i].PT) {
      // de-assign all particles from jet "i" !
      pjets[i].itrack.clear();
    }
  }

  for (vector<Particle>::iterator it=pjets.begin(); it!=pjets.end(); ++it)
    RefreshJet(*it);

  _isotrackid.clear();
  // any particles in more than 1 jet are assigned to the CLOSEST jet in angle
  for (int i=0; i<_ptracks.size(); ++i) {
    double dRmin = 100;
    int jmin = -1;
    vector<int> jids;
    for (int j=0; j<pjets.size(); ++j) {
      if (pjets[j].itrack.size()==0) continue;
      if (find(pjets[j].itrack.begin(), pjets[j].itrack.end(), i) != pjets[j].itrack.end()) jids.push_back(j);
      double dR = kinem::delta_R(_ptracks[i].eta, _ptracks[i].phi, pjets[j].eta, pjets[j].phi);
      if (jmin == -1) {jmin=j; dRmin=dR; continue;}
      if (dR < dRmin) {jmin=j; dRmin=dR;}
    }
    if (jids.size() == 0) {
      //cout<<"WARNING: an isolated track not in any jet"<<endl;
      _isotrackid[i] = dRmin;
    }
    if (jids.size() <= 1) continue;
    for (vector<int>::iterator jt=jids.begin(); jt!=jids.end(); ++jt) {
      if ((*jt) != jmin) {
	pjets[*jt].itrack.erase(find(pjets[*jt].itrack.begin(), pjets[*jt].itrack.end(), i));
      }
    }
  }

  for (vector<Particle>::iterator it=pjets.begin(); it!=pjets.end(); ++it)
    RefreshJet(*it);

  if (_debug) {
    cout<<" ==== ISOLATED TRACKS NOT IN ANY JETS : "<<_isotrackid.size()<<" ===="<<endl;
    for (map<int,double>::iterator iso=_isotrackid.begin(); iso!=_isotrackid.end(); ++iso)
      cout<<iso->first<<" : "<<iso->second<<"  "<<_ptracks[iso->first].PT<<endl;
  }

  return;
}


// refresh a jet variables from itrack
void ConeJet::RefreshJet(Particle& jet){
  jet.Px = 0;
  jet.Py = 0;
  jet.Pz = 0;
  jet.E = 0;
  jet.eta = -10;
  jet.phi = -1;
  jet.PT = 0;
  if (jet.itrack.size()==0) return;
  for (vector<int>::iterator it=jet.itrack.begin(); it!=jet.itrack.end(); ++it) {
    jet.Px += _ptracks[*it].Px;
    jet.Py += _ptracks[*it].Py;
    jet.Pz += _ptracks[*it].Pz;
    jet.E += _ptracks[*it].E;
  }
  jet.eta = kinem::eta(jet.Px, jet.Py, jet.Pz);
  jet.phi = kinem::phi(jet.Px, jet.Py);
  jet.PT = sqrt(jet.Px*jet.Px + jet.Py*jet.Py);

  return;
}


void ConeJet::PrintParts(vector<Particle>& parts){
  cout<<"**** size = "<<parts.size()<<" **** Px, Py, Pz, E, eta, phi, PT"<<endl;
  int ii = 0;
  for (vector<Particle>::iterator it = parts.begin(); it != parts.end(); ++it) {
    cout<<"("<<setw(3)<<ii<<") "<<setprecision(4)<<setw(8)<<(*it).Px<<"  "<<setprecision(4)<<setw(8)<<(*it).Py<<"  "<<setprecision(4)<<setw(8)<<(*it).Pz<<"  "<<"  "<<setprecision(4)<<setw(8)<<(*it).E<<"  "<<setprecision(4)<<setw(8)<<(*it).eta<<"  "<<setprecision(4)<<setw(8)<<(*it).phi<<"  "<<setprecision(4)<<setw(8)<<(*it).PT<<endl;
    ++ii;
    if ((*it).itrack.size() > 0) {
      cout<<"        itrack "<<(*it).itrack.size()<<" : ";
      for (vector<int>::iterator j = (*it).itrack.begin(); j != (*it).itrack.end(); ++j) cout<<(*j)<<"  ";
      cout<<endl;
    }
  }
}