#include "CbmKFParticleInterface.h"
#include "CbmKFParticleFinder.h"

#define cnst static const fvec

//for particle finding
#include "CbmKFTrackFitter.h"
#include "CbmL1Track.h"
#include "CbmKFTrack.h"
#include "L1AlgoInputData.h"
#include "L1Algo.h"

const float CbmKFParticleFinder::DefaultCuts[2][3]  = {{3.,3.,-100.},{3.,3.,-100.}};

template<class T> 
void CbmKFParticleFinder::FindParticlesT(vector<T> &vRTracks, vector<CbmKFParticle>& Particles, 
                                         CbmKFVertex& PrimVtx, const float cuts[2][3])
{
  //* Finds particles (K0s and Lambda) from a given set of tracks

  CbmKFTrackFitter fitter;
  fitter.SetHits(fHits);
  fitter.SetAlgo(fAlgo);

  vector<T> vPos;
  vector<T> vNeg;

  vector<short> idPos;
  vector<short> idNeg;

  vector<T> vPosPrim;
  vector<T> vNegPrim;

  vector<short> idPosPrim;
  vector<short> idNegPrim;

  // fitter.Fit(vRTracks); TODO copy from L1Algo

  vector<float> ChiToPrimVtx;
  fitter.GetChiToVertex(vRTracks,ChiToPrimVtx,PrimVtx);

  for(unsigned short iTr=0; iTr < vRTracks.size(); iTr++)
  {
    CbmKFTrack kfTrack(vRTracks[iTr]);
    bool ok = 1;
    for(unsigned short iT=0; iT<5; iT++)
      ok = ok && finite(kfTrack.GetTrack()[iT]);
    for(unsigned short iC=0; iC<15; iC++)
      ok = ok && finite(kfTrack.GetCovMatrix()[iC]);
      ok = ok && (kfTrack.GetCovMatrix()[0] < 1. && kfTrack.GetCovMatrix()[0] > 0.)
              && (kfTrack.GetCovMatrix()[2] < 1. && kfTrack.GetCovMatrix()[2] > 0.)
              && (kfTrack.GetCovMatrix()[5] < 1. && kfTrack.GetCovMatrix()[5] > 0.)
              && (kfTrack.GetCovMatrix()[9] < 1. && kfTrack.GetCovMatrix()[9] > 0.)
              && (kfTrack.GetCovMatrix()[14] < 1. && kfTrack.GetCovMatrix()[14] > 0.);
      ok = ok && kfTrack.GetRefChi2() < 10*kfTrack.GetRefNDF();
    if(!ok) continue;

    if( ChiToPrimVtx[iTr] < cuts[0][0] )
    {
      if(kfTrack.GetTrack()[4] >= 0.)
      {
        vPosPrim.push_back(vRTracks[iTr]);
        idPosPrim.push_back(iTr);
      }
      if(kfTrack.GetTrack()[4] < 0.)
      {
        vNegPrim.push_back(vRTracks[iTr]);
        idNegPrim.push_back(iTr);
      }
    }
    else
    {
      if(kfTrack.GetTrack()[4] >= 0.)
      {
        vPos.push_back(vRTracks[iTr]);
        idPos.push_back(iTr);
      }
      if(kfTrack.GetTrack()[4] < 0.)
      {
        vNeg.push_back(vRTracks[iTr]);
        idNeg.push_back(iTr);
      }
    }
  }

  const Int_t PiPlusPDG = 211;
  const Int_t PiMinusPDG =-211;

  const Int_t KMinusPDG =-321;

  const Int_t PPlusPDG = 2212;
//  const Int_t PMinusPDG =-2212;

  const float massLambdaPDG = 1.115683;


  vector<L1FieldRegion> posB;
  vector<L1FieldRegion> negB;

  fitter.CalculateFieldRegion(vPos,posB);
  fitter.CalculateFieldRegion(vNeg,negB);

  for(unsigned short iTrP=0; iTrP < vPos.size(); iTrP++) {
    CbmKFTrack kfTrack(vPos[iTrP]);
    kfTrack.SetId(idPos[iTrP]);
    CbmKFParticle tmp(&kfTrack);
    tmp.SetPDG(211);
    tmp.SetId(Particles.size());
    int entrSIMD = iTrP % fvecLen;
    int entrVec  = iTrP / fvecLen;
    posB[entrVec].GetOneEntry(tmp.fieldRegion,entrSIMD);
    Particles.push_back(tmp);
  }
  for(unsigned short iTrN=0; iTrN < vNeg.size(); iTrN++) {
    CbmKFTrack kfTrack(vNeg[iTrN]);
    kfTrack.SetId(idNeg[iTrN]);
    CbmKFParticle tmp(&kfTrack);
    tmp.SetPDG(-211);
    tmp.SetId(Particles.size());
    int entrSIMD = iTrN % fvecLen;
    int entrVec  = iTrN / fvecLen;
    negB[entrVec].GetOneEntry(tmp.fieldRegion,entrSIMD);
    Particles.push_back(tmp);
  }

  vector<L1FieldRegion> posPrimB;
  vector<L1FieldRegion> negPrimB;

  fitter.CalculateFieldRegion(vPosPrim,posPrimB);
  fitter.CalculateFieldRegion(vNegPrim,negPrimB);

  for(unsigned short iTrP=0; iTrP < vPosPrim.size(); iTrP++) {
    CbmKFTrack kfTrack(vPosPrim[iTrP]);
    kfTrack.SetId(idPosPrim[iTrP]);
    CbmKFParticle tmp(&kfTrack);
    tmp.SetPDG(211);
    tmp.SetId(Particles.size());
    int entrSIMD = iTrP % fvecLen;
    int entrVec  = iTrP / fvecLen;
    posPrimB[entrVec].GetOneEntry(tmp.fieldRegion,entrSIMD);
    Particles.push_back(tmp);
  }
  for(unsigned short iTrN=0; iTrN < vNegPrim.size(); iTrN++) {
    CbmKFTrack kfTrack(vNegPrim[iTrN]);
    kfTrack.SetId(idNegPrim[iTrN]);
    CbmKFParticle tmp(&kfTrack);
    tmp.SetPDG(-211);
    tmp.SetId(Particles.size());
    int entrSIMD = iTrN % fvecLen;
    int entrVec  = iTrN / fvecLen;
    negPrimB[entrVec].GetOneEntry(tmp.fieldRegion,entrSIMD);
    Particles.push_back(tmp);
  }

  vector<CbmKFParticle> vK0s;
  vector<CbmKFParticle> vLambda;
  vector<CbmKFParticle> vKsiPlus;
  vector<CbmKFParticle> vKsiMinus;
  vector<CbmKFParticle> vSigmaStarPlus;
  vector<CbmKFParticle> vSigmaStarMinus;
  vector<CbmKFParticle> vOmegaMinus;
  vector<CbmKFParticle> vHdibarion;

  vector<float> vLambdaTopoChi2Ndf;
  vector<CbmKFParticle> vLambdaSec;
  vector<CbmKFParticle> vLambdaPrim;

  unsigned short NPositive = vPos.size();
  unsigned short NNegative = vNeg.size();
  unsigned short NPositivePrim = vPosPrim.size();

  for(unsigned short iTrN=0; iTrN < vNeg.size(); iTrN++)
  {
    CbmKFTrack kfTrackNeg(vNeg[iTrN]);
    CbmKFParticle_simd vDaughters[2] = {CbmKFParticle_simd(kfTrackNeg,0,&PiMinusPDG),
                                        CbmKFParticle_simd()};
    int entrSIMD = iTrN % fvecLen;
    int entrVec  = iTrN / fvecLen;
    vDaughters[0].SetField(negB[entrVec],1,entrSIMD);
    vDaughters[0].SetId(float(iTrN+NPositive));

    CbmKFTrack kfTrackPos[fvecLen];
    CbmKFTrackInterface* vvPos[fvecLen];

    for(unsigned short iTrP=0; iTrP < NPositive; iTrP += fvecLen)
    {
      int NTracks = (iTrP + fvecLen < NPositive) ? fvecLen : (NPositive - iTrP);

      for(unsigned short iv=0; iv<NTracks; iv++)
      {
        kfTrackPos[iv] = CbmKFTrack(vPos[iTrP+iv]);
        vvPos[iv] = &kfTrackPos[iv];
      }
      vDaughters[1].Create(vvPos,NTracks,0,&PiPlusPDG);
      vDaughters[1].SetField(posB[iTrP/fvecLen]);
#ifdef VC
      fvec posId(int_v::IndexesFromZero());
      posId += float(iTrP);
#else
      fvec posId(iTrP,iTrP+1,iTrP+2,iTrP+3);
#endif
      vDaughters[1].SetId(posId);

      CbmKFParticleInterface Ks;
      Ks.SetPDG( 310.f );
      Ks.Construct(vDaughters, 2, 0);

      for(int iv=0; iv<NTracks; iv++)
      {
        if(!finite(float(Ks.GetChi2()[iv]))) continue;
        if(!(Ks.GetChi2()[iv] > 0.0f)) continue;
        if(!(Ks.GetChi2()[iv]==Ks.GetChi2()[iv])) continue;
        if(Ks.GetZ()[iv] < cuts[0][2]) continue;

        if( Ks.GetChi2()[iv]/Ks.GetNDF()[iv] < cuts[0][1] )
        {
          CbmKFParticle Ks_temp;
          Ks.GetKFParticle(Ks_temp, iv);
          vK0s.push_back(Ks_temp);
        }
      }

      vDaughters[1].Create(vvPos,NTracks,0,&PPlusPDG);
      vDaughters[1].SetId(posId);

      CbmKFParticleInterface Lambda;
      Lambda.SetPDG(3122.f);
      Lambda.Construct(vDaughters, 2, 0);

      CbmKFParticleInterface LambdaTopo(Lambda);
      LambdaTopo.MeasureProductionVertex(LambdaTopo.KFPart, LambdaTopo.GetParameters(), &PrimVtx);

//       CbmKFParticleInterface LambdaTopo2;
//       LambdaTopo2.Construct(vDaughters, 2, &PrimVtx);

      for(int iv=0; iv<NTracks; iv++)
      {
        if(!finite(float(Lambda.GetChi2()[iv]))) continue;
        if(!(Lambda.GetChi2()[iv] > 0.0f)) continue;
        if(!(Lambda.GetChi2()[iv]==Lambda.GetChi2()[iv])) continue;
        if(Lambda.GetZ()[iv] < cuts[1][2]) continue;

        if( Lambda.GetChi2()[iv]/Lambda.GetNDF()[iv] < cuts[1][1])
        {
          CbmKFParticle Lambda_temp;
          Lambda.GetKFParticle(Lambda_temp, iv);
          vLambda.push_back(Lambda_temp);

          vLambdaTopoChi2Ndf.push_back(LambdaTopo.GetChi2()[iv]/LambdaTopo.GetNDF()[iv]);
        }
      }
    }
  }

  for(unsigned int iK0=0; iK0<vK0s.size(); iK0++)
  {
    vK0s[iK0].SetId(Particles.size());
    Particles.push_back(vK0s[iK0]);
  }

  for(unsigned int iLambda=0; iLambda<vLambda.size(); iLambda++)
  {
    vLambda[iLambda].SetId(Particles.size());
    Particles.push_back(vLambda[iLambda]);

    Double_t mass, errMass;

    vLambda[iLambda].GetMass(mass, errMass);
    if( (fabs(mass - massLambdaPDG)/0.0012) > 3 ) continue;

    if( vLambdaTopoChi2Ndf[iLambda] < 5 )
    {
      vLambdaPrim.push_back(vLambda[iLambda]);
    }
    else
    {
      if( vLambda[iLambda].GetZ() < 4 ) continue;
      vLambdaSec.push_back(vLambda[iLambda]);
    }
  }

  // Find Ksi-
  float cutKsi[3] = {3.,5.,6.};
  for(unsigned short iTrN=0; iTrN < vNeg.size(); iTrN++)
  {
    CbmKFTrack kfTrackNeg(vNeg[iTrN]);
    CbmKFParticle_simd vDaughters[2] = {CbmKFParticle_simd(),CbmKFParticle_simd(kfTrackNeg,0,&PiMinusPDG)};
    int entrSIMD = iTrN % fvecLen;
    int entrVec  = iTrN / fvecLen;

    vDaughters[1].SetField(negB[entrVec],1,entrSIMD);
    vDaughters[1].SetId(float(iTrN+NPositive));

    vector<int> daughterIds;
    daughterIds.push_back(iTrN+NPositive);
    FindHyperons(3312, vDaughters, daughterIds, vLambdaSec, vKsiMinus, PrimVtx, cutKsi);
  }

  for(unsigned int iKsi=0; iKsi<vKsiMinus.size(); iKsi++)
  {
    vKsiMinus[iKsi].SetId(Particles.size());
    Particles.push_back(vKsiMinus[iKsi]);
  }

  //Find Omega-
  float cutOmega[3] = {3.,3.,3.};
  for(unsigned short iTrN=0; iTrN < vNeg.size(); iTrN++)
  {
    if( ChiToPrimVtx[idNeg[iTrN]] < 7 ) continue;
    CbmKFTrack kfTrackNeg(vNeg[iTrN]);
    CbmKFParticle_simd vDaughters[2] = {CbmKFParticle_simd(),CbmKFParticle_simd(kfTrackNeg,0,&KMinusPDG)};
    int entrSIMD = iTrN % fvecLen;
    int entrVec  = iTrN / fvecLen;

    vDaughters[1].SetField(negB[entrVec],1,entrSIMD);
    vDaughters[1].SetId(float(iTrN+NPositive));

    vector<int> daughterIds;
    daughterIds.push_back(iTrN+NPositive);
    FindHyperons(3334, vDaughters, daughterIds, vLambdaSec, vOmegaMinus, PrimVtx, cutOmega);
  }

  for(unsigned int iOm=0; iOm<vOmegaMinus.size(); iOm++)
  {
    vOmegaMinus[iOm].SetId(Particles.size());
    Particles.push_back(vOmegaMinus[iOm]);
  }

  // Find Sigma*+
  float cutSigmaStarPlus[3] = {-100.,10000.,3.};
  for(unsigned short iTrP=0; iTrP < vPosPrim.size(); iTrP++)
  {
    CbmKFTrack kfTrackPos(vPosPrim[iTrP]);
    CbmKFParticle_simd vDaughters[2] = {CbmKFParticle_simd(),CbmKFParticle_simd(kfTrackPos,0,&PiMinusPDG)};
    int entrSIMD = iTrP % fvecLen;
    int entrVec  = iTrP / fvecLen;

    vDaughters[1].SetField(posPrimB[entrVec],1,entrSIMD);
    vDaughters[1].SetId(float(iTrP+NPositive+NNegative));

    vector<int> daughterIds;
    daughterIds.push_back(iTrP+NPositive+NNegative);
    FindHyperons(3224, vDaughters, daughterIds, vLambdaPrim, vSigmaStarPlus, PrimVtx, cutSigmaStarPlus);
  }

  for(unsigned int iSigma=0; iSigma<vSigmaStarPlus.size(); iSigma++)
  {
    vSigmaStarPlus[iSigma].SetId(Particles.size());
    Particles.push_back(vSigmaStarPlus[iSigma]);
  }

  // Find Sigma*-
  float cutSigmaStarMinus[3] = {-100.,10000.,3.};
  for(unsigned short iTrN=0; iTrN < vNegPrim.size(); iTrN++)
  {
    CbmKFTrack kfTrackNeg(vNegPrim[iTrN]);
    CbmKFParticle_simd vDaughters[2] = {CbmKFParticle_simd(),CbmKFParticle_simd(kfTrackNeg,0,&PiMinusPDG)};
    int entrSIMD = iTrN % fvecLen;
    int entrVec  = iTrN / fvecLen;

    vDaughters[1].SetField(negPrimB[entrVec],1,entrSIMD);
    vDaughters[1].SetId(float(iTrN+NPositive+NNegative+NPositivePrim));

    vector<int> daughterIds;
    daughterIds.push_back(iTrN+NPositive+NNegative+NPositivePrim);
    FindHyperons(3114, vDaughters, daughterIds, vLambdaPrim, vSigmaStarMinus, PrimVtx, cutSigmaStarMinus);
  }

  for(unsigned int iSigma=0; iSigma<vSigmaStarMinus.size(); iSigma++)
  {
    vSigmaStarMinus[iSigma].SetId(Particles.size());
    Particles.push_back(vSigmaStarMinus[iSigma]);
  }

  // Find H-dibarion
  float cutHdb[3] = {3.,3.,3.};
  for(unsigned short iL=0; iL < vLambdaSec.size(); iL++)
  {
    CbmKFParticle_simd vDaughters[2] = {CbmKFParticle_simd(),CbmKFParticle_simd(vLambdaSec[iL])};

    vector<int> daughterIds;
    for(unsigned int iD=0; iD<vLambdaSec[iL].DaughterIds().size(); iD++)
      daughterIds.push_back(vLambdaSec[iL].DaughterIds()[iD]);
    FindHyperons(3000, vDaughters, daughterIds, vLambdaSec, vHdibarion, PrimVtx, cutHdb, iL+1);
  }

  for(unsigned int iH=0; iH<vHdibarion.size(); iH++)
  {
    vHdibarion[iH].SetId(Particles.size());
    Particles.push_back(vHdibarion[iH]);
  }
}

void CbmKFParticleFinder::FindParticles(vector<CbmL1Track>& vRTracks, vector<CbmKFParticle>& Particles,
                                           CbmKFVertex& PrimVtx, const float cuts[2][3])
{
  FindParticlesT(vRTracks, Particles, PrimVtx, cuts);
}

void CbmKFParticleFinder::FindHyperons(int PDG,
                                          CbmKFParticle_simd vDaughters[2],
                                          vector<int>& daughterIds,
                                          vector<CbmKFParticle>& vLambdaSec,
                                          vector<CbmKFParticle>& vHyperon,
                                          CbmKFVertex& PrimVtx,
                                          const float *cuts,
                                          int startIndex)
{
  CbmKFParticle* lambdas[fvecLen];
  int nLambdasSec = vLambdaSec.size();

  for(unsigned short iL=startIndex; iL < vLambdaSec.size(); iL += fvecLen)
  {
    unsigned int nEntries = (iL + fvecLen < nLambdasSec) ? fvecLen : (nLambdasSec - iL);

    for(unsigned short iv=0; iv<nEntries; iv++)
      lambdas[iv] = &vLambdaSec[iL];

    vDaughters[0].Create(lambdas,nEntries);

    CbmKFParticleInterface Hyperon;

    Hyperon.SetPDG( float(PDG) );
    Hyperon.Construct(vDaughters, 2, 0);

    CbmKFParticleInterface HyperonTopo(Hyperon);
    HyperonTopo.MeasureProductionVertex(HyperonTopo.KFPart, HyperonTopo.GetParameters(), &PrimVtx);

    for(unsigned int iv=0; iv<nEntries; iv++)
    {
      bool isSameTrack = 0;
      for(unsigned short iD=0; iD<lambdas[iv]->DaughterIds().size(); iD++)
        for(unsigned short iD0=0; iD0<daughterIds.size(); iD0++)
          if(lambdas[iv]->DaughterIds()[iD] == daughterIds[iD0]) isSameTrack=1;

      if(isSameTrack) continue;
       if(!finite(float(Hyperon.GetChi2()[iv]))) continue;
      if(!(Hyperon.GetChi2()[iv] > 0.0f)) continue;
      if(!(Hyperon.GetChi2()[iv]==Hyperon.GetChi2()[iv])) continue;
      if((Hyperon.GetZ()[iv] < cuts[0]) || ((lambdas[iv]->GetZ() - Hyperon.GetZ()[iv]) < 0)) continue;

      if(HyperonTopo.GetChi2()[iv]/HyperonTopo.GetNDF()[iv] > cuts[1] ) continue;

      if( Hyperon.GetChi2()[iv]/Hyperon.GetNDF()[iv] > cuts[2] ) continue;
      CbmKFParticle Hyperon_temp;
      Hyperon.GetKFParticle(Hyperon_temp, iv);
      vHyperon.push_back(Hyperon_temp);
    }
  }
}

template<class T> 
void CbmKFParticleFinder::ConstructPVT(const vector<T>& vRTracks, CbmKFVertex& pv) const
{
  CbmKFTrackFitter fitter;
  fitter.SetHits(fHits);
  fitter.SetAlgo(fAlgo);

  const unsigned int NTracks = vRTracks.size();
  
  vector<L1FieldRegion> field;

  fitter.CalculateFieldRegion(vRTracks,field);

  vector<CbmKFParticle_simd> vTracks;
  vTracks.resize(NTracks);

  const Int_t PiMinusPDG =-211;

  for(unsigned int iTr=0; iTr<NTracks; ++iTr)
  {
    CbmKFTrack kfTrack(vRTracks[iTr]);

    CbmKFParticle_simd &t = vTracks[iTr];
    t = CbmKFParticle_simd(kfTrack,0,&PiMinusPDG);
    int entrSIMD = iTr % fvecLen;
    int entrVec  = iTr / fvecLen;
    t.SetField(field[entrVec],1,entrSIMD);
  }
  
  CbmKFParticleInterface pvSIMD;

  pvSIMD.SetVtxGuess(0.f,0.f,0.f);
  pvSIMD.SetVtxErrGuess(0.01f, 0.01f, 0.01f);
  pvSIMD.Construct(&(vTracks[0]), vTracks.size(), 0, -1, 3.5*3.5);

  pv.GetRefX() = pvSIMD.GetX()[0];
  pv.GetRefY() = pvSIMD.GetY()[0];
  pv.GetRefZ() = pvSIMD.GetZ()[0];

  for(int iC=0; iC<6; iC++)
    pv.GetCovMatrix()[iC] = pvSIMD.GetCovMatrix()[iC][0];

  pv.GetRefChi2() = pvSIMD.GetChi2()[0];
  pv.GetRefNDF() = pvSIMD.GetNDF()[0];
}

void CbmKFParticleFinder::ConstructPV(const vector<CbmL1Track>& vRTracks, CbmKFVertex& pv) const
{
  ConstructPVT(vRTracks, pv);
}

#undef cnst