// from ROOT
#include "TMath.h"
#include "TMatrixD.h"
#include "TVector3.h"
#include "TVectorD.h"

// from hydra
#include "hmdctrackgfield.h"
#include "hphysicsconstants.h"

#include "hkalplane.h"
#include "hkalsystem.h"

#include <iostream>
using namespace std;

ClassImp (HKalSystem)

//_HADES_CLASS_DESCRIPTION
///////////////////////////////////////////////////////////////////////////////
//
// Class: HKalSystem
//
// This class is responsible for the Kalman filter.
// Before running the filter for the first time, the filter options should
// be set with the function setFilterOptions().
// Possible filter options:
// - number of iterations
// - propagation direction (kIterForward, kIterBackward)
// - do smoothing
// - tilt the coordinate system to avoid large values for the track state
//   parameters tx or ty.
//
// The Kalman filter is started via the fitTrack() functions. These get a
// series of measurement points (an array of class HKalMdcHit) as input.
// The user may pass an initial track state vector and covariance matrix that
// the filter will use by using the function:
// Bool_t fitTrack(const TObjArray &hits,  const TVectorD &iniStateVec, const TMatrixD &iniCovMat)
// Alternatively the filter may also be started with only an intial momentum guess:
// Bool_t fitTrack(const TObjArray &hits, Double_t iniQP)
// It will then make an estimation for the initial track state based on the
// measurement hits and fill the covariance matrix with large values.
// The function will return a boolean to indicate whether problems (high chi^2 or
// unrealistic track parameters) were encountered during track fitting.
//
// Track propagation will be done with the Runge-Kutta method.
//
// After the Kalman filter is done, a measurement site may be eliminated
// from the calculation with the function
// excludeSite(Int_t iSite)
//
//
//-----------------------------------------------------------------------
///////////////////////////////////////////////////////////////////////////////


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

HKalSystem::HKalSystem(Int_t nHits, Int_t measDim, Int_t stateDim, Double_t chiMax)
: TObject(),
  fProj(measDim, stateDim),
  fProjTrans(stateDim, measDim) {
    // Constructor for a Kalman system.
    // nHits:    the number of hits in a track (may be changed afterwards)
    // measDim:  the dimension of a measurement vector (must be the same for all tracks)
    // stateDim: the dimension of the track state vector (must be the same for all tracks)
    // chiMax:   the maximum allowed chi^2 for which a track is accepted.
    nHitsInTrack = nHits;

    sites = new HKalTrackSite* [nHitsInTrack];
    for(Int_t i = 0; i < nHitsInTrack; i++) {
        sites[i] = new HKalTrackSite(measDim, stateDim);
    }

    setFilterOptions(1, kIterForward, kTRUE, kNoRot);
    pRotMat = new TRotation();
    trackPropagator.setRotationMatrix(pRotMat);

    calcProjector(fProj);
    fProjTrans = TMatrixD(TMatrixD::kTransposed, fProj);

    chi2High   = chiMax;
    bPrintWarn = kTRUE;

    hDafMonitor = new TNtupleD("daf", "", "weight1:weight2:chi21:chi22:temp:iter:site");
}

HKalSystem::~HKalSystem() {
    //? daf test
    delete hDafMonitor;

    delete sites;
    delete pRotMat;
}

//  -----------------------------------
//  Implementation of protected methods
//  -----------------------------------

void HKalSystem::calcIniStateVec(TVectorD &stateVec, Double_t iniQP) const {
    // Precondition: updateSites() called.
    // Make an estimation on the initial track state based on the track sites
    // stored in the Kalman system.
    // stateVec: the estimated intial track state (return value)
    // iniQP:    an initial guess for the particle charge in elementary charges
    //           divided by total momentum in MeV/c

    Int_t iStart;
    if(getDirection() == kIterForward) {
        iStart = 0;
    } else {
        iStart = getNhits() - 1;
    }
    TVector3 start;
    getSite(iStart)->getHitVec3(start);
    TVector3 dir;
    calcIniTrackDir(dir);

#ifdef kalDebug
    if(stateVec.GetNrows() < getStateDim()) {
        Warning("calcIniStateVec()", "Wrong dimension of input state vector.");
    } else {
        HKalTrackState::calcStateVec(stateVec, iniQP, start, dir);
    }
#else
    HKalTrackState::calcStateVec(stateVec, iniQP, start, dir);
#endif
}

void HKalSystem::calcIniTrackDir(TVector3 &dir) const {
    // Precondition: updateSites() called.
    // Make an estimation for the initial track direction based on the
    // track sites stored in the Kalman system.
    // The direction will be a straight line either between the first
    // two sites when iterating in forward direction or the last two
    // sites when iterating backwards. The vector will always point in
    // positive z (beam direction).
    // dir: direction vector (return value)

    Int_t idxStart, idxSecond;
    if(getDirection() == kIterForward) {
        idxStart  = 0;
        idxSecond = 1;
    } else {
        idxStart  = getNhits() - 1;
        idxSecond = getNhits() - 2;
    }
    TVector3 start;
    getSite(idxStart)->getHitVec3(start);
    TVector3 second;
    getSite(idxSecond)->getHitVec3(second);
    dir = second - start;
    if(getDirection() == kIterBackward) {
        dir *= -1.;  // direction vector should always point in positive z direction
    }
    dir = dir.Unit();
}

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

void HKalSystem::calcIniCovMat(TMatrixD &fCov) const {
    // Fill diagonal elements of parameter fCov with large values.

    for(Int_t ci = 0; ci < fCov.GetNrows(); ci++) {
        fCov(ci,ci) = 10000.;
    }
}

void HKalSystem::calcIniStateVec(TVectorD &stateVec, Double_t iniQP, const TObjArray &allhits) const {
    // Make an estimation on the initial track state based on the measurement hits
    // stored in parameter allhits.
    // stateVec: the estimated intial track state (return value)
    // iniQP:    an initial guess for the particle charge in elementary charges
    //           divided by total momentum in MeV/c
    // allhits:  array with measurement hits.

    Int_t iStart;
    if(getDirection() == kIterForward) {
        iStart = 0;
    } else {
        iStart = allhits.GetEntries() - 1;
    }

#ifdef kalDebug
    if(!allhits.At(iStart)->InheritsFrom("HKalMdcHit")) {
        Error("calcIniStateVec", Form("Object at index %i in allhits array is of class %s. Need class HKalMdcHit.", iStart, allhits.At(iStart)->ClassName()));
        exit(1);
    }
#endif

    TVector3 start;
    ((HKalMdcHit*)allhits.At(iStart))->getHitVec3(start);
    TVector3 dir;
    calcIniTrackDir(dir, allhits);

#ifdef kalDebug
    if(stateVec.GetNrows() < getStateDim()) {
        Warning("calcIniStateVec()", "Wrong dimension of input state vector.");
    } else {
        HKalTrackState::calcStateVec(stateVec, iniQP, start, dir);
    }
#else
    HKalTrackState::calcStateVec(stateVec, iniQP, start, dir);
#endif
}

void HKalSystem::calcIniTrackDir(TVector3 &dir, const TObjArray &allhits) const {
    // Make an estimation for the initial track direction based on the
    // measurement hits stored in the parameter allhits.
    // The direction will be a straight line either between the first
    // two hits when iterating in forward direction or the last two
    // hits when iterating backwards. The vector will always point in
    // positive z (beam direction).
    // dir:     direction vector (return value)
    // allhits: array with measurement hits.

    Int_t idxStart, idxSecond;
    if(getDirection() == kIterForward) {
        idxStart  = 0;
        idxSecond = 1;

        Bool_t bIsCompetitor = kTRUE;
#ifdef kalDebug
        if(!allhits.At(idxStart)->InheritsFrom("HKalMdcHit")) {
            Error("calcIniTrackDir", Form("Object at index %i allhits array is of class %s. Need class HKalMdcHit.", idxStart, allhits.At(idxStart)->ClassName()));
            exit(1);
        }
#endif
        HKalMdcHit *hitstart = (HKalMdcHit*)allhits.At(idxStart);

        while(bIsCompetitor && idxSecond < allhits.GetEntries()) {

#ifdef kalDebug
            if(!allhits.At(idxSecond)->InheritsFrom("HKalMdcHit")) {
                Error("calcIniTrackDir", Form("Object at index %i allhits array is of class %s. Need class HKalMdcHit.", idxSecond, allhits.At(idxSecond)->ClassName()));
                exit(1);
            }
#endif
            HKalMdcHit *hitsecond = (HKalMdcHit*)allhits.At(idxSecond);

            // If next hit is a competitor then go to next site.
            if(hitstart->areCompetitors(*hitstart, *hitsecond)) {
                idxSecond++;
            } else {
                bIsCompetitor = kFALSE;
            }
        }
    } else {
        idxStart  = getNhits() - 1;
        idxSecond = getNhits() - 2;

        Bool_t bIsCompetitor = kTRUE;
#ifdef kalDebug
        if(!allhits.At(idxStart)->InheritsFrom("HKalMdcHit")) {
            Error("calcIniTrackDir", Form("Object at index %i allhits array is of class %s. Need class HKalMdcHit.", idxStart, allhits.At(idxStart)->ClassName()));
            exit(1);
        }
#endif
        HKalMdcHit *hitstart = (HKalMdcHit*)allhits.At(idxStart);

        while(bIsCompetitor && idxSecond >= 0) {
#ifdef kalDebug
            if(!allhits.At(idxSecond)->InheritsFrom("HKalMdcHit")) {
                Error("calcIniTrackDir", Form("Object at index %i allhits array is of class %s. Need class HKalMdcHit.", idxSecond, allhits.At(idxSecond)->ClassName()));
                exit(1);
            }
#endif
            HKalMdcHit *hitsecond = (HKalMdcHit*)allhits.At(idxSecond);

            // If next hit is a competitor then go to next site.
            if(hitstart->areCompetitors(*hitstart, *hitsecond)) {
                idxSecond--;
            } else {
                bIsCompetitor = kFALSE;
            }
        }
    }

    TVector3 start;
    ((HKalMdcHit*)allhits.At(idxStart))->getHitVec3(start);
    TVector3 second;
    ((HKalMdcHit*)allhits.At(idxSecond))->getHitVec3(second);
    dir = second - start;
    if(getDirection() == kIterBackward) {
        dir *= -1.;  // direction vector should always point in positive z direction
    }
    dir = dir.Unit();
}

void HKalSystem::calcProjector(TMatrixD &fProj) {
    // Returns the values for the Projector Matrix, i.e. the derivation of the
    // position vector (with coordinates x,y,z) by the track state parameters
    // q/p, x0, y0, tx, ty where
    // q/p: particle charge divided by momentum
    // (x0,y0) are coordinates on a reference planes parallel to the xy-plane.
    // tx = dx/dz, ty = dy/dz
    // @x/@x0  ||  @x/@y0 ||  @x/@tx  ||  @x/@ty  || @x/@qp
    // @y/@x0  ||  @y/@y0 ||  @y/@tx  ||  @y/@ty  || @y/@qp
    // =
    //    1    ||    0    ||    0     ||    0     ||   0
    //    0    ||    1    ||    0     ||    0     ||   0
    // fProj: the projector matrix (return value)

#ifdef kalDebug
    if(fProj.GetNrows() != getMeasDim() || fProj.GetNcols() != getStateDim()) {
        Error("calcProjector()", Form("Input is a %ix%i matrix. Projector is a %ix%i matrix.",
                fProj.GetNrows(), fProj.GetNcols(), getMeasDim(), getStateDim()));
        exit(1);
    }
#endif
    fProj.UnitMatrix();
}

Int_t HKalSystem::getNDF() const {
    // Returns number degrees of freedom.
    return (getNhits() * getMeasDim() - getStateDim()); //? subtract 1?
}

Bool_t HKalSystem::fitTrack(const TObjArray &hits, Double_t iniQP, Int_t pid) {
    // Runs the Kalman filter over a set of measurement hits.
    // Returns true if the Kalman filter was successful.
    // Returns false if the Kalman filter encountered problems during
    // track fitting due to high chi^2 or unrealistically large track parameters.
    // hits:  the array with the measurement hits. Elements must be of class
    //        HKalMdcHit.
    // iniQP: initial guess for particle charge divided by momentum

    TVectorD iniStateVec(getStateDim());
    calcIniStateVec(iniStateVec, iniQP, hits);
    TMatrixD iniCovMat(getStateDim(), getStateDim());
    calcIniCovMat(iniCovMat);
    return fitTrack(hits, iniStateVec, iniCovMat, pid);
}

Bool_t HKalSystem::fitTrack(const TObjArray &hits, const TVectorD &iniStateVec, const TMatrixD &iniCovMat, Int_t pId) {
    // Runs the Kalman filter over a set of measurement hits.
    // Returns true if the Kalman filter was successful.
    // Returns false if the Kalman filter encountered problems during
    // track fitting due to high chi^2 or unrealistically large track parameters.
    // hits: the array with the measurement hits. Elements must be of class
    //       HKalMdcHit.
    // iniStateVec: initial track state that the filter will use
    // iniCovMat: initial covariance matrix that the filter will use
    // iniQP: initial guess for particle charge divided by momentum.

    bTrackAccepted = kTRUE;
    chi2 = 0.;
    updateSites(hits);
    pid = pId;

    // Depending on the direction for propagation, set the index of the first site (istart)
    // and last site (igoal).
    // idir: increment (forward propagation) or decrement (backward propagation) sites indices.
    Int_t fromSite, toSite;
    if(getDirection() == kIterForward) {
        fromSite = 0;
        toSite   = getNhits() - 1;
    } else {
        fromSite = getNhits() - 1;
        toSite   = 0;
    }

    // Set the state vector and covariance matrix for the first site.
    getSite(fromSite)->setStateVec   (kPredicted, iniStateVec);
    getSite(fromSite)->setStateVec   (kFiltered,  iniStateVec);
    getSite(fromSite)->setStateCovMat(kPredicted, iniCovMat);
    getSite(fromSite)->setStateCovMat(kFiltered , iniCovMat);

    //cout<<"predicted state first site: "<<endl;
    //getSite(fromSite)->getStateVec(kPredicted).Print();
    //getSite(fromSite)->getStateVec(kFiltered).Print();

    // Tilt coordinate system.
    // There are three options:
    // 1. kNoRot:
    //    Don't tilt the coordinate system.
    // 2. kVarRot:
    //    Tilt the coordinate system so that the z-axis will always
    //    point in the initial track direction. Assuming track curvature
    //    is not large this ensures that the track state parameters
    //    tx = tan(p_x / p_z) and ty = tan(p_y / p_z) remain small.
    //    Angles close to 90° would result in large values for tx and ty
    //    due to the tangens and poor results for the Kalman filter.
    // 3. kFixedRot:
    //    Always use the same rotation matrix for all tracks. The rotation
    //    matrix has to be set by the user before starting track fitting
    //    via the setRotationAngles() function.
    //    It is recommended to rotate by 45° around the x-Axis and 0°
    //    around the y-Axis.
    if(rotCoords == kVarRot) {
        // Calculate rotation matrix depending on track direction.
        TVector3 dir;
        calcIniTrackDir(dir);
        Double_t rotXangle = dir.Theta();
        Double_t rotYangle = - TMath::Pi()/2 + dir.Phi();
        setRotationAngles(rotXangle, rotYangle);
    }
    if(rotCoords == kVarRot || rotCoords == kFixedRot) {
        // Rotate coordinates of all sites.
        transformFromSectorToTrack();
    }

    // Propagation dir may change when doing several Kalman Filter iterations.
    Bool_t orgDirection = getDirection();

    // At least one site has competing hits. Run the Deterministic Annealing Filter.
    // For explanation of DAF, see for example:
    //
    // R. Fruehwirth, A. Strandlie, Comput. Phys. Commun. 120 (1999) 197.
    // Track fitting with ambiguities and noise:
    // a study of elastic tracking and nonlinear filters
    //
    // Sebastian Fleischmann
    // Track Reconstruction in the ATLAS Experiment - the Deterministic Annealing Filter
    if(bDoDaf) {
        const Int_t nDafs   = 5;                        // Number of DAF iterations.
        Double_t    T[]     = { 81., 9., 4., 1., 1. };  // Annealing factors (temperatures).
        Double_t    chi2Cut = 35.;                      // Cut-off parameter for DAF.

        for(Int_t iDaf = 0; iDaf < nDafs; iDaf++) {
            //cout<<"\n**** Daf run number "<< iDaf<<" **** \n"<<endl;

            // In each iteration do Kalman prediction/filter steps and then smooth back.
            predictAndFilter(fromSite, toSite);
            smoothAll();
            //getSite(fromSite)->setStateVec   (kPredicted, getSite(fromSite)->getStateVec   (kSmoothed));
            //getSite(fromSite)->setStateCovMat(kPredicted, getSite(fromSite)->getStateCovMat(kSmoothed));

            for(Int_t iSite = 0; iSite < getNhits(); iSite++) {
                HKalTrackSite *site = getSite(iSite);
                Double_t weightNorm = 0.;

                for(Int_t iHit = 0; iHit < site->getNcompetitors(); iHit++) {
                    const TVectorD &measVec = site->getHitVec(iHit);
                    TVectorD smooVec(getMeasDim());
                    site->calcMeasVecFromState(smooVec, kSmoothed);

                    //cout<<"meas vec: "<<endl;
                    //measVec.Print();
                    //cout<<"smoothed state: "<<endl;
                    //smooVec.Print();
                    //cout<<"predicted state of site "<<iSite<<": "<<endl;
                    //site->getStateVec(kPredicted).Print();

                    TVectorD  residual     = measVec - smooVec;                          // Residual r
                    //cout<<"residual: "<<endl;
                    //residual.Print();
                    TMatrixD  residualTrans(1, getMeasDim(), residual.GetMatrixArray()); // r^T
                    const TMatrixD &errMat = site->getErrMat(iHit);                      // Error matrix V
                    TMatrixD  invErrMat    = TMatrixD(TMatrixD::kInverted, errMat);      // G = V^-1
                    //invErrMat.Print();

                    //
                    // Normalization of hit's chi2 distribution: 1 / ((2*pi)^n/2 * sqrt(T * det(V)))
                    Double_t  normChi2     = 1. / (TMath::Power(2*TMath::Pi(), getMeasDim()/2.) * TMath::Sqrt(T[iDaf] * errMat.Determinant()));

                    // Standardized distance between smoothed track state and measurement vector:
                    //                     =  r^T           * G         * r
                    Double_t  chi2         = (residualTrans * invErrMat * residual)(0);
                    //cout<<"site "<<iSite<<", hit "<<iHit<<" has chi2: "<<chi2<<", norm chi2: "<<normChi2<<endl;
                    site->setHitChi2(chi2, iHit);

                    // Normalization for hit's weight recalculation:
                    // sum( normchi2 * exp(-chi2/(2T)) + normchi2 * exp(-chi2cut/(2T))
                    weightNorm            += normChi2 * TMath::Exp(-chi2    / (2. * T[iDaf]));
                    weightNorm            += normChi2 * TMath::Exp(-chi2Cut / (2. * T[iDaf]));
                }

                // Recalculate weight of each hit.
                for(Int_t iHit = 0; iHit < site->getNcompetitors(); iHit++) {
                    const TMatrixD &errMat = site->getErrMat(iHit);
                    // New weight of hit.
                    // 1 / ((2*pi)^n/2 * sqrt(T * det(V))) * exp(-chi2/(2T)) / weightNorm
                    Double_t normChi2      = 1. / (TMath::Power(2*TMath::Pi(), getMeasDim()/2.) * TMath::Sqrt(T[iDaf] * errMat.Determinant()));
                    Double_t hitWeight     = normChi2 * TMath::Exp(-site->getHitChi2(iHit) / (2*T[iDaf])) / weightNorm;
                    //cout<<"updating weight for site "<<iSite<<", hit "<<iHit<<", old weight "<<site->getHitWeight(iHit)<<", new weight "<<hitWeight<<endl;
                    site->setHitWeight(hitWeight, iHit);
                }
            }
            //? print daf results

            for(Int_t iSite = 0; iSite < getNhits(); iSite++) {
                HKalTrackSite *site = getSite(iSite);
                Double_t arr[] = {
                        site->getHitWeight(0), site->getHitWeight(1), site->getHitChi2(0), site->getHitChi2(1), T[iDaf], iDaf, iSite
                };
                hDafMonitor->Fill(arr);
            }
        }
        for(Int_t iSite = 0; iSite < getNhits(); iSite++) {
            //getSite(iSite)->sortHits();
        }
        // Do standard Kalman Filter.
    } else {
        // Run the prediction and filter steps.
        for(Int_t i = 0; i < nKalRuns; i++) {
            bTrackAccepted &= predictAndFilter(fromSite, toSite);

            /*
            // Filter then smooth in each iteration.
            smoothAll();
            //getSite(fromSite)->setStateVec(   kPredicted, getSite(fromSite)->getStateVec(kSmoothed));
            getSite(fromSite)->setStateCovMat(kPredicted, getSite(fromSite)->getStateCovMat(kSmoothed));
            TVectorD filtStateVecFrom = getSite(fromSite)->getStateVec(kFiltered);
            filtStateVecFrom(kIdxQP) = getSite(fromSite)->getStateParam(kSmoothed, kIdxQP);
            getSite(fromSite)->setStateVec(kPredicted, filtStateVecFrom);
            getSite(fromSite)->setStateVec(kFiltered, filtStateVecFrom);
             */

            /*
            if(nKalRuns > 1) {
                // After the first iteration, replace the qp value of the
                // first (last) site when iterating forward (backward) with
                // the qp value that the Kalman filter has calculated in the
                // previous iteration.
                Double_t filtQPLast = getSite(toSite)->getStateParam(kFiltered, kIdxQP);

                if(bDoDeDx) {
                    // Correct new input momentum for energy loss.
                    // E' = particle energy including energy loss due to radiation and ionization.
                    // E  = particle energy without energy loss.
                    // E' = E + delta(E), delta(E) is negative for forward iteration
                    // E  = E' - delta(E)
                    //    = E' * (1 - delta(E)/E')
                    // qp = qp' / (1 - delta(E)*qp) for e+/e-
                    Double_t dE     = getSite(toSite)->getEnergyLoss();

                    if(pid == 2 || pid == 3) {
                        filtQPLast = filtQPLast / (1. -  dE * filtQPLast);
                    } else {
                        Double_t mass   = HPhysicsConstants::mass(pid);
                        Double_t ePrime = TMath::Sqrt(mass*mass + 1./(filtQPLast*filtQPLast));
                        Double_t E      = ePrime * (1. - dE/ePrime);
                        filtQPLast = 1./TMath::Sqrt(E*E - mass*mass);
                    }
                }
                TVectorD filtStateVecFrom = getSite(fromSite)->getStateVec(kFiltered);
                filtStateVecFrom(kIdxQP) = filtQPLast;
                getSite(fromSite)->setStateVec(kPredicted, filtStateVecFrom);
                getSite(fromSite)->setStateVec(kFiltered, filtStateVecFrom);
            }
             */

            if(i < nKalRuns) { // At least one more iteration to do.
                // Set the filtered state of the last site from the previous iteration as the
                // prediction for the next iteration.
                getSite(toSite)->setStateVec(kPredicted, getSite(toSite)->getStateVec(kFiltered));
                //getSite(toSite)->setStateCovMat(kPredicted , getSite(toSite)->getStateCovMat(kFiltered));

                // Change propagation direction in each iteration.
                direction   = !direction;
                trackPropagator.setDirection(direction);
                Int_t oldTo = toSite;
                toSite      = fromSite;
                fromSite    = oldTo;
            }
        }
        // Do smoothing.
        if(bDoSmooth) {
            smoothAll();
        }
    }

    // Restore direction passed by user.
    direction = orgDirection;
    trackPropagator.setDirection(direction);

    // Transform hit and state vectors from track back to sector coordinates.
    if(rotCoords == Kalman::kVarRot || rotCoords == Kalman::kFixedRot) {
        transformFromTrackToSector();
    }

    // Check for large chi^2 values.
    if(getChi2() > chi2High) {
        bTrackAccepted = kFALSE;
        if(bPrintWarn) {
            Warning("", "High chi^2.");
        }
    }

    return bTrackAccepted;
}

Bool_t HKalSystem::predictAndFilter(Int_t fromSite, Int_t toSite) {
    // Do the prediction and filter steps of the Kalman filter by propagating
    // the track between sites.
    // fromSite: index of track site to start propagation from
    // toSite:   index of track site to propagate to.

    Bool_t propagationOk = kTRUE;
    Int_t iDir;
    if(fromSite < toSite) { // propagate in forward direction
        iDir = +1;
    } else {                // propagate in backward direction
        iDir = -1;
    }
    Double_t energyLoss = 0.;
    for(Int_t iCurSite = fromSite; (iDir > 0 && iCurSite < toSite) || (iDir < 0 && iCurSite > toSite); iCurSite += iDir) {
        // iCurSite is index of current site k-1
        Int_t iNextSite = iCurSite + iDir; // index of next site k

        //cout<<"############## Prediction from site "<<iCurSite<<" to site "<<iNextSite<<" ############"<<endl;

        HKalTrackSite    *curSite              = getSite(iCurSite);
        HKalTrackSite    *nextSite             = getSite(iNextSite);
        const TVectorD   &filtStateCurSite     = curSite ->getStateVec(kFiltered);          // a_(k-1): filtered state vector of current site
        TVectorD          predStateNextSite    = nextSite->getStateVec(kPredicted);         // a^(k-1)_k: predicted state for next site, to be calculated
        const HKalMdcMeasLayer &planeCurSite   = curSite ->getMeasLayer();                  // measurement layer of current site
        const HKalMdcMeasLayer &planeNextSite  = nextSite->getMeasLayer();                  // measurement layer of next site
        TMatrixD          propMatCurSite       = curSite ->getStatePropMat(kFiltered);      // F_k-1: propagation matrix, to be calculated
        TMatrixD          procMatCurSite       = curSite ->getStateProcNoiseMat(kFiltered); // Q_k-1: process noise matrix, to be calculated

#ifdef kalDebug
        Int_t nRows, nCols;
        nRows = filtStateCurSite.GetNrows();
        if(nRows != getStateDim()) {
            Error(Form("Prediction to site %i.", iNextSite), Form("State vector must have dimension %i, but has dimension %i", nRows, getStateDim()));
            exit(1);
        }
        nRows = predStateNextSite.GetNrows();
        if(nRows != getStateDim()) {
            Error(Form("Prediction to site %i.", iNextSite), Form("State vector must have dimension %i, but has dimension %i", predStateNextSite.GetNrows(), getStateDim()));
            exit(1);
        }
        nRows = propMatCurSite.GetNrows();
        nCols = propMatCurSite.GetNcols();
        if(nRows != getStateDim() || nCols != getStateDim()) {
            Error(Form("Prediction to site %i.", iNextSite), Form("Propagator matrix is %ix%i, but must be %ix%i matrix.", nRows, nCols, getStateDim(), getStateDim() ));
            exit(1);
        }
        nRows = procMatCurSite.GetNrows();
        nCols = procMatCurSite.GetNcols();
        if(nRows != getStateDim() || nCols != getStateDim()) {
            Error(Form("Prediction to site %i.", iNextSite), Form("Process noise matrix is %ix%i, but must be %ix%i matrix.", nRows, nCols, getStateDim(), getStateDim() ));
            exit(1);
        }
#endif

        // Propagate track to next site. Calculate f(a_k-1) and F_k-1.
        propagationOk &= trackPropagator.propagateToPlane(propMatCurSite, procMatCurSite,
                predStateNextSite, filtStateCurSite,
                planeCurSite, planeNextSite, pid, getDirection());
        nextSite->setStateVec(kPredicted, predStateNextSite);  // a^(k-1)_k = f(a_k-1), store prediction for next site
        curSite ->setStatePropMat(kFiltered, propMatCurSite);  // F_k-1, store propagator matrix
        energyLoss += trackPropagator.getEnergyLoss();
        nextSite->setEnergyLoss(energyLoss);

        // Calculate covariance matrix for prediction of next site.
        TMatrixD covMatCurSite       = curSite->getStateCovMat(kFiltered);              // C_k-1
        TMatrixD propMatCurSiteTrans = TMatrixD(TMatrixD::kTransposed, propMatCurSite); // F_k-1

#ifdef kalDebug
        nRows = covMatCurSite.GetNrows();
        nCols = covMatCurSite.GetNcols();
        if(nRows != getStateDim() || nCols != getStateDim()) {
            Error(Form("Prediction to site %i.", iNextSite), Form("Covariance matrix is %ix%i, but must be %ix%i matrix.", nRows, nCols, getStateDim(), getStateDim() ));
            exit(1);
        }
#endif
        //        C^k-1_k        = F_k-1          * C_k-1         * F^T_k-1             + Q_k-1
        TMatrixD  covMatNextSite = propMatCurSite * covMatCurSite * propMatCurSiteTrans + procMatCurSite;
        nextSite->setStateCovMat(kPredicted, covMatNextSite);

        // Filter next site.
        filter(iNextSite);
    }
    return propagationOk;
}

void HKalSystem::filter(Int_t iSite) {
    // Do the filter step in Kalman filter for a site. The prediction must have been done
    // for the site before calling this.
    // iSite: index of site to filter

    //cout<<"############## Filter site "<<iSite<<"#####################"<<endl;
    HKalTrackSite *site  = getSite(iSite);

    // Inverted covariance matrix from the prediction.
    TMatrixD predCovInv = TMatrixD(TMatrixD::kInverted, site->getStateCovMat(kPredicted)); // (C^k-1_k)^-1
    // Inverted measurement error matrix.
    TMatrixD errInv     = TMatrixD(TMatrixD::kInverted, site->getErrMat());                // G_k
    // Predicted state vector for this site.
    TVectorD predState  = site->getStateVec(kPredicted);                                   // a^k-1_k
    // Measurement vector.
    TVectorD measVec    = site->getHitVec();                                               // m_k

    Double_t hitsTotalWeight = 1.;
    if(bDoDaf) {
        hitsTotalWeight = site->getHitsTotalWeight();
    }

#ifdef kalDebug
    Int_t nRows, nCols;
    nRows = predState.GetNrows();
    if(nRows != getStateDim()) {
        Error(Form("Filter site %i.", iSite), Form("State vector must have dimension %i, but has dimension %i", nRows, getStateDim()));
        exit(1);
    }

    nRows = measVec.GetNrows();
    if(nRows != getMeasDim()) {
        Error(Form("Filter site %i.", iSite), Form("Measurement vector must have dimension %i, but has dimension %i", nRows, getMeasDim()));
        exit(1);
    }

    nRows = predCovInv.GetNrows();
    nCols = predCovInv.GetNcols();
    if(nRows != getStateDim() || nCols != getStateDim()) {
        Error(Form("Filter site %i.", iSite), Form("Inverse covariance matrix is %ix%i, but must be %ix%i matrix.", nRows, nCols, getStateDim(), getStateDim() ));
        exit(1);
    }

    nRows = errInv.GetNrows();
    nCols = errInv.GetNcols();
    if(nRows != getMeasDim() || nCols != getMeasDim()) {
        Error(Form("Filter site %i.", iSite), Form("Inverse error matrix is %ix%i, but must be %ix%i matrix.", nRows, nCols, getMeasDim(), getMeasDim() ));
        exit(1);
    }
#endif

    // Calculate covariance matrix for filtered state.
    //       C_k        = [                           (C^k-1_k)^-1 + H^T_k      * G_k    * H_k  ]^-1
    TMatrixD filtCov    = TMatrixD(TMatrixD::kInverted, predCovInv + hitsTotalWeight * fProjTrans * errInv * fProj);
    site->setStateCovMat(kFiltered, filtCov);

    // Kalman gain matrix.
    //       K_k        = C_k     * H^T_k      * G_k
    TMatrixD kalmanGain = filtCov * fProjTrans * errInv;

    // Calculate a measurement vector from the predicted state.
    TVectorD projMeasVec(getMeasDim());
    site->calcMeasVecFromState(projMeasVec, kPredicted);    // h_k(a^k-1_k)
    TVectorD filtState(getStateDim());
    TVectorD residual(getMeasDim());

    if(bDoDaf) {
        // Go through all competitors of a site. The residual is the weighted difference
        // between measurement vector and the projected measurement from the
        // predicted state vector.
        for(Int_t i = 0; i < site->getNcompetitors(); i++) {
            measVec   = site->getHitVec(i);
            residual += site->getHitWeight(i) * (measVec - projMeasVec);
        }
    } else {
        measVec  = site->getHitVec();
        residual = measVec - projMeasVec;
    }

    // Calculate filtered state vector.
    // a_k      = a^k-1_k   + K_k        * (m_k     - h_k(a^k-1_k) )
    filtState = predState + kalmanGain * residual;

    site->setStateVec(kFiltered, filtState);

    // Update chi^2.
    TVectorD projFiltMeasVec(getMeasDim());
    site->calcMeasVecFromState(projFiltMeasVec, kFiltered); // h_k(a_k)
    // Transform TVectorD objects into matrices for matrix multiplications.
    TMatrixD diffStateTrans(1, getStateDim(), (filtState - predState).GetMatrixArray()); // (a_k - a^k-1_k )^T
    TMatrixD resTrans      (1, getMeasDim(),  (measVec - projFiltMeasVec).GetMatrixArray());    // (a_k - a^k-1_k )^T. Transposed residual.

    //        (a_k - a^k-1_k  )^T * (C^k-1_k)^-1 * (a_k       - a^k-1_k  )
    chi2 +=   (diffStateTrans     * predCovInv   * (filtState - predState)    )(0)
	                //        (m_k - h_k(a_k))^T  * G            * (m_k     - h_k(a_k)       )
			                + (resTrans           * errInv       * (measVec - projFiltMeasVec))(0);
}

void HKalSystem::smoothAll() {
    // Do smoothing for all sites. Requires that the prediction and filter steps have
    // been done before.

    Int_t iStartSite, iGoalSite, iDir;
    // For smoothing, iterate through the sites in opposite direction as in the
    // prediction/filter steps.
    if(getDirection() == kIterForward) {
        iStartSite = getNhits() - 1;
        iGoalSite  = 0;
        iDir       = -1;
    } else {
        iStartSite = 0;
        iGoalSite  = getNhits() - 1;
        iDir       = +1;
    }

    // The result of filtering for the last site coincides with that of smoothing.
    HKalTrackSite *cur = getSite(iStartSite);
    cur->setStateVec   (kSmoothed,    cur->getStateVec(kFiltered));
    cur->setStateCovMat(kSmoothed,    cur->getStateCovMat(kFiltered));
    cur->setStateCovMat(kInvFiltered, cur->getStateCovMat(kFiltered));
    iStartSite += iDir;

    for(Int_t iSite = iStartSite; (iDir > 0 && iSite <= iGoalSite) || (iDir < 0 && iSite >= iGoalSite); iSite += iDir) {
        //cout<<"Smoothing site "<<iSite<<endl;
        HKalTrackSite* pre     = getSite(iSite - iDir);     // Already smoothed site k+1.
        HKalTrackSite* cur     = getSite(iSite);            // Site k to be smoothed.

        // Filtered state vector of site k+1.
        const TVectorD &prePredState  = pre->getStateVec    (kPredicted);          // a^k_k+1
        // Smoothed state vector of site k+1.
        const TVectorD &preSmooState  = pre->getStateVec    (kSmoothed);           // a^n_k+1
        // Filtered state vector of site k.
        const TVectorD &curFiltState  = cur->getStateVec    (kFiltered);           // a_k
        // Covariance from filter step for site k.
        const TMatrixD &curFiltCov    = cur->getStateCovMat (kFiltered);           // C_k
        // Covariance and inverted covariance from prediction for site k+1.
        const TMatrixD &prePredCov    = pre->getStateCovMat (kPredicted);          // C^k_k+1
        TMatrixD  prePredCovInv = TMatrixD(TMatrixD::kInverted, prePredCov); // (C^k_k-1)^-1
        // Propagator and transposed propagator from predicion for site k.
        const TMatrixD &curProp       = cur->getStatePropMat(kFiltered);           // F_k
        TMatrixD  curPropTrans  = TMatrixD(TMatrixD::kTransposed, curProp);  // F^T_k
        // Covariance from smoothing step for site k+1.
        const TMatrixD &preSmooCov    = pre->getStateCovMat (kSmoothed);           // C^n_k+1

        // Gain matrix for smoothing.
        //       A_k          = C_k        * F^T_k        * (C^k_k+1)^-1
        TMatrixD smootherGain = curFiltCov * curPropTrans * prePredCovInv;

        // Calculate smoothed state vector.
        //       a^n_k        = a_k          + A_k          * (a^n_k+1 - a^k_k+1)
        TVectorD curSmooState = curFiltState + smootherGain * (preSmooState - prePredState);
        cur->setStateVec(kSmoothed,    curSmooState);
        cur->setStateVec(kInvFiltered, curSmooState); // Inverse filtering works on smoothed states.

        // Calculate covariance for smoothed state. Needed for possible inverse filtering.
        //       C^n_k      = C_k        + A_k          * (C^n_k-1    - C^k_k+1   ) * A^T_k
        TMatrixD curSmooCov = curFiltCov + smootherGain * (preSmooCov - prePredCov) * TMatrixD(TMatrixD::kTransposed, smootherGain);
        cur->setStateCovMat(kSmoothed,    curSmooCov);
        cur->setStateCovMat(kInvFiltered, curSmooCov);
    }
}

void HKalSystem::excludeSite(Int_t iSite) {
    // Exclude a site from the measurement. Requires that prediction, filtering and smoothing
    // has been done before. If smoothing has been disabled it will be done in this function.
    // iSite: index in sites array of the site to exclude

    if(iSite < 0 || iSite > getNhits() - 1) {
        Warning("excludeSite()", "Invalid site index.");
        return;
    }

    // Transform state vector from sector to track coordinates.
    if(rotCoords == Kalman::kVarRot || rotCoords == Kalman::kFixedRot) {
        transformFromSectorToTrack();
    }

    // Do smoothing including site k if it hasn't been done before.
    if(!bDoSmooth) {
        smoothAll();
    }

    HKalTrackSite* exSite = getSite(iSite);                           // Site k to exclude.
    // Smoothed state vector of site to exclude.
    TVectorD exSmooState  = exSite->getStateVec(kSmoothed);           // a^n_k
    // Covariance and inverse covariance from smoothing step.
    TMatrixD exSmooCov    = exSite->getStateCovMat(kSmoothed);        // C^n_k
    TMatrixD exSmooCovInv = TMatrixD(TMatrixD::kInverted, exSmooCov); // (C^n_k)^-1
    // Measurement and inverse measurement matrices.
    TMatrixD exErr        = exSite->getErrMat();                      // V_k
    TMatrixD exErrInv     = TMatrixD(TMatrixD::kInverted, exErr);     // G_k
    // Measurement vector.
    TVectorD exMeasVec    = exSite->getHitVec();                      // m_k
    // Calculate a measurement vector from smoothed state vector.
    TVectorD exSmooProjMeasVec(getMeasDim());
    exSite->calcMeasVecFromState(exSmooProjMeasVec, kSmoothed);       // h_k(a^n_k)

    //                             - V_k   + H_k   * C^n_k     * H_k
    TMatrixD kalmanGainHelper = -1.* exErr + fProj * exSmooCov * fProj;
    //       K^n*_k     = C^n_k     * H^T_k      * (                             - V_k   + H_k   * C^n_k     * H_k  )^-1
    TMatrixD kalmanGain = exSmooCov * fProjTrans * TMatrixD(TMatrixD::kInverted, kalmanGainHelper);

    // Inverse filter: calculate new state vector at site k excluding measurement at this site.
    //       a^n*_k         = a^n_k       + K^n*_k     * (m_k       - h_k(a^n_k)   )
    TVectorD exInvFiltState = exSmooState + kalmanGain * (exMeasVec - exSmooProjMeasVec);
    exSite->setStateVec(kInvFiltered, exInvFiltState);

    // Calculate new covariance at site k.
    //       C^n*_k       =         (                     (C^n_k)^-1   - H^T_k      * G_k      * H_k  )^-1
    TMatrixD exInvFiltCov = TMatrixD(TMatrixD::kInverted, exSmooCovInv - fProjTrans * exErrInv * fProj);
    exSite->setStateCovMat(kInvFiltered, exInvFiltCov);

    // Update chi^2.
    TVectorD projMeasVecInvFilt(getMeasDim());
    exSite->calcMeasVecFromState(projMeasVecInvFilt, kInvFiltered);

    // Transform TVectorD objects into column vectors for matrix multiplication.
    TMatrixD diffStateTrans(1, getStateDim(), (exInvFiltState - exSmooState).GetMatrixArray()); // (a^n*_k - a^n_k)^T
    TMatrixD resTrans(1, getMeasDim(), (exMeasVec - projMeasVecInvFilt).GetMatrixArray()); // (m_k - h_k(a^n*_k). Transposed residual.

    //        (a^n*_k - a^n_k)^T    * (C^n_k)^-1 * (a^n*_k         - a^n_k      )
    chi2 -=   (diffStateTrans       * exSmooCov  * (exInvFiltState - exSmooState)  )(0)
	                //        (m_k - h_k(a^n*_k))^T * G          * (m_k       - h_k(a^n*_k)       )
			                - (resTrans             * exErrInv   * (exMeasVec - projMeasVecInvFilt))(0);

    // Transform state vectors back from track to sector coordinates.
    if(rotCoords == Kalman::kVarRot || rotCoords == Kalman::kFixedRot) {
        transformFromTrackToSector();
    }
}

void HKalSystem::transformFromSectorToTrack() {
    // Transform all sites using the Kalman system's rotation matrix.
#ifdef kalDebug
    if(pRotMat->IsIdentity()) {
        Warning("transformFromSectorToTrack()", "Transformation matrix has not been set.");
    }
#endif
    for(Int_t i = 0; i < getNhits(); i++) {
        getSite(i)->transform(pRotMat);
    }
}

void HKalSystem::transformFromTrackToSector() {
    // Transform all sites using the inverse rotation matrix stored in the Kalman system.
    // This will transform the sites back to sector coordinates after a call of
    // transformFromSectorToTrack() and if the rotation matrix has not been modified
    // in the meantime.
    TRotation *rotInv = new TRotation(pRotMat->Inverse());
    for(Int_t i = 0; i < getNhits(); i++) {
        getSite(i)->transform(rotInv);
    }
}

void HKalSystem::updateSites(const TObjArray &hits) {
    // Replace sites of the Kalman system with new hits.
    // hits: array with new measurement hits.

    for(Int_t i = 0; i < getNhits(); i++) {
        getSite(i)->clearHits();
    }

    bDoDaf = kFALSE;

    Int_t iHit  = 0;
    Int_t iSite = 0;
    while(iHit < hits.GetEntries() - 1) {
#ifdef kalDebug
        if(!hits.At(iHit)->InheritsFrom("HKalMdcHit")) {
            Error("updateSites()", Form("Object at index %i in hits array is of class %s. Expected class is HKalMdcHit.", iHit, hits.At(iHit)->ClassName()));
            exit(1);
        }
#endif

        HKalMdcHit *hit     = (HKalMdcHit*)hits.At(iHit);
        HKalMdcHit *nexthit = (HKalMdcHit*)hits.At(iHit + 1);

        //cout<<"adding hit "<<iHit<<" to site "<<iSite<<endl;
        //hit->print("hit");

        getSite(iSite)->addHit(hit);

        // nexthit is not a competitor. Go to next site.
        if(hit->areCompetitors(*hit, *nexthit)) {
            bDoDaf = kTRUE;
        } else {
            iSite++;
        }
        iHit++;
    }
    // Add last hit in array to a site.
    HKalMdcHit* lasthit = (HKalMdcHit*)hits.At(hits.GetEntries()-1);
    getSite(iSite)->addHit(lasthit);

    //cout<<"adding hit "<<hits.GetEntries()-1<<" to site "<<iSite<<endl;
    //lasthit->print("hit");


    //cout<<"updated sites"<<endl;
    //for(Int_t i = 0; i < getNhits(); i++) {
    //    getSite(i)->print("hit");
    //}
}

void HKalSystem::setFilterOptions(Int_t kalRuns, Bool_t dir, Bool_t smooth, kalRotateOptions rotate) {
    // Set options for the Kalman filter.
    // kalRuns: number of iterations
    // dir:    propagation direction (kIterForward, kIterBackward)
    // smooth: smoothing/no smoothing
    // kalRotateOptions: Tilt coordinate system.
    // There are three options:
    // 1. kNoRot:
    //    Don't tilt the coordinate system.
    // 2. kVarRot:
    //    Tilt the coordinate system so that the z-axis will always
    //    point in the initial track direction. Assuming track curvature
    //    is not large this ensures that the track state parameters
    //    tx = tan(p_x / p_z) and ty = tan(p_y / p_z) remain small.
    //    Angles close to 90° would result in large values for tx and ty
    //    due to the tangens and poor results for the Kalman filter.
    // 3. kFixedRot:
    //    Always use the same rotation matrix for all tracks. The rotation
    //    matrix has to be set by the user before starting track fitting
    //    via the setRotationAngles() function.
    //    It is recommended to rotate by 45° around the x-Axis and 0°
    //    around the y-Axis.

    nKalRuns  = kalRuns;
    direction = dir;
    bDoSmooth = smooth;
    rotCoords = rotate;
    if(rotCoords == kFixedRot || rotCoords == kVarRot) {
        trackPropagator.setRotateBfieldVecs(kTRUE);
    } else {
        trackPropagator.setRotateBfieldVecs(kFALSE);
    }
}

void HKalSystem::setPrintWarnings(Bool_t print) {
    // Print or suppress warning messages.
    bPrintWarn = print;
    trackPropagator.setPrintWarnings(print);
}

void HKalSystem::setRotationAngles(Double_t rotXangle, Double_t rotYangle) {
    // Set the rotation matrix.
    // rotXangle: rotation angle around x axis in radians.
    // rotYangle: rotation angle around y axis in radians.
    pRotMat->SetToIdentity();
    pRotMat->RotateX(rotXangle);
    pRotMat->RotateY(rotYangle);
}