// -------------------------------------------------------------------------
// -----               PndRhoFitProbPresorter source file              -----
// -------------------------------------------------------------------------

#include <iostream>

#include "TVector3.h"

#include "PndRhoPropDistToIP.h"

PndRhoPropDistToIP::PndRhoPropDistToIP( int verboseLevel ) : fVerbose(verboseLevel),
        fPropSuccessful(kFALSE), fPropPocaIPSuccessful(kFALSE), fPropHelixSuccessful(kFALSE),
        fCandWorks(kFALSE)
{
    Init();
}

PndRhoPropDistToIP::~PndRhoPropDistToIP()
{
    delete fGeaneProp;
}

// initialize the propagators
void PndRhoPropDistToIP::Init()
{
    if ( fVerbose > 0 ) std::cout << "PndRhoPropDistToIP:  Init()" << std::endl;

    // vectors
    fNullVector = TVector3(0,0,0);
    fDirectionX = TVector3(1,0,0);
    fDirectionY = TVector3(0,1,0);

    // Geane propagator
    fGeaneProp = new FairGeanePro();
}

Bool_t PndRhoPropDistToIP::SetParticleTrackPar( RhoCandidate * cand, TVector3 vtxPos )
{
    fCand = cand;

    // can we use this candidate?
    if ( TMath::IsNaN(fCand->GetPosition().X()) || TMath::IsNaN(fCand->GetPosition().Y()) || TMath::IsNaN(fCand->GetPosition().Z()) ||
         TMath::IsNaN(fCand->GetMomentum().X()) || TMath::IsNaN(fCand->GetMomentum().Y()) || TMath::IsNaN(fCand->GetMomentum().Z()) ||
         TMath::IsNaN(fCand->GetCharge()) ) {
        fCandWorks = kFALSE;
        return kFALSE;
    }
    fCandWorks = kTRUE;


    // define some temporary vectors
    TVector3 pos = vtxPos; //fCand->GetPosition();
    TVector3 posErr = fNullVector;
    TVector3 mom = fCand->GetMomentum();
    TVector3 momErr = fNullVector;
    TVector3 planeO = pos;
    TVector3 planeU = fDirectionX;
    TVector3 planeV = fDirectionY;
    Double_t charge = fCand->GetCharge();
    Int_t ierr = 0;

    // combine parameters to FairTrack objects
    fTrackPoint = FairTrackParP(pos, mom, posErr, momErr, charge, planeO, planeU, planeV);
    fTrackHelix = FairTrackParH(&fTrackPoint, ierr);

    // set options for helix propagator
    fHelixProp = PndHelixPropagator(2.0, pos, mom, charge);

    // reset the results
    fPropResult = FairTrackParP();
    fPropResultPocaIP = FairTrackParH();
    fPropResultHelix = FairTrackPar();

    if ( fVerbose > 2 ) {
        printf("\nPndRhoPropDistToIP:  SetParticleTrackPar(RhoCandidate *)\n");
        printf("CandVtx  x: %7.2f  y: %7.2f  z: %7.2f\n", fCand->GetPosition().X(), fCand->GetPosition().Y(), fCand->GetPosition().Z() );
        printf("CandMom  x: %7.2f  y: %7.2f  z: %7.2f\n", fCand->Px(), fCand->Py(), fCand->Pz() );
    }

    return kTRUE;
}

Bool_t PndRhoPropDistToIP::PropagateWithHelix()
{
    if ( !fCandWorks )
        return kFALSE;

    // check, if the candidates momentum and z coordinate show in the same direction (i.e. flying away from IP)
    if ( fTrackPoint.GetZ() * fTrackPoint.GetPz() <= 0 ) {
        fPropHelixSuccessful = kFALSE;
    } else {
        fPropHelixSuccessful = kTRUE;
        fPropResultHelix = fHelixProp.PropagateToZ(0.0);
    }

    if ( fVerbose > 0 ) {
        std::cout << "PndRhoPropDistToIP:  PropagateWithHelix() " << ( fPropHelixSuccessful ? "succeeded" : "failed" ) << std::endl;
        if ( fVerbose > 2 ) {
            printf("PropVtx  x: %7.2f  y: %7.2f  z: %7.2f\n", fPropResultHelix.GetX(), fPropResultHelix.GetY(), fPropResultHelix.GetZ() );
        }
    }

    return fPropHelixSuccessful;
}

Bool_t PndRhoPropDistToIP::PropagateWithGeane()
{
    if ( !fCandWorks )
        return kFALSE;

    // back-propagate to the xy plane at z=0
    fGeaneProp->PropagateToPlane( fNullVector, fDirectionX, fDirectionY );
    fGeaneProp->setBackProp();  // needs to be called after PropagateToPlane()

    // do the propagation
    fPropSuccessful = fGeaneProp->Propagate(&fTrackPoint, &fPropResult, fCand->PdgCode());
    if ( fVerbose > 0 ) {
        std::cout << "PndRhoPropDistToIP:  PropagateWithGeane() " << ( fPropSuccessful ? "succeeded" : "failed" ) << std::endl;
        if ( fVerbose > 2 && fPropSuccessful ) {
            printf("PropVtx  x: %7.2f  y: %7.2f  z: %7.2f\n", fPropResult.GetX(), fPropResult.GetY(), fPropResult.GetZ() );
        }
    }

    return fPropSuccessful;
}

Bool_t PndRhoPropDistToIP::PropagateWithGeanePocaIP()
{
    if ( !fCandWorks )
        return kFALSE;

    // propagate to the closest point to IP
    fGeaneProp->BackTrackToVertex();

    // do the propagation
    fPropPocaIPSuccessful = fGeaneProp->Propagate(&fTrackHelix, &fPropResultPocaIP, fCand->PdgCode());
    if ( fVerbose > 0 ) {
        std::cout << "PndRhoPropDistToIP:  PropagateWithGeanePocaIP() " << ( fPropPocaIPSuccessful ? "succeeded" : "failed" ) << std::endl;
        if ( fVerbose > 2 && fPropPocaIPSuccessful ) {
            printf("PropVtxH x: %7.2f  y: %7.2f  z: %7.2f\n", fPropResultPocaIP.GetX(), fPropResultPocaIP.GetY(), fPropResultPocaIP.GetZ() );
        }
    }
    return fPropPocaIPSuccessful;
}

Double_t PndRhoPropDistToIP::GetDistanceToIP()
{
    if ( fCandWorks && fPropSuccessful ) {
        return fPropResult.GetPosition().Mag();
    } else {
        return -1;
    }
}

Double_t PndRhoPropDistToIP::GetDistanceToIPPocaIP()
{
    if ( fCandWorks && fPropPocaIPSuccessful ) {
        return fPropResultPocaIP.GetPosition().Mag();
    } else {
        return -1;
    }
}

Double_t PndRhoPropDistToIP::GetDistanceToIPHelix()
{
    if ( fCandWorks && fPropHelixSuccessful ) {
        return fPropResultHelix.GetPosition().Mag();
    } else {
        return -1;
    }
}

void PndRhoPropDistToIP::qaPropResult(TString pre, RhoTuple * n, bool skip)
{
    TString pre2;

    // propagate with Geane to POCA
    pre2 = "gpo_";
    if ( !skip && fCandWorks && fPropPocaIPSuccessful ) {
        qaTrackParResult(pre+pre2, n, (FairTrackPar*)&fPropResultPocaIP);
    } else {
        qaTrackParResult(pre+pre2, n, (FairTrackPar*)0x0, true);
    }

    // propagate with Geane to plane
    pre2 = "gpl_";
    if ( !skip && fCandWorks && fPropSuccessful ) {
        qaTrackParResult(pre+pre2, n, (FairTrackPar*)&fPropResult);
    } else {
        qaTrackParResult(pre+pre2, n, (FairTrackPar*)0x0, true);
    }

    // propagate with helix to plane
    pre2 = "hpl_";
    if ( !skip && fCandWorks && fPropHelixSuccessful ) {
        qaTrackParResult(pre+pre2, n, (FairTrackPar*)&fPropResultHelix);
    } else {
        qaTrackParResult(pre+pre2, n, (FairTrackPar*)0x0, true);
    }
}

void PndRhoPropDistToIP::qaTrackParResult(TString pre, RhoTuple * n, FairTrackPar * trPar, bool skip)
{
    if ( !skip ) {
        // position
        n->Column(pre+"x",   (Float_t) trPar->GetX(),   0.0f );
        n->Column(pre+"y",   (Float_t) trPar->GetY(),   0.0f );
        n->Column(pre+"z",   (Float_t) trPar->GetZ(),   0.0f );
        n->Column(pre+"dx",  (Float_t) trPar->GetDX(),  0.0f );
        n->Column(pre+"dy",  (Float_t) trPar->GetDY(),  0.0f );
        n->Column(pre+"dz",  (Float_t) trPar->GetDZ(),  0.0f );

        // momentum
        n->Column(pre+"px",  (Float_t) trPar->GetPx(),  0.0f );
        n->Column(pre+"py",  (Float_t) trPar->GetPy(),  0.0f );
        n->Column(pre+"pz",  (Float_t) trPar->GetPz(),  0.0f );
        n->Column(pre+"dpx", (Float_t) trPar->GetDPx(), 0.0f );
        n->Column(pre+"dpy", (Float_t) trPar->GetDPy(), 0.0f );
        n->Column(pre+"dpz", (Float_t) trPar->GetDPz(), 0.0f );

        // distance to IP
        n->Column(pre+"dIP", (Float_t) trPar->GetPosition().Mag(), 0.0f );
    } else {
        // position
        n->Column(pre+"x",   (Float_t) -999., 0.0f );
        n->Column(pre+"y",   (Float_t) -999., 0.0f );
        n->Column(pre+"z",   (Float_t) -999., 0.0f );
        n->Column(pre+"dx",  (Float_t) -999., 0.0f );
        n->Column(pre+"dy",  (Float_t) -999., 0.0f );
        n->Column(pre+"dz",  (Float_t) -999., 0.0f );

        // momentum
        n->Column(pre+"px",  (Float_t) -999., 0.0f );
        n->Column(pre+"py",  (Float_t) -999., 0.0f );
        n->Column(pre+"pz",  (Float_t) -999., 0.0f );
        n->Column(pre+"dpx", (Float_t) -999., 0.0f );
        n->Column(pre+"dpy", (Float_t) -999., 0.0f );
        n->Column(pre+"dpz", (Float_t) -999., 0.0f );

        // distance to IP
        n->Column(pre+"dIP", (Float_t) -999., 0.0f );
    }
}