/************************************************************************
 * TCtvmft class implementation file                                  *
 *                                                                      *
 *                                                                      *
 * Author: Craig Blocker, CDF Group                                     *
 *         Phone 781-736-2920                                           *
 *                                                                      *
 * Description: Routines to implement TCtvmft class                   *
 *
 * See TCtvmft.hh in the Sin2BetaMods/TCtvmft area and
 * the file TCtvmft.txt in the /doc area of Sin2BetaMods for
 * documentation.
 *                                                                      *
 * Revision History:                                                    *
 *    May 27, 1999   Craig Blocker   Initial creation                   * 
 ************************************************************************/

// Apr 21 2001 P.Murat: convert to ROOT-compliant format

#include "TMath.h"

#include "Stntuple/fit/TCtvmft.hh"
#include "Stntuple/obj/TStnTrack.hh"

extern "C" { 
  void  ctvmft_(int&,int&,int&);
  bool  mcalc_(int&,int*,float&,float&,double*);
  void  dcalc_(int&,int&,float*,float&,float&,float*);
 }

ClassImp(TCtvmft)
//_____________________________________________________________________________
TCtvmft::TCtvmft() {
  // Initialize various arrays
  // First get pointers to various FORTRAN common blocks.

  _ctvmq       = (CTVMQ*)  ctvmq_address_ ();
  _ctvmfr      = (CTVMFR*) ctvmfr_address_();
  _fiddle      = (FIDDLE*) fiddle_address_();
  _trkprm      = (TRKPRM*) trkprm_address_();
  _print_level = (PRINT_LEVEL*) print_level_address_();
  
					// Now initialize CTVMQ common.
  _ctvmq->runnum=1;
  _ctvmq->trgnum=100;
  float bmag=14.116;			// Eventually, we have to get the
					// magnetic field from the right place.
  _ctvmq->pscale=0.000149896*bmag;
  _print_level->print_level = 0;

  Init();
}

//_____________________________________________________________________________
void TCtvmft::Init() {
  // Aug 15 1999 P.Murat: move initializations which can be done multiple times
  // into Init() function

  _ctvmq->ntrack=0;
  _ctvmq->nvertx=0;
  _ctvmq->nmassc=0;
  for(int j=0; j<_maxtrk; ++j) {
    _ctvmq->list[j]=0;
    for(int jv=0; jv<_maxvtx; ++jv) {
      _ctvmq->trkvtx[jv][j]=false;
    }
    for(int jmc=0; jmc<_maxmcn; ++jmc) {
      _ctvmq->trkmcn[jmc][j]=false;
    }
  }
  for(int jv=0; jv<_maxvtx; ++jv) {
    _ctvmq->cvtx[jv]=0;
    _ctvmq->vtxpnt[0][jv]=-1;
    _ctvmq->vtxpnt[1][jv]=0;
  }
				// Don't bomb program on error.
  _fiddle->excuse=1;
  				// Set stat to -999 to symbolize that no fit 
				// has yet been done.
  stat=-999;
}

//_____________________________________________________________________________
TCtvmft::~TCtvmft(){ 
}

//_____________________________________________________________________________
bool TCtvmft::AddTrack(const TStnTrack* trk, float mass, int jv)
{
  // Check that this vertex number is within the allowed range.

  if (jv<1 || jv>_maxvtx) return false;
  _ctvmq->nvertx = jv>_ctvmq->nvertx ? jv : _ctvmq->nvertx;

   // Check that we have exceeded the maximum number of tracks.

  if(_ctvmq->ntrack>=_maxtrk) return false;

					// Add this track.

  _ctvmq->list[_ctvmq->ntrack]=trk->Number();

  _ctvmq->tmass[_ctvmq->ntrack]=mass;
  _ctvmq->trkvtx[jv-1][_ctvmq->ntrack]=true;

// Put this track's helix parameters and error matrix into
// a fortran common block so that they can  be accessed
// by gettrk.  This is a dummy for now.

  _trkprm->trhelix[_ctvmq->ntrack][0]=trk->Lam0();
  _trkprm->trhelix[_ctvmq->ntrack][1]=trk->C0();
  _trkprm->trhelix[_ctvmq->ntrack][2]=trk->Z0();
  _trkprm->trhelix[_ctvmq->ntrack][3]=trk->D0();
  _trkprm->trhelix[_ctvmq->ntrack][4]=trk->Phi0();

  TMatrix55& cov = *((TStnTrack*)trk)->Cov();
  for(int j=0; j<5; ++j) {
    for(int k=0; k<5; ++k) {
      _trkprm->trem[_ctvmq->ntrack][j][k]=cov(j,k);
    }
  }

  _ctvmq->ntrack++;

  return true;
}

//_____________________________________________________________________________
bool TCtvmft::VertexPoint(int jv1, int jv2, int type) {
  // Check that these vertex numbers are within allowed range.

  if(jv1>_maxvtx || jv1<1) return false;
  if(jv2>_maxvtx || jv2<0) return false;
					// Setup vertex pointing.
  _ctvmq->vtxpnt[0][jv1]=jv2;
  _ctvmq->vtxpnt[1][jv1]=type;

  return true;
}

//_____________________________________________________________________________
bool TCtvmft::Conversion2D(int jv) {

  if (jv<1 || jv>_ctvmq->nvertx) return false;

  _ctvmq->cvtx[jv-1]=1;
  return true;
}

//_____________________________________________________________________________
bool TCtvmft::Conversion3D(int jv) {
  if(jv<1 || jv>_ctvmq->nvertx) return false;

  _ctvmq->cvtx[jv-1]=2;
  return true;
}

//_____________________________________________________________________________
bool TCtvmft::MassConstrain(const TStnTrack* trk1,
                              const TStnTrack* trk2, 
			      float            mass)
{
   int ntrk=2;
   const TStnTrack* tracks[2];
   tracks[0]=trk1;
   tracks[1]=trk2;
   return MassConstrain(ntrk,tracks,mass);
}

//_____________________________________________________________________________
bool TCtvmft::MassConstrain(const TStnTrack* trk1, 
			      const TStnTrack* trk2,
			      const TStnTrack* trk3, 
			      float mass)
{
   int ntrk=3;
   const TStnTrack* tracks[3];
   tracks[0]=trk1;
   tracks[1]=trk2;
   tracks[2]=trk3;
   return MassConstrain(ntrk,tracks,mass);
}

//_____________________________________________________________________________
bool TCtvmft::MassConstrain(const TStnTrack* trk1, 
			      const TStnTrack* trk2,
			      const TStnTrack* trk3, 
			      const TStnTrack* trk4,
			      float mass)
{
   int ntrk=4;
   const TStnTrack* tracks[4];
   tracks[0]=trk1;
   tracks[1]=trk2;
   tracks[2]=trk3;
   tracks[3]=trk4;
   return MassConstrain(ntrk,tracks,mass);
}

//_____________________________________________________________________________
bool TCtvmft::MassConstrain(int ntrk, const TStnTrack* tracks[], float mass)
{
  // Check that we have not exceeded the allowed number of mass constraints.

  if(_ctvmq->nmassc>=_maxmcn) return false;

					// Set constraint mass
  _ctvmq->cmass[_ctvmq->nmassc]=mass;
     
// For each track in contraint, set trkmcn true.  Since the number in
// tracks[] is the track number, we have to find each track in the
// list of tracks.

  for(int jt=0; jt<ntrk; ++jt) {
    bool found=false;
    for(int kt=0; kt<_ctvmq->ntrack; ++kt) {
      if (tracks[jt]->Number() == _ctvmq->list[kt]) {
	_ctvmq->trkmcn[_ctvmq->nmassc][kt]=true;
	found=true;
      }
    }
    if (!found) return false;
  }
					// Increment number of mass constraints
  _ctvmq->nmassc++;

  return true;
}

//_____________________________________________________________________________
bool TCtvmft::Fit() {
  // Do the vertex fit.

  int print=0;
  int level=0;
  ctvmft_(print,level,stat);

  return (stat == 0);
}

//_____________________________________________________________________________
void TCtvmft::PrintAll(Option_t* opt) const {
   printf("****************************** TCtvmft ");
   printf("******************************\n");
   printf("Number of tracks: %5i\n",_ctvmq->ntrack);
   printf("   Tracks: ");
   for (int jt=0; jt<_ctvmq->ntrack; ++jt) {
     if (jt != 0) printf(",  ");
     printf("%i5",_ctvmq->list[jt]);
   }
   printf("\n");

   printf("Number of vertices: %5i\n",_ctvmq->nvertx);
   for (int jv=0; jv<_ctvmq->nvertx; ++jv) {
     if(_ctvmq->vtxpnt[0][jv]==0) {
       printf("   Vertex %5i points to the primary vertex ",jv+1);
     }
     else if(_ctvmq->vtxpnt[0][jv]>0) {
       printf("   Vertex %5i points to the vertex %5i",
	      jv+1,_ctvmq->vtxpnt[0][jv]);
     }
     if (_ctvmq->vtxpnt[1][jv]==1) {
       printf("in 2 dimensions\n");
     }
     else if(_ctvmq->vtxpnt[1][jv]==2) {
       printf("in 3 dimensions\n");
     }
     else if(_ctvmq->vtxpnt[1][jv]==3) {
       printf(", a single track vertex\n");
     }

     if(_ctvmq->cvtx[jv]>0) {
       printf("   Vertex %5i is a conversion\n",jv+1);
     }
   }

   printf("Number of mass constraints: %5i\n",_ctvmq->nmassc);

   for(int jmc=0; jmc<_ctvmq->nmassc; ++jmc) {
     printf("   Tracks ");
     for(int jt=0; jt<_ctvmq->ntrack; ++jt) {
       if(_ctvmq->trkmcn[jmc][jt]) {
	 printf(" %5i",_ctvmq->list[jt]);
       }
     }
     printf(" constrained to mass %10.5f Gev/c^2\n",_ctvmq->cmass[jmc]);
   }

   if (stat == -999) {
     printf("No fit has been done.\n");
   }
   else {
      printf("***** Results of Fit *****\n");
      PrintErr();
      printf("   Status = %5i\n",stat);
      printf("   Chi-square = %10.3e for %3i degrees of freedom\n",
	     _ctvmq->chisqr[0],
	     _ctvmq->ndof);
      printf("   => probability = %10.3e\n",GetProb());
      for(int jv=0; jv<_ctvmq->nvertx; ++jv) {
	printf("Vertex %3i position: %10.4f %10.4f %10.4f\n",
	       jv+1,
	       _ctvmq->xyzvrt[jv+1][0],
	       _ctvmq->xyzvrt[jv+1][1],
	       _ctvmq->xyzvrt[jv+1][2]);
      }

      for(int jt=0; jt<_ctvmq->ntrack; ++jt) {
          printf("Track %3i PG: %10.4f %10.4f %10.4f %10.4f %10.4f %10.4f\n",
		 _ctvmq->list[jt],
		 _ctvmq->trkp4[0][jt],
		 _ctvmq->trkp4[1][jt],
		 _ctvmq->trkp4[2][jt],
		 _ctvmq->trkp4[3][jt],
		 _ctvmq->trkp4[4][jt],
		 _ctvmq->trkp4[5][jt]);
      }
   }
   printf("****************************************");
   printf("**************************\n");
}

//_____________________________________________________________________________
void TCtvmft::PrintErr() const {
  // print CTVMFT errors

  printf("   IJKERR = %5i, %5i, %5i\n",
	 _ctvmq->ijkerr[0],
	 _ctvmq->ijkerr[1],
	 _ctvmq->ijkerr[2]);
}

//_____________________________________________________________________________
float TCtvmft::GetProb() const {
  // Probability of chi-square of fit with given number of degrees of freedom

  if (_ctvmq->chisqr[0] > 0.) {
    return TMath::Prob(_ctvmq->chisqr[0],_ctvmq->ndof);
  }
  else {
    return 0.;
  }
}

//_____________________________________________________________________________
int TCtvmft::GetTrackP4(int It, TLorentzVector& P4) {
  // return four momentum of fit track 

  int rc = -1;

  for(int jt=0; jt<_ctvmq->ntrack; ++jt) {

					// Find which track matches this Id.
    if (It == _ctvmq->list[jt]) {
      P4.SetXYZT(_ctvmq->trkp4[0][jt], _ctvmq->trkp4[1][jt],
		 _ctvmq->trkp4[2][jt], _ctvmq->trkp4[3][jt]);
      rc = 0;
      break;
    }
  }
  return rc;
}

//_____________________________________________________________________________
int TCtvmft::GetTrackP4(TStnTrack* Trk, TLorentzVector& P4) {
  // return four momentum of fit track 

  return GetTrackP4(Trk->Number(),P4);
}

//_____________________________________________________________________________
float TCtvmft::GetMass(const TStnTrack* trk1,
                         const TStnTrack* trk2, 
			 float&           dmass) 
{
  // Get the fit invariant mass of tracks trk1 and trk2.
  // dmass is the error on the mass.

   int ntrk=2;
   const TStnTrack* trks[2];
   trks[0]=trk1;
   trks[1]=trk2;
   return GetMass(ntrk,trks,dmass);
}

//_____________________________________________________________________________
float TCtvmft::GetMass(const TStnTrack* trk1, 
			 const TStnTrack* trk2,
			 const TStnTrack* trk3,
			 float&           dmass)
{
  // Get the fit invariant mass of tracks trk1, trk2, and trk3.
  // dmass is the error on the mass.

   int ntrk=3;
   const TStnTrack* trks[3];
   trks[0]=trk1;
   trks[1]=trk2;
   trks[2]=trk3;
   return GetMass(ntrk,trks,dmass);
}


//_____________________________________________________________________________
float TCtvmft::GetMass(const TStnTrack* trk1, 
			 const TStnTrack* trk2,
			 const TStnTrack* trk3, 
			 const TStnTrack* trk4, 
			 float& dmass)
{
  // Get the fit invariant mass of tracks trk1, trk2, trk3, and trk4.
  // dmass is the error on the mass.

   int ntrk=4;
   const TStnTrack* trks[4];
   trks[0]=trk1;
   trks[1]=trk2;
   trks[2]=trk3;
   trks[3]=trk4;
   return GetMass(ntrk,trks,dmass);
}

//_____________________________________________________________________________
float TCtvmft::GetMass(int ntrk, const TStnTrack* tracks[], float& dmass)
{
  // Get fit invariant mass of ntrk tracks listed in array tracks.
  // dmass is the error on the mass.

  dmass=-999.;
  if (ntrk <= 0) return 0;

  int jtrks[_maxtrk];
  for(int jt=0; jt<ntrk; ++jt) {
    bool found=false;
    for(int kt=0; kt<_ctvmq->ntrack; ++kt) {
      if(tracks[jt]->Number() == _ctvmq->list[kt]) {
	found=true;
	jtrks[jt]=kt+1;
      }
    }
    if(!found) return 0;
  }
  int ntr=ntrk;
  float mass;
  double p4[4];
  mcalc_(ntr, jtrks, mass, dmass, p4);
  return mass;
}


//_____________________________________________________________________________
float TCtvmft::GetDecayLength(int             nv, 
				int             mv, 
				const TVector3& dir, 
				float&          dlerr)
{
// Get the signed decay length from vertex nv (primary) 
// to vertex mv (secondary)
// along the x-y direction defined by the 3 vector dir.
// dlerr is the error on the decay length including the
// full error matrix.
// Check that vertices are within range.

  if (nv<0 || nv >= _ctvmq->nvertx) return -999.;
  if (mv<1 || mv >  _ctvmq->nvertx) return -999.;
  if (nv >= mv)                     return -999.;
  
  float dir_t = dir.Pt();

  if (dir_t <= 0.) return -999.;

  TVector3  v1;
  TVector3  v2;
  TVector3  dv;
  
  GetVertex(nv,v1);
  GetVertex(mv,v2);
  dv  = v2-v1;
  dv.SetZ(0.0);

  float dl = dv*dir/dir_t;
					// Set up the column matrix of derivatives
  TMatrix A(4,1);
  A(0,0) =  dir.X()/dir_t;
  A(1,0) =  dir.Y()/dir_t;
  A(2,0) = -dir.X()/dir_t;
  A(3,0) = -dir.Y()/dir_t;
//-----------------------------------------------------------------------------
// Check if first vertex (nv) is primary vertex.  If it is,
// check if it was used in the primary vertex.  If not,
// all of the corresponding error matrix elements are those
// supplied for the primary vertex.
//-----------------------------------------------------------------------------
  int nvf=0;
  if (nv==0) {
    nvf=-1;
    for (int jv=0; jv<_ctvmq->nvertx; ++jv) {
      if (_ctvmq->vtxpnt[jv][0]==0) nvf=0;
    }
  }
					// Get the relevant submatrix of the 
					// full error matrix.
  TMatrix V(4,4);
  if (nvf < 0) {
    V(0,0)=GetErrorMatrix(_ctvmq->voff[mv-1]+1,_ctvmq->voff[mv-1]+1);
    V(0,1)=GetErrorMatrix(_ctvmq->voff[mv-1]+1,_ctvmq->voff[mv-1]+2);
    V(1,0)=GetErrorMatrix(_ctvmq->voff[mv-1]+2,_ctvmq->voff[mv-1]+1);
    V(1,1)=GetErrorMatrix(_ctvmq->voff[mv-1]+2,_ctvmq->voff[mv-1]+2);
    V(2,2)=_ctvmq->exyzpv[0][0];
    V(2,3)=_ctvmq->exyzpv[0][1];
    V(3,2)=_ctvmq->exyzpv[1][0];
    V(3,3)=_ctvmq->exyzpv[1][1];
  }
  else {
					// Get the indices into the error 
					// matrix vmat

    int index[4]={_ctvmq->voff[mv-1]+1,_ctvmq->voff[mv-1]+2,0,0};
    if (nv==0) {
      index[2]=1;
      index[3]=2;
    }
    for(int j=0; j<0; ++j) {
      for(int k=0; k<0; ++k) {
	V(j,k)=GetErrorMatrix(index[j],index[k]);
      }
    }
  }
					// Calculate square of dlerr
  TMatrix temp(V,TMatrix::kMult,A);
  TMatrix mat(A,TMatrix::kTransposeMult,temp);
  //TMatrix mat(A,TMatrix::kAtBA,V);
  dlerr = mat(0,0);
  if (dlerr >= 0.) {
    dlerr = sqrt(dlerr);
  }
  else {
    dlerr = -sqrt(-dlerr);
  }
  
  return dl;
}

//_____________________________________________________________________________
float TCtvmft::GetDr(int nv, int mv, float& drerr) {
// Get the transvese distance between vertices nv and mv
// and return it as the function value.  drerr is the
// uncertainty on the transverse distance, calculated from the
// full error matrix including correlations.

  float dxyz[3];
  float dr;
  float dz;
  float dl[3];
  
  dcalc_(mv,nv,dxyz,dr,dz,dl);
  drerr=dl[0];
  return dr;
}

//_____________________________________________________________________________
float TCtvmft::GetDz(int nv, int mv, float& dzerr)
{
// Get the longitudinal distance between vertices nv and mv
// and return it as the function value.  dzerr is the
// uncertainty on the longitudinal distance, calculated from the
// full error matrix including correlations.

  float dxyz[3];
  float dr;
  float dz;
  float dl[3];

  dcalc_(mv,nv,dxyz,dr,dz,dl);
  dzerr=dl[1];
  return dz;
}

//_____________________________________________________________________________
int TCtvmft::GetVertex(int nv, TVector3& vertex) {
  // Return x,y,z position of vertex nv.

  if (nv<0 || nv>_ctvmq->nvertx) return -1;
//-----------------------------------------------------------------------------
// Check if first vertex (nv) is primary vertex.  If it is,
// check if it was used in the primary vertex.  If not,
// all of the corresponding error matrix elements are those
// supplied for the primary vertex.
//-----------------------------------------------------------------------------
  int nvf=0;
  if (nv==0) {
    nvf=-1;
    for(int jv=0; jv<_ctvmq->nvertx; ++jv) {
      if(_ctvmq->vtxpnt[jv][0]==0) nvf=0;
    }
  }
					// If primary vertex was not part of 
					// fit, take vertex location as 
					// supplied.
  if (nvf < 0) {
    vertex.SetXYZ(_ctvmq->xyzpv0[0],_ctvmq->xyzpv0[1],_ctvmq->xyzpv0[2]);
  }
  else {
    vertex.SetXYZ(_ctvmq->xyzvrt[nv][0],
		  _ctvmq->xyzvrt[nv][1],
		  _ctvmq->xyzvrt[nv][2]);
  }
  return 0;
}

//_____________________________________________________________________________
void TCtvmft::SetPrimaryVertex(float xv, float yv, float zv) {
  // Set x,y,z position of the primary vertex.

  _ctvmq->xyzpv0[0]=xv;
  _ctvmq->xyzpv0[1]=yv;
  _ctvmq->xyzpv0[2]=zv;
}

//_____________________________________________________________________________
void TCtvmft::SetPrimaryVertex(const TVector3* v) {
  // Set x,y,z position of the primary vertex.

  _ctvmq->xyzpv0[0]=v->X();
  _ctvmq->xyzpv0[1]=v->Y();
  _ctvmq->xyzpv0[2]=v->Z();
}

//_____________________________________________________________________________
void TCtvmft::SetPrimaryVertexError(const float* Xverr) {
  // Set the error matrix for the primary vertex.

  for(int j=0; j<3; ++j) {
    for(int k=0; k<3; ++k) {
      _ctvmq->exyzpv[j][k]=Xverr[3*j+k];
    }
  }
}

//_____________________________________________________________________________
double TCtvmft::GetErrorMatrix(int j, int k) {
  // Return the j,k element of the full error matrix VMAT.
  // Since the CTVMFT documentation assumes the indices
  // start from 1 (ala Fortran), we will also assume this
  // and convert the C++ indices.  Note also the that order of
  // Fortran and C++ indices is different.  We assume that
  // j and k are in the Fortran order.

  if(j<1 || k<1 || j>_maxdim+1 || k>_maxdim) {
    return -999.;
  }
  return _ctvmfr->vmat[k-1][j-1];
}