#include "fmcLoopProcess.h"
#include "fmcConstants.h"

#include <TMath.h>

ClassImp(fmcLoopProcess);

fmcLoopProcess::fmcLoopProcess() : fOptions(),
                                   fHeader(),
                                   fParticle()
{
}

fmcLoopProcess::~fmcLoopProcess()
{
    delete fOptions;
    delete fHeader;
    delete fParticle;
    delete fReader;
}

fmcParticleLoopInfo *fmcLoopProcess::ProcessParticleLoop(fmcLong iev)
{
    if (!fReader)
        return nullptr;
    if (!fOptions)
        return nullptr;

    fmcInt refMultSTAR = 0, refMultEastPHENIX = 0, refMultWestPHENIX = 0;
    fmcInt refMultPHENIX = 0;
    fmcDouble eta, pt, phi;
    fmcDouble qvTpcEastX[fmc::nTpcEtaGaps], qvTpcEastY[fmc::nTpcEtaGaps];
    fmcDouble qvTpcWestX[fmc::nTpcEtaGaps], qvTpcWestY[fmc::nTpcEtaGaps];
    fmcDouble qvFHCalEastX, qvFHCalEastY;
    fmcDouble qvFHCalWestX, qvFHCalWestY;
    fmcDouble qvRXNEastX, qvRXNEastY;
    fmcDouble qvRXNWestX, qvRXNWestY;
    fmcDouble qvWeight;

    fmcQvector *qv = nullptr;

    if (fOptions->qv())
    {
        qvFHCalEastX = 0.; qvFHCalEastY = 0.;
        qvFHCalWestX = 0.; qvFHCalWestY = 0.;
        qvRXNEastX = 0.;   qvRXNEastY = 0.;
        qvRXNWestX = 0.;   qvRXNWestY = 0.;
        for (int i = 0; i < fmc::nTpcEtaGaps; i++)
        {
            qvTpcEastX[i] = 0.;
            qvTpcEastY[i] = 0.;
            qvTpcWestX[i] = 0.;
            qvTpcWestY[i] = 0.;
        }
    }

    fmcParticleLoopInfo *info = new fmcParticleLoopInfo();

    fHeader = fReader->ReadEventHeader(iev);
    if (!fHeader)
        return nullptr;

    // Start main particle loop
    for (int iTr = 0; iTr < fHeader->GetNparts(); iTr++)
    {
        fParticle = fReader->ReadParticle(iTr);
        if (!fParticle)
            continue;
        if (fParticle->GetCharge() == 0)
        {
            delete fParticle;
            continue;
        }

        if (fOptions->refMult() || fOptions->qv())
        {
            eta = fParticle->GetEta();
            pt  = fParticle->GetPt();
        }

        if (fOptions->qv())
        {
            phi = fParticle->GetPhi();
        }

        // Calculate refMult
        if (fOptions->refMult())
        {
            if (TMath::Abs(eta) <= fmc::refMultSTAREtaCut && pt >= fmc::ptMinCut)
                refMultSTAR++;
            if (TMath::Abs(eta) >= fmc::refMultRXNEtaCut.first &&
                TMath::Abs(eta) <= fmc::refMultRXNEtaCut.second &&
                eta < 0)
                refMultEastPHENIX++;
            if (TMath::Abs(eta) >= fmc::refMultRXNEtaCut.first &&
                TMath::Abs(eta) <= fmc::refMultRXNEtaCut.second &&
                eta > 0)
                refMultWestPHENIX++;
        }

        // Calculate Q-vectors
        if (fOptions->qv())
        {
            qvWeight = 1.;

            // Qv from TPC East
            if (TMath::Abs(eta) <= fmc::TpcEtaCut &&
                eta < 0 && pt < fmc::ptMaxCutEP &&
                pt > fmc::ptMinCut)
            {
                for (int i = 0; i < fmc::nTpcEtaGaps; i++)
                {
                    if (TMath::Abs(eta) >= fmc::TpcEtaGap.at(i))
                    {
                        qvTpcEastX[i] += qvWeight * TMath::Cos(2 * phi);
                        qvTpcEastY[i] += qvWeight * TMath::Sin(2 * phi);
                    }
                }
            }
            // Qv from TPC West
            if (TMath::Abs(eta) <= fmc::TpcEtaCut &&
                eta > 0 && pt < fmc::ptMaxCutEP &&
                pt > fmc::ptMinCut)
            {
                for (int i = 0; i < fmc::nTpcEtaGaps; i++)
                {
                    if (TMath::Abs(eta) >= fmc::TpcEtaGap.at(i))
                    {
                        qvTpcWestX[i] += qvWeight * TMath::Cos(2 * phi);
                        qvTpcWestY[i] += qvWeight * TMath::Sin(2 * phi);
                    }
                }
            }

            // Qv from RXN East
            if (TMath::Abs(eta) >= fmc::refMultRXNEtaCut.first &&
                TMath::Abs(eta) <= fmc::refMultRXNEtaCut.second &&
                eta < 0)
            {
                qvRXNEastX += qvWeight * TMath::Cos(2 * phi);
                qvRXNEastY += qvWeight * TMath::Sin(2 * phi);
            }
            // Qv from RXN West
            if (TMath::Abs(eta) >= fmc::refMultRXNEtaCut.first &&
                TMath::Abs(eta) <= fmc::refMultRXNEtaCut.second &&
                eta > 0)
            {
                qvRXNWestX += qvWeight * TMath::Cos(2 * phi);
                qvRXNWestY += qvWeight * TMath::Sin(2 * phi);
            }

            // Qv from FHCal East
            if (TMath::Abs(eta) >= fmc::refMultFHCalEtaCut.first &&
                TMath::Abs(eta) <= fmc::refMultFHCalEtaCut.second &&
                eta < 0)
            {
                qvFHCalEastX += qvWeight * TMath::Cos(1 * phi);
                qvFHCalEastY += qvWeight * TMath::Sin(1 * phi);
            }
            // Qv from FHCal West
            if (TMath::Abs(eta) >= fmc::refMultFHCalEtaCut.first &&
                TMath::Abs(eta) <= fmc::refMultFHCalEtaCut.second &&
                eta > 0)
            {
                qvFHCalWestX += -1 * qvWeight * TMath::Cos(1 * phi);
                qvFHCalWestY += -1 * qvWeight * TMath::Sin(1 * phi);
            }
        }
        delete fParticle;
    } // end of particle loop

    if (fOptions->refMult())
    {
        refMultPHENIX = (refMultEastPHENIX >= 2 && refMultWestPHENIX >= 2) ?
            refMultEastPHENIX + refMultWestPHENIX : 0;
        info->SetMultSTAR(refMultSTAR);
        info->SetMultPHENIX(refMultPHENIX);
    }

    if (fOptions->qv())
    {
        qv = new fmcQvector();
        for (int i = 0; i < fmc::nTpcEtaGaps; i++)
        {
            qv->Set(qvTpcEastX[i], qvTpcEastY[i]);
            qv->SetHarmonic(2.);
            info->AddQv(qv);
            qv->Set(qvTpcWestX[i], qvTpcWestY[i]);
            qv->SetHarmonic(2.);
            info->AddQv(qv);
        }
        qv->Set(qvRXNEastX, qvRXNEastY);
        qv->SetHarmonic(2.);
        info->AddQv(qv);
        qv->Set(qvRXNWestX, qvRXNWestY);
        qv->SetHarmonic(2.);
        info->AddQv(qv);
        qv->Set(qvFHCalEastX, qvFHCalEastY);
        qv->SetHarmonic(1.);
        info->AddQv(qv);
        qv->Set(qvFHCalWestX, qvFHCalWestY);
        qv->SetHarmonic(1.);
        info->AddQv(qv);
        delete qv;
    }

    delete fHeader;
    return info;
}