#include "hemcneutralcandfinder.h"

#include "hades.h"
#include "hcategory.h"
#include "hcategorymanager.h"
#include "hevent.h"
#include "hlinearcategory.h"
#include "hlocation.h"

#include "emcdef.h"

#include "hparticletool.h"
#include "hparticlecand.h"
#include "hemccluster.h"
#include "hemcclustersim.h"
#include "hgeantkine.h"
#include "hemcneutralcand.h"
#include "hemcneutralcandsim.h"

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

//_HADES_CLASS_DESCRIPTION
/////////////////////////////////////////////////////////////////////////
//
//  HEmcNeutralCandFinder:
//  creates HEmcNeutralCand/Sim categories for emc clusters unmatched with
//  reconstructed HParticleCand for analysis of neutral particles with EMC,
//  By default PID is assumed to be 0 (photon). The TLorentz 4vector of
//  HEmcNeutralCand is precalculated with mass 0. If the particle should
//  be used as neutron cand->calc4vectorProperties(HPhysicsConstants::mass("n"))
//  should be called to use proper TLorentzVector arithmetics with the
//  neutral candidate. Needs HParticleCand and HEmcCluster as input.
//
/////////////////////////////////////////////////////////////////////////

ClassImp(HEmcNeutralCandFinder);

HEmcNeutralCandFinder::HEmcNeutralCandFinder(const Text_t* name, const Text_t* title)
    : HReconstructor(name, title)
{
    initVariables();
}

HEmcNeutralCandFinder::~HEmcNeutralCandFinder()
{

}

Float_t HEmcNeutralCandFinder::calcMomentum(Float_t length, Float_t tof, Float_t mass)
{
    /* 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;

    Float_t gamma_mo_sqrt = zz2 / (tof * tof * c2 - zz2);

    if (gamma_mo_sqrt <= 0) { return -1.0; }

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

void HEmcNeutralCandFinder::initVariables()
{
    isSimulation = kFALSE;
    pKine = nullptr;
    pEmcCluster = nullptr;
    pParticleCand = nullptr;
    pEmcNeutralCand = nullptr;

    fTimeMin = 5.00f;
    fTimeMax = 100.0f;

    femcMQCut = 2.5;

    fPhotonBetaCutMin = 0.8;
    fPhotonBetaCutMax = 1.2;
}

Bool_t HEmcNeutralCandFinder::init()
{
    // GEANT input data
    pKine = gHades->getCurrentEvent()->getCategory(catGeantKine);
    if (pKine) { isSimulation = kTRUE; }

    pEmcCluster   = HCategoryManager::getCategory(catEmcCluster  , 0, "catEmcCluster, from HEmcNeutralCandFinder::init()");
    pParticleCand = HCategoryManager::getCategory(catParticleCand, 0, "catParticleCand, from HEmcNeutralCandFinder::init()");

    // build the Straw vector category
    if (isSimulation) { pEmcNeutralCand = new HLinearCategory("HEmcNeutralCandSim"); }
    else              { pEmcNeutralCand = new HLinearCategory("HEmcNeutralCand");    }

    if (!pEmcNeutralCand)
    {
        Error("HEmcNeutralCandFinder::init()", "EmcNeutralCand category not created");
        return kFALSE;
    } else {
	gHades->getCurrentEvent()->addCategory(catEmcNeutralCand, pEmcNeutralCand, "Emc");
    }

    return kTRUE;
}

Bool_t HEmcNeutralCandFinder::reinit()
{
    return kTRUE;
}

Int_t HEmcNeutralCandFinder::execute()
{

    if(!pEmcCluster || !pParticleCand) {  return 0; }


    // From all the clusters pool, remove those being attributed with particle cand.
    Int_t nPartCand   = pParticleCand->getEntries();
    Int_t nEmcCluster = pEmcCluster  ->getEntries();

    for (Int_t i = 0; i < nPartCand; i++)
    {
        HParticleCand* pcand = nullptr;
        pcand = HCategoryManager::getObject(pcand, pParticleCand, i);
        Int_t ecalIndex = pcand->getEmcInd();

        if (ecalIndex != -1)
        {
            HEmcCluster* clus = dynamic_cast<HEmcCluster*>(pEmcCluster->getObject(ecalIndex));

            // if it has a very good matching with EMC, most likely charged particle
            Double_t emcMQ = pcand->getMetaMatchQualityEmc();
            if (emcMQ < femcMQCut)
            {
                if (clus->getNMatchedTracks() == 0) clus->addMatchedTrack();   // a track has been matched to the particular EMC cluster
            }
        }
    }

    HLocation loc;
    loc.set(1, 0);

    HEventHeader* header = gHades->getCurrentEvent()->getHeader();

    // get Event vertex
    HVertex primVertex = header->getVertexReco();
    if (primVertex.getChi2() < 0)
    { // vertex reco has not been reconstructed: fall back solution : Vertex cluster
        primVertex = header->getVertexCluster();
    }
    TVector3 vertex(primVertex.getX(), primVertex.getY(), primVertex.getZ());

    for (Int_t i = 0; i < nEmcCluster; i++)
    { // loop over ECAL clusters
        HEmcCluster* clus = dynamic_cast<HEmcCluster*>(pEmcCluster->getObject(i));

        // if it has matched charged track, ignore
        if (clus->getNMatchedTracks() > 0) continue;
        if (clus->getNMatchedCells() > 0) continue;

        // we discard clusters from outside of the optimal time window
        Float_t tof = clus->getTime();
        if (tof < fTimeMin || tof > fTimeMax) continue;

        // calc path length
        TVector3 emc_hit(clus->getXLab(), clus->getYLab(), clus->getZLab());
        TVector3 dist = emc_hit - vertex;

        Float_t beta = (dist.Mag() / 1000.) / (tof * 1.e-9 * TMath::C());
        Char_t pid = -1;
        if (beta < fPhotonBetaCutMin)
        {
            // consider as neutron
            pid = 13;
        }
        else if (beta < fPhotonBetaCutMax)
        {
            // consider as photon
            Int_t cell = emc_cell2tile_map[(Int_t)clus->getCell()]; // cell in format 0-162
            // the photon cluster cannot originate in first 2 rows of ECAL(closest to the beam),
            // there is no acceptance for photons due to RICH/MDC material
            if (cell < 10) continue;
            pid = 1;
        }

        HEmcNeutralCand* cand = (HEmcNeutralCand*)pEmcNeutralCand->getNewSlot(loc);

        if (cand)
        {
            if (isSimulation)
            {
                HEmcNeutralCandSim* candsim = new (cand) HEmcNeutralCandSim();
                HEmcClusterSim* clussim = dynamic_cast<HEmcClusterSim*>(clus);
		if (!clussim) return kFALSE; // something wrong : reject  event

                Int_t track = clussim->getTrack();

                fillVirtualCandSimInfo(candsim, pKine, track);

                cand = dynamic_cast<HEmcNeutralCand*>(candsim);
            } // end simulation
            else { cand = new (cand) HEmcNeutralCand(); }

            cand->setEmcClusterIndex(i);

            Double_t p0, p1, p2, p3;
            HParticleTool::calcMdcSeg(vertex.X(), vertex.Y(), vertex.Z(), emc_hit.X(), emc_hit.Y(), emc_hit.Z(), p0, p1, p2, p3);
            if (p3 < 0.) p3 += TMath::Pi() * 2.0;

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

            cand->setDistanceToEmc(dist.Mag());
            cand->setTof(tof);
            cand->setBeta(beta);
            cand->setPID(pid);

            if (pid == 13)
            {
                // consider as neutron
                // momentum must be calculated from tof and distance
                cand->setMomentum(calcMomentum(cand->getDistanceToEmc(), cand->getTof(),HPhysicsConstants::mass(13)));
                cand->calc4vectorProperties(HPhysicsConstants::mass(13));
            }
            else if (pid == 1)
            {
                // consider as photon
                cand->setMomentum(clus->getEnergy()); // energy deposited in EMC is equivalent of the photon momentum
                cand->calc4vectorProperties(0.0);
            }
            cand->SetUniqueID(i);
        }
    } // end of photon part

    return 0;
}

Bool_t HEmcNeutralCandFinder::finalize()
{ 
    return kTRUE; 
}


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

    HGeantKine* kine = 0;
    kine = HCategoryManager::getObject(kine, pKine, track - 1);

    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);
    Float_t geninfo, geninfo1, geninfo2, 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);
        parentPid = kine->getID();
        grandparentTr = kine->getParentTrack();
        if (grandparentTr > 0)
        {
            kine = HCategoryManager::getObject(kine, pKine, grandparentTr - 1);
            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;
}