//*--- Author: Erik Krebs

#include "TMatrixD.h"
#include "TMath.h"
#include "TVector3.h"

#include "hkalplane.h"
#include "hkaltrackstate.h"

#include <iostream>
using namespace std;

ClassImp(HKalTrackState)

//_HADES_CLASS_DESCRIPTION
///////////////////////////////////////////////////////////////////////////////
//
// Class: HKalTrackState
//
// The state of the particle track at a given site k can be described by a vector
// a_k containing a set of track parameters:
// 1. x0: x position
// 2. y0: y position
// 3. tx: p_x / p_z
// 4. ty: p_y / p_z
// 5. qp: particle charage in elementary charges divided by
//        magnitude of momentum in MeV/c
//
// The class also stores the matrices needed for the Kalman filter:
// 1. propagator: F_k = @f_k(a_k) / @a_k
// 2. projector : H_k = @h_k(a_k) / @a^(k-1)_k
// 3. covariance
// 4. process noise
// where f_k(a_k) is the propagator function that propagates the track to
// the next site and h_k(a_k) projects the track state onto a measurement vector.
//
//-----------------------------------------------------------------------
///////////////////////////////////////////////////////////////////////////////


//  -----------------------------------
//  Ctors and Dtor
//  -----------------------------------

HKalTrackState::HKalTrackState(Kalman::kalFilterTypes stateType, Int_t measDim, Int_t stateDim)
: TObject(),
  fPropagator(stateDim, stateDim),
  fProjector(measDim, stateDim),
  fCovariance(stateDim, stateDim),
  fProcessNoise(stateDim, stateDim),
  stateVec(stateDim) {
    // Create state with dummy state vector and matrices.
    // stateType: state is either for the prediction, filtering, smoothing or inverse filtering step in the Kalman filter
    // measDim:  dimension of measurement vector
    // stateDim: dimension of state vector.

    type = stateType;
    fPropagator.UnitMatrix();
    fProjector.UnitMatrix();
    fCovariance.UnitMatrix();
}

HKalTrackState::HKalTrackState(Kalman::kalFilterTypes stateType, const TVectorD &sv, Int_t measDim)
: TObject(),
  fPropagator(sv.GetNrows(), sv.GetNrows()),
  fProjector(measDim, sv.GetNrows()),
  fCovariance(sv.GetNrows(), sv.GetNrows()),
  fProcessNoise(sv.GetNrows(), sv.GetNrows()),
  stateVec(sv) {
    // Create a state with dummy matrices and sets the state vector equal to function
    // parameter sv.
    // stateType: state is either for the prediction, filtering, smoothing or inverse filtering step in the Kalman filter
    // sv:        state vector.
    // measDim:   dimension of measurement vector.

    type = stateType;
    fPropagator.UnitMatrix();
    fProjector.UnitMatrix();
    fCovariance.UnitMatrix();
}

HKalTrackState::~HKalTrackState() {
}

//  -----------------------------------
//  Implementation of public methods
//  -----------------------------------

void HKalTrackState::calcMeasVec(TVectorD &projMeasVec) const {
    // Calculates a measurement vector from this state.
    // projMeasVec: return value.

#ifdef kalDebug
    if(projMeasVec.GetNrows() != getMeasDim()) {
        Warning("calcMeasVec", Form("Needed to resize TVectorD. Dimension of measurement vector (%i) does not match that of function parameter (%i).", getMeasDim(), projMeasVec.GetNrows()));
        projMeasVec.ResizeTo(getMeasDim());
    }
#endif
    projMeasVec(0) = stateVec(kIdxX0);
    projMeasVec(1) = stateVec(kIdxY0);
}

void HKalTrackState::calcDir(TVector3 &dir) const {
    // Calculates a direction unit vector from this state.
    // dir: track direction (return value).

    calcDir(dir, stateVec);
}

void HKalTrackState::calcDir(TVector3 &dir, const TVectorD &sv) {
    // Calculates a direction unit vector from a state vector given by function
    // parameter sv.
    // dir: track direction (return value)
    // sv:  state vector to calculate direction from.

    Double_t tanx = sv(kIdxTanPhi);
    Double_t tany = sv(kIdxTanTheta);
    Double_t qp = sv(kIdxQP);
    dir.SetZ( 1./(TMath::Abs(qp) * TMath::Sqrt(tanx*tanx + tany*tany + 1)) );
    dir.SetX(tanx * dir.Z());
    dir.SetY(tany * dir.Z());
    dir = dir.Unit();
}

void HKalTrackState::calcMomVec(TVector3 &dir) const {
    // Calculates the momentum vector from this state.
    // dir: track momentum (return value).

    calcMomVec(dir, stateVec);
}

void HKalTrackState::calcMomVec(TVector3 &dir, const TVectorD &sv) {
    // Calculates momentum vector from a state vector given by function
    // parameter sv.
    // dir: track momentum (return value)
    // sv:  state vector to calculate direction from.

    Double_t tanx = sv(kIdxTanPhi);
    Double_t tany = sv(kIdxTanTheta);
    Double_t qp = sv(kIdxQP);
    dir.SetZ( 1./(TMath::Abs(qp) * TMath::Sqrt(tanx*tanx + tany*tany + 1)) );
    dir.SetX(tanx * dir.Z());
    dir.SetY(tany * dir.Z());
}

void HKalTrackState::calcPosAtPlane(TVector3 &pos, const HKalPlane &plane) const {
    // Calculate the track position as a three dimensional vector from this state.
    // Only two coordinates are stored in the state vector. The z-coordinate is
    // calculated by going in z-direction and find the intersection with a plane.
    // pos:   track position (return value).
    // plane: the plane the track is on.

    calcPosAtPlane(pos, plane, stateVec);
}

void HKalTrackState::calcPosAtPlane(TVector3 &pos, const HKalPlane &plane, const TVectorD &sv) {
    // Calculate the track position as a three dimensional vector from the state
    // given by function parameter sv.
    // Only two coordinates are stored in the state vector. The z-coordinate is
    // calculated by going in z-direction and find the intersection with a plane.
    // pos:   track position (return value).
    // plane: the plane the track is on.

    TVector3 posFrom;
    posFrom.SetX(sv(kIdxX0));
    posFrom.SetY(sv(kIdxY0));
    posFrom.SetZ(0.);
    TVector3 dir(0., 0., 1.);
    plane.findIntersection(pos, posFrom, dir);
}

void HKalTrackState::calcPosAndDirAtPlane(TVector3 &pos, TVector3 &dir, const HKalPlane &plane) const {
    // Calculate the track position and direction as three dimensional vectors
    // from this state.
    // pos:   track position (return value)
    // dir:   track direction (return value)
    // plane: the plane the track is on
    calcPosAtPlane(pos, plane);
    calcDir(dir);
}

void HKalTrackState::calcStateVec(TVectorD &sv, Double_t qp, const TVector3 &pos, const TVector3 &dir) {
    // Return the state vector from a position, direction and momentum.
    // sv:  the state vector (return value)
    // qp:  particle charage in elementary charges divided by magnitude of momentum in MeV/c
    // pos: track position
    // dir: track direction.

#ifdef kalDebug
    Int_t retDim = sv.GetNrows();
    Int_t stateDim = 5;
    if(retDim != stateDim) {
        sv.ResizeTo(stateDim);
        cout<<"Needed to resize TVectorD. Dimension of function parameter ("<<retDim<<") does not match state vector dimension ("<<stateDim<<")."<<endl;
        //Warning("calcStateVec", Form("Needed to resize TVectorD. Dimension of function parameter (%i) does not match state vector dimension (%i).", retDim, stateDim));
    }
#endif
    sv(kIdxX0)       = pos.x();
    sv(kIdxY0)       = pos.y();
    sv(kIdxTanPhi)   = dir.x() / dir.z();
    sv(kIdxTanTheta) = dir.y() / dir.z();
    sv(kIdxQP)       = qp;
}

void HKalTrackState::clear() {
    // Resets state vector to zeroes, matrices to unit matrix.

    stateVec.Zero();
    fPropagator.UnitMatrix();
    fProjector.UnitMatrix();
    fCovariance.UnitMatrix();
    fProcessNoise.UnitMatrix();
}

void HKalTrackState::print(const Option_t *opt) const {
    // Prints information about the state.
    // opt: Determines what information about the object will be printed.
    // Capitalization does not matter.
    // If opt contains:
    // "S": print state vector
    // "F": print propagator matrix
    // "C": print covariance matrix
    // "H": print projector matrix
    // "Q": print process noise matrix
    // If opt is empty all of the above will be printed.

    TString stropt(opt);

    switch (type) {
    case kPredicted:
        cout<<"**** Predicted state: ****"<<endl;
        break;
    case kFiltered:
        cout<<"**** Filtered state: ****"<<endl;
        break;
    case kSmoothed:
        cout<<"**** Smoothed state: ****"<<endl;
        break;
    case kInvFiltered:
        cout<<"**** Inverse filtered vector: ****"<<endl;
        break;
    }
    if(stropt.Contains("S", TString::kIgnoreCase) || stropt.IsNull()) {
        cout<<"State vector:"<<endl;
        stateVec.Print();
    }
    if(stropt.Contains("F", TString::kIgnoreCase) || stropt.IsNull()) {
        cout<<"Propagator matrix:"<<endl;
        fPropagator.Print();
    }
    if(stropt.Contains("C", TString::kIgnoreCase) || stropt.IsNull()) {
        cout<<"Covariance matrix:"<<endl;
        fCovariance.Print();
    }
    if(stropt.Contains("H", TString::kIgnoreCase) || stropt.IsNull()) {
        cout<<"Projector matrix:"<<endl;
        fProjector.Print();
    }
    if(stropt.Contains("Q", TString::kIgnoreCase) || stropt.IsNull()) {
        cout<<"Process noise matrix:"<<endl;
        fProcessNoise.Print();
    }
    cout<<"**** End of track state print. ****"<<endl;
}

void HKalTrackState::transform(const TRotation *transMat, const HKalPlane &plane) {
    // Transform the state vector of this object using a rotation matrix.
    // transMat: the rotation matrix
    // plane:    the plane the track is on. Needed to calculate the z-coordinate of the
    //           track position.

    TVector3 pos; // track position
    TVector3 dir; // track direction
    calcPosAndDirAtPlane(pos, dir, plane);
    pos.Transform(*transMat);
    dir.Transform(*transMat);
    Double_t qp = getStateParam(kIdxQP);
    setStateVec(qp, pos, dir);
}

void HKalTrackState::setCovMat(const TMatrixD &fCov) {
    // Replace the state's covariance matrix.
#ifdef kalDebug
    Int_t nRowsOld = getCovMat().GetNrows();
    Int_t nColsOld = getCovMat().GetNcols();
    Int_t nRowsNew = fCov.GetNrows();
    Int_t nColsNew = fCov.GetNcols();
    if(nRowsOld != nRowsNew || nColsOld != nColsNew) {
        Error("setCovMat()", Form("Matrices not compatible. Cannot replace %ix%i matrix with %ix%i matrix.", nRowsOld, nColsOld, nRowsNew, nColsNew));
        exit(1);
    }
#endif
    fCovariance = fCov;
}

void HKalTrackState::setPropMat(const TMatrixD &fProp) {
    // Replace the state's propagator matrix.

#ifdef kalDebug
    Int_t nRowsOld = getPropMat().GetNrows();
    Int_t nColsOld = getPropMat().GetNcols();
    Int_t nRowsNew = fProp.GetNrows();
    Int_t nColsNew = fProp.GetNcols();
    if(nRowsOld != nRowsNew || nColsOld != nColsNew) {
        Error("setPropMat()", Form("Matrices not compatible. Cannot replace %ix%i matrix with %ix%i matrix.", nRowsOld, nColsOld, nRowsNew, nColsNew));
        exit(1);
    }
#endif
    fPropagator = fProp;
}

void HKalTrackState::setProjMat(const TMatrixD &fProj) {
    // Replace the state's projector matrix.

#ifdef kalDebug
    Int_t nRowsOld = getProjMat().GetNrows();
    Int_t nColsOld = getProjMat().GetNcols();
    Int_t nRowsNew = fProj.GetNrows();
    Int_t nColsNew = fProj.GetNcols();
    if(nRowsOld != nRowsNew || nColsOld != nColsNew) {
        Error("setProjMat()", Form("Matrices not compatible. Cannot replace %ix%i matrix with %ix%i matrix.", nRowsOld, nColsOld, nRowsNew, nColsNew));
        exit(1);
    }
#endif
    fProjector = fProj;
}

void HKalTrackState::setProcNoiseMat(const TMatrixD &fProc) {
    // Replace the state's process noise matrix.

#ifdef kalDebug
    Int_t nRowsOld = getProcNoiseMat().GetNrows();
    Int_t nColsOld = getProcNoiseMat().GetNcols();
    Int_t nRowsNew = fProc.GetNrows();
    Int_t nColsNew = fProc.GetNcols();
    if(nRowsOld != nRowsNew || nColsOld != nColsNew) {
        Error("setProcNoiseMat()", Form("Matrices not compatible. Cannot replace %ix%i matrix with %ix%i matrix.", nRowsOld, nColsOld, nRowsNew, nColsNew));
        exit(1);
    }
#endif
    fProcessNoise = fProc;
}

void HKalTrackState::setStateVec(const TVectorD &sv) {
    // Replace the state's track parameter vector.

#ifdef kalDebug
    Int_t oldStateDim = getStateVec().GetNrows();
    Int_t newStateDim = sv.GetNrows();
    if(oldStateDim != newStateDim) {
        Error("setStateVec()", Form("Dimension of new state vector (%i) not the same as dimension of current state vector (%i).", newStateDim, oldStateDim));
        exit(1);
    }
#endif
    stateVec = sv;
}