#include "hforwardtools.h"

#include "frpcdef.h"
#include "hcategory.h"
#include "hcategorymanager.h"
#include "hforwardcand.h"
#include "hfrpchit.h"
#include "hgeantkine.h"
#include "hmdcsizescells.h"
#include "hparticletool.h"
#include "hvirtualcandsim.h"

#include <TVector2.h>

namespace HForwardTools
{

Float_t calcFRpcHitDistance(HForwardCand* fcand, HFRpcHit* hit)
{
    if (!fcand or !hit) return -100.0f;

    Double_t params[5];
    HForwardTools::getParams(fcand, params);
    Double_t z_sts = params[4];

    Float_t x_rpc = hit->getX();
    Float_t y_rpc = hit->getY();
    Float_t z_rpc = hit->getZ();

    Double_t dz = z_rpc - z_sts;
    Double_t x_atz_rpc = params[0] + params[2] * dz;
    Double_t y_atz_rpc = params[1] + params[3] * dz;

    TVector2 v_ext(x_atz_rpc, y_atz_rpc);
    TVector2 v_rpc(x_rpc, y_rpc);
    TVector2 v_dist = v_ext - v_rpc;
    return v_dist.Mod();
}

void setHadesParams(HVirtualCand* fcand, TVector3 base, TVector3 dir)
{
    // Convert to HADES track parameters
    Double_t x1 = base.X(), y1 = base.Y(), z1 = base.Z();
    Double_t dz = 100.0;
    Double_t z2 = z1 + dz, x2 = x1 + dir.X() * dz, y2 = y1 + dir.Y() * dz;

    Double_t p0, p1, p2, p3;
    HMdcSizesCells::calcMdcSeg(x1, y1, z1, x2, y2, z2, p0, p1, p2, p3);
    if (p3 < 0.) p3 += TMath::Pi() * 2.0;

    fcand->setZ(p0);
    fcand->setR(p1);
    fcand->setTheta(p2 * TMath::RadToDeg());
    fcand->setPhi(p3 * TMath::RadToDeg());
}

TVector2 getPlaneIntersect(HGeomVector base, HGeomVector dir, Float_t z)
{
    Float_t x2_proj = base.X() + dir.X() / dir.Z() * (z - base.Z());
    Float_t y2_proj = base.Y() + dir.Y() / dir.Z() * (z - base.Z());

    return TVector2(x2_proj, y2_proj);
}

TVector2 getPlaneIntersect(HForwardCand* fcand, Float_t z)
{
    if (!fcand) throw("");

    HGeomVector base, dir;
    HParticleTool::calcSegVector(fcand->getZ(), fcand->getR(), fcand->getPhi() * TMath::DegToRad(),
                                 fcand->getTheta() * TMath::DegToRad(), base, dir);

    return getPlaneIntersect(base, dir, z);
}

void getParams(HForwardCand* const fcand, Double_t* pars)
{
    if (!fcand or !pars) throw("");

    HGeomVector base, dir;
    HParticleTool::calcSegVector(fcand->getZ(), fcand->getR(), fcand->getPhi() * TMath::DegToRad(),
                                 fcand->getTheta() * TMath::DegToRad(), base, dir);

    // get X, Y, Tx, Ty
    pars[0] = base.X();
    pars[1] = base.Y();
    pars[2] = dir.X() / dir.Z();
    pars[3] = dir.Y() / dir.Z();
    pars[4] = base.Z();
}

void setParams(HForwardCand* fcand, Double_t* pars)
{
    if (!fcand or !pars) throw("");

    setHadesParams(fcand, TVector3(pars[0], pars[1], pars[4]), TVector3(pars[2], pars[3], 1));
}

Float_t calcGammaSquaredMinusOne(Float_t length, Float_t tof)
{
    /* for tof^2 in ns:
     * [10^-9 * 10^8  = 10^-1]
     */
    const Double_t c = 2.99792458e-1;
    const Double_t c2 = c * c;

    Float_t zz = length * 1.0e-03;
    Float_t zz2 = zz * zz;

    return zz2 / (tof * tof * c2 - zz2);
}

Float_t calcMomentum(Float_t length, Float_t tof, Float_t mass)
{
    Float_t gamma_mo_sqrt = calcGammaSquaredMinusOne(length, tof);
    if (gamma_mo_sqrt <= 0) { return -1.0; }

    return mass * TMath::Sqrt(gamma_mo_sqrt);
}

void correctPathLength(HForwardCand* fcand, HGeomVector vertex, HFRpcHit* frpchit)
{
    if (!frpchit) return;

    // clear flags which may be re-evaluated
    fcand->unsetFlagBit(HForward::kErrorTofRec);
    fcand->unsetFlagBit(HForward::kErrorLength);
    fcand->unsetFlagBit(HForward::kIsRejected);

    if (vertex.Z() == -1000.0f)
    {
        fcand->setDistance(-10.0f);
        fcand->setFlagBit(HForward::kErrorLength);
        return;
    }

    HGeomVector frpc_pos = make_vector<HGeomVector>(frpchit);
    fcand->setDistance((frpc_pos - vertex).length());
}

void correctPathTof(HForwardCand* fcand, Float_t start_time, HFRpcHit* frpchit, Float_t tof_wl, Float_t tof_wu)
{
    if (!frpchit) return;

    // clear flags which may be re-evaluated
    fcand->unsetFlagBit(HForward::kIsRejected);
    fcand->unsetFlagBit(HForward::kErrorTofRec);
    fcand->unsetFlagBit(HForward::kErrorTof);

    if (start_time == -1000.0f)
    {
        fcand->setTof(-10.0f);
        fcand->setFlagBit(HForward::kErrorTof);
        return;
    }

    Float_t tof = frpchit->getTof();
    if (tof < tof_wl or tof > tof_wu)
    {
        fcand->setFlagBit(HForward::kErrorTof);
        fcand->setFlagBit(HForward::kIsRejected);
    }
    tof -= start_time;
    fcand->setTof(tof);
}

Bool_t fillVirtualCandSimInfo(HVirtualCandSim* part, HCategory* pKine, Int_t track)
{
    if (!pKine or track < 1) return kFALSE;

    HGeantKine* kine = HCategoryManager::getObject(kine, pKine, track - 1);
    if (!kine)
    {
        cout << "ERROR :: HForwardTools::fillVirtualCandSimInfo() : Zero HGeantKine pointer retrieved for cand " << part
             << " track Number " << track << " size of KineCat " << pKine->getEntries() << endl;
        return kFALSE;
    }

    Int_t pid = kine->getID();

    Float_t px, py, pz;
    kine->getMomentum(px, py, pz);

    Float_t vx, vy, vz;
    kine->getVertex(vx, vy, vz);

    Int_t parentTr, medium, mechanism;
    kine->getCreator(parentTr, medium, mechanism);
    Int_t geninfo, geninfo1, geninfo2;
    Float_t genweight;
    kine->getGenerator(geninfo, genweight);
    kine->getGenerator(geninfo, geninfo1, geninfo2);

    Int_t parentPid = -1, grandparentTr = 0, grandparentPid = -1;
    if (parentTr > 0)
    {
        kine = HCategoryManager::getObject(kine, pKine, parentTr - 1);

        if (!kine)
        {
            cout << "ERROR :: HForwardTools::fillVirtualCandSimInfo() : Zero HGeantKine pointer retrieved for cand "
                 << part << " parent track Number " << track << " size of KineCat " << pKine->getEntries() << endl;
        }
        else
        {
	    parentPid = kine->getID();
	    grandparentTr = kine->getParentTrack();
	    if (grandparentTr > 0)
	    {
		kine = HCategoryManager::getObject(kine, pKine, grandparentTr - 1);
                if (!kine)
                {
                    cout << "ERROR :: HForwardTools::fillVirtualCandSimInfo() : Zero HGeantKine pointer retrieved for "
                            "cand "
                         << part << " grand parent track Number " << track << " size of KineCat " << pKine->getEntries()
                         << endl;
		}
                else { grandparentPid = kine->getID(); }
	    }
	}
    }

    part->setGeantPID(pid);
    part->setGeantTrack(track);
    part->setGeantxMom(px);
    part->setGeantyMom(py);
    part->setGeantzMom(pz);
    part->setGeantxVertex(vx);
    part->setGeantyVertex(vy);
    part->setGeantzVertex(vz);
    part->setGeantParentTrackNum(parentTr);
    part->setGeantParentPID(parentPid);
    part->setGeantGrandParentTrackNum(grandparentTr);
    part->setGeantGrandParentPID(grandparentPid);
    part->setGeantCreationMechanism(mechanism);
    part->setGeantMediumNumber(medium);
    part->setGeantGeninfo(geninfo);
    part->setGeantGeninfo1(geninfo1);
    part->setGeantGeninfo2(geninfo2);
    part->setGeantGenweight(genweight);

    return kTRUE;
}

Float_t getFRpcDistance(HForwardCand* fcand, HFRpcHit* hit)
{
    Float_t x_rpc = hit->getX();
    Float_t y_rpc = hit->getY();
    Float_t z_rpc = hit->getZ();
    TVector2 v_ext = HForwardTools::getPlaneIntersect(fcand, z_rpc);
    TVector2 v_rpc(x_rpc, y_rpc);

    TVector2 v_dist = v_ext - v_rpc;
    return v_dist.Mod();
}

} // namespace HForwardTools