// ******************************************************
// DecayTreeFitter Package
// We thank the original author Wouter Hulsbergen 
// (BaBar, LHCb) for providing the sources. 
// http://arxiv.org/abs/physics/0503191v1 (2005)
// Adaptation & Development for PANDA: Ralf Kliemt (2015)
// ******************************************************
#include "FitParams.h"
#include "InteractionPoint.h"
//#include "Event/Particle.h"
//#include "Event/VertexBase.h"
#include <assert.h>
#include "TVector3.h"
#include "TMath.h"
#include "RhoVector3Err.h"

using namespace DecayTreeFitter;

extern int vtxverbose ;
ClassImp(InteractionPoint);

DecayTreeFitter::InteractionPoint::InteractionPoint(const RhoVector3Err& ipvertex,
                                   RhoCandidate* cand,
                                   const Configuration& config)
: ParticleBase("IP"), m_ipPosCov(3), m_ipPosCovInv(3), m_ipMomCov(4), m_ipMomCovInv(4), m_haspos(true), m_hasmom(false)
{
  if(vtxverbose>=2) std::cout << "InteractionPoint: start"<<std::endl;
  addDaughter( cand, config ) ;
  if(vtxverbose>=2) std::cout << "InteractionPoint: daughters done"<<std::endl;
  
  m_ipPos(0) = ipvertex.x() ;
  m_ipPos(1) = ipvertex.y() ;
  m_ipPos(2) = ipvertex.z() ;
  for(int ii=0;ii<3;ii++) for(int jj=0;jj<3;jj++){
    m_ipPosCov(ii,jj) = ipvertex.CovMatrix()(ii,jj);
  }
  m_hasmomcov = false;
  m_hasposcov = (0!=m_ipPosCov.NonZeros()); // a zero protection flag
  if(m_hasposcov)
  {
    double ierr;
    if(vtxverbose>=2) std::cout << "InteractionPoint: inverting matrix"<<std::endl;
    m_ipPosCovInv = m_ipPosCov.Invert(&ierr) ;
  }
  if(vtxverbose>=2) {
    std::cout << "InteractionPoint: initial beam spot = ("
    <<m_ipPos(0) << "," << m_ipPos(1) << "," << m_ipPos(2) << ")       " <<std::endl ;
    std::cout << "daughter: " << daughters().front()->name() << std::endl ;
  }
}

DecayTreeFitter::InteractionPoint::InteractionPoint(const RhoLorentzVectorErr& ipmom, 
                                   RhoCandidate* cand,
                                   const Configuration& config)
: ParticleBase("IP"), m_ipPosCov(3), m_ipPosCovInv(3), m_ipMomCov(4), m_ipMomCovInv(4), m_haspos(false), m_hasmom(true)
{
  if(vtxverbose>=2) std::cout << "InteractionPoint: start"<<std::endl;
  addDaughter( cand, config ) ;
  if(vtxverbose>=2) std::cout << "InteractionPoint: daughters done"<<std::endl;
  
  m_ipMom(0) = ipmom.X();
  m_ipMom(1) = ipmom.Y();
  m_ipMom(2) = ipmom.Z();
  m_ipMom(3) = ipmom.E();
  for(int ii=0;ii<4;ii++) for(int jj=0;jj<4;jj++){
    m_ipMomCov(ii,jj) = ipmom.CovMatrix()(ii,jj);
  }
  
  m_hasposcov = false;
  m_hasmomcov = (0!=m_ipMomCov.NonZeros()); // a zero protection flag
  if(m_hasmomcov)
  {
    double ierr;
    if(vtxverbose>=2) std::cout << "InteractionPoint: inverting matrix"<<std::endl;
    m_ipMomCovInv = m_ipMomCov.Invert(&ierr) ;
  }
  if(vtxverbose>=2) {
    std::cout << "InteractionPoint: initial mom = (" <<m_ipMom(0) << "," << m_ipMom(1) << "," << m_ipMom(2)<< "," << m_ipMom(3) << ")" <<std::endl ;
    std::cout << "daughter: " << daughters().front()->name() << std::endl ;
  }
}

DecayTreeFitter::InteractionPoint::InteractionPoint(const RhoLorentzVectorErr& ipmom, 
                                   const RhoVector3Err& ipvertex,
                                   RhoCandidate* cand,
                                   const Configuration& config)
: ParticleBase("IP"), m_ipPosCov(3), m_ipPosCovInv(3), m_ipMomCov(4), m_ipMomCovInv(4), m_haspos(true), m_hasmom(true)
{
  if(vtxverbose>=2) std::cout << "InteractionPoint: start"<<std::endl;
  addDaughter( cand, config ) ;
  if(vtxverbose>=2) std::cout << "InteractionPoint: daughters done"<<std::endl;
  
  m_ipPos(0) = ipvertex.x() ;
  m_ipPos(1) = ipvertex.y() ;
  m_ipPos(2) = ipvertex.z() ;
  for(int ii=0;ii<3;ii++) for(int jj=0;jj<3;jj++){
    m_ipPosCov(ii,jj) = ipvertex.CovMatrix()(ii,jj);
  }
  
  m_ipMom(0) = ipmom.X();
  m_ipMom(1) = ipmom.Y();
  m_ipMom(2) = ipmom.Z();
  m_ipMom(3) = ipmom.E();
  for(int ii=0;ii<4;ii++) for(int jj=0;jj<4;jj++){
    m_ipMomCov(ii,jj) = ipmom.CovMatrix()(ii,jj);
  }
  
  m_hasposcov = (0!=m_ipPosCov.NonZeros()); // a zero protection flag
  if(m_hasposcov)
  {
    double ierr;
    if(vtxverbose>=2) std::cout << "InteractionPoint: inverting matrix"<<std::endl;
    m_ipPosCovInv = m_ipPosCov.Invert(&ierr) ;
  }
  m_hasmomcov = (0!=m_ipMomCov.NonZeros()); // a zero protection flag
  if(m_hasmomcov)
  {
    double ierr;
    if(vtxverbose>=2) std::cout << "InteractionPoint: inverting matrix"<<std::endl;
    m_ipMomCovInv = m_ipMomCov.Invert(&ierr) ;
  }
  if(vtxverbose>=2) {
    std::cout << "InteractionPoint: initial beam spot = ("
    <<m_ipPos(0) << "," << m_ipPos(1) << "," << m_ipPos(2) << ")       " 
    << "mom = (" <<m_ipMom(0) << "," << m_ipMom(1) << "," << m_ipMom(2)<< "," << m_ipMom(3) << ")" <<std::endl ;
    std::cout << "daughter: " << daughters().front()->name() << std::endl ;
  }
}

ErrCode
DecayTreeFitter::InteractionPoint::initPar1(FitParams* fitparams)
{
  ErrCode status ;
  if(vtxverbose>5){std::cout<<"InteractionPoint::initPar1: - start"<<std::endl;}

  // initialize our IP measurement
  if(m_haspos)
  {
    int posindex = posIndex() ;
    for(int row=0; row<3; row++)
      fitparams->par()(posindex+row) = m_ipPos(row) ;
  }
  
  // Initialize Daugters and the whole tree
  for(daucontainer::iterator idau = daughters().begin() ;
      idau != daughters().end(); ++idau ) {
  if(vtxverbose>5){std::cout<<"InteractionPoint::initPar1: - calling daughter initPar1"<<std::endl;}
    status |= (*idau)->initPar1(fitparams) ;
  if(vtxverbose>5){std::cout<<"InteractionPoint::initPar1: - calling daughter initPar2"<<std::endl;}
    status |= (*idau)->initPar2(fitparams) ;
  }

  return status ;
}

ErrCode
DecayTreeFitter::InteractionPoint::initPar2(FitParams* /*fitparams*/)
{
  if(vtxverbose>5){std::cout<<"InteractionPoint::initPar2: - "<<std::endl;}
  // nothing left to do: actually, should never be called
  assert(0) ;
  return ErrCode::success ;
}


ErrCode
DecayTreeFitter::InteractionPoint::initCov(FitParams* fitpar) const
{
  ErrCode status ;
  int posindex = posIndex() ;
  for(int row=0; row<3; row++)
  {
    if(m_hasposcov) fitpar->cov()(posindex+row,posindex+row) = 1000*m_ipPosCov(row,row) ;
    else fitpar->cov()(posindex+row,posindex+row) = 1.; //1cm default?
  }

  //int momindex = momIndex() ; 
  //if(m_hasmomcov) for(int row=0; row<4; row++)
    //fitpar->cov()(momindex+row,momindex+row) = 1000*m_ipMomCov(row,row) ;

  for(daucontainer::const_iterator it = daughters().begin() ;
      it != daughters().end() ; ++it)
    status |= (*it)->initCov(fitpar) ;
  return status ;
}

ErrCode
DecayTreeFitter::InteractionPoint::projectIPConstraint(const FitParams* fitparams,
                                      Projection& p) const
{
  int posindex = posIndex() ;
  if(vtxverbose>6){std::cout<<"InteractionPoint::projectIPConstraint(): posindex="<<posindex<<std::endl;}
  if(m_haspos){
    for(int row=0; row<3; ++row) {
      p.r(row) =  fitparams->par()(posindex+row) - m_ipPos(row); 
      p.H(row,posindex+row) = 1 ;
      //if(m_hasposcov)
        for(int col=0; col<3; ++col)
          p.Vfast(row,col) = m_ipPosCov(row,col) ;
    }
  }
  if(vtxverbose>6){
    std::cout<<"InteractionPoint::projectIPConstraint(): projection is:"<<posindex<<std::endl;
    p.r().Print();
    p.V().Print();
    p.H().Print();
    }
  return ErrCode::success ;
}

ErrCode
DecayTreeFitter::InteractionPoint::projectBeamConstraint(const FitParams* fitparams,
                                      Projection& p) const
{
  int momindex = momIndex() ; // get momentum part of the only daughter (head of tree)
  if(vtxverbose>6){std::cout<<"InteractionPoint::projectIPConstraint(): daumomindex="<<momindex<<std::endl;}
  if(m_hasmom){
    for(int row=0; row<4; ++row) {
      p.r(row) = fitparams->par()(momindex+row) - m_ipMom(row);
      p.H(row,momindex+row) = 1 ;
      if(m_hasmomcov)
        for(int col=0; col<4; ++col)
          p.Vfast(row,col) = m_ipMomCov(row,col) ;
    }
  }
  if(vtxverbose>6){
    std::cout<<"InteractionPoint::projectBeamConstraint(): projection is:"<<momindex<<std::endl;
    p.r().Print();
    p.V().Print();
    p.H().Print();
    }
  return ErrCode::success ;
}

ErrCode
DecayTreeFitter::InteractionPoint::projectConstraint(Constraint::Type aType,
                                    const FitParams* fitparams,
                                    Projection& p) const
{
  ErrCode status ;
  switch(aType) {
    case Constraint::beamspot:
      status |= projectIPConstraint(fitparams,p) ;
      break ;
    case Constraint::beamenergy:
      status |= projectBeamConstraint(fitparams,p) ;
      break ;
    default:
      status |= ParticleBase::projectConstraint(aType,fitparams,p) ;
  }
  return status ;
}

double
DecayTreeFitter::InteractionPoint::chiSquare(const FitParams* fitparams) const
{
  double chisq=0.;
  if(m_haspos&&m_hasposcov) 
  { // no covariance, no chi2 contribution!
    int posindex = posIndex() ;
    TVectorD residualpos(3) ;
    for(int row=0; row<3; ++row)
      residualpos(row) = fitparams->par()(posindex+row) - m_ipPos(row) ;
    chisq += m_ipPosCovInv.Similarity(residualpos);
  }
  if(m_hasmom&&m_hasmomcov) 
  { // no covariance, no chi2 contribution!
    int momindex = momIndex() ;
    TVectorD residualmom(4) ;
    for(int row=0; row<4; ++row)
      residualmom(row) = m_ipMom(row) - fitparams->par()(momindex+row) ;
    chisq += m_ipMomCovInv.Similarity(residualmom);
  }
  // add the daughters' chi2
  chisq += ParticleBase::chiSquareD(fitparams) ;

  return chisq ;
}

void DecayTreeFitter::InteractionPoint::addToConstraintList(constraintlist& alist, int depth) const
{
  // first the daughters
  for(daucontainer::const_iterator it = daughters().begin() ;
      it != daughters().end() ; ++it)
    (*it)->addToConstraintList(alist,depth-1) ;

  // then the beamspot
  if(m_haspos) alist.push_back(Constraint(this,Constraint::beamspot,depth,3)) ;
  if(m_hasmom) alist.push_back(Constraint(this,Constraint::beamenergy,depth,4)) ;
}


std::string DecayTreeFitter::InteractionPoint::parname(int thisindex) const
{
  int id = thisindex ;
  // skip the lifetime parameter name
  if(id>=3) ++id ;
  return ParticleBase::parname(id) ;
}