@@ -17,15 +17,18 @@
 /// \author Mingze li <mingze.li@cern.ch>, CCNU/UniTo
 
 #include "PWGHF/Core/CentralityEstimation.h"
+#include "PWGHF/Core/DecayChannels.h"
 #include "PWGHF/D2H/Utils/utilsFlow.h"
 #include "PWGHF/DataModel/CandidateReconstructionTables.h"
 #include "PWGHF/DataModel/CandidateSelectionTables.h"
 
 #include "Common/Core/RecoDecay.h"
 #include "Common/DataModel/Centrality.h"
 #include "Common/DataModel/EventSelection.h"
+#include "Common/DataModel/Multiplicity.h"
 #include "Common/DataModel/Qvectors.h"
 
+#include <CommonConstants/MathConstants.h>
 #include <CommonConstants/PhysicsConstants.h>
 #include <Framework/ASoA.h>
 #include <Framework/AnalysisDataModel.h>
@@ -35,6 +38,16 @@
 #include <Framework/InitContext.h>
 #include <Framework/runDataProcessing.h>
 
+#include <Math/GenVector/Boost.h>
+#include <Math/GenVector/RotationZ.h>
+#include <Math/Vector3D.h> // IWYU pragma: keep (do not replace with Math/Vector3Dfwd.h)
+#include <Math/Vector3Dfwd.h>
+#include <Math/Vector4D.h> // IWYU pragma: keep (do not replace with Math/Vector4Dfwd.h)
+#include <Math/Vector4Dfwd.h>
+#include <TPDGCode.h>
+#include <TRandom3.h>
+
+#include <cmath>
 #include <cstdint>
 #include <vector>
 
@@ -44,6 +57,7 @@ using namespace o2::framework::expressions;
 using namespace o2::hf_centrality;
 using namespace o2::hf_occupancy;
 using namespace o2::analysis::hf_flow_utils;
+using namespace o2::constants::physics;
 
 namespace o2::aod
 {
@@ -70,6 +84,7 @@ DECLARE_SOA_COLUMN(Pt, pt, float);
 DECLARE_SOA_COLUMN(Eta, eta, float);
 DECLARE_SOA_COLUMN(Phi, phi, float);
 DECLARE_SOA_COLUMN(Y, y, float);
+DECLARE_SOA_COLUMN(CosThetaStarEP, cosThetaStarEP, float);
 DECLARE_SOA_COLUMN(MlProbDstarToD0PiBkg, mlProbDstarToD0PiBkg, float);
 DECLARE_SOA_COLUMN(MlProbDstarToD0PiPrompt, mlProbDstarToD0PiPrompt, float);
 DECLARE_SOA_COLUMN(MlProbDstarToD0PiNonPrompt, mlProbDstarToD0PiNonPrompt, float);
@@ -112,24 +127,45 @@ DECLARE_SOA_TABLE(HfCandDstMix, "AOD", "HFCANDDSTMIX",
                   mixing_dstar::MinNumTpcCls,
                   mixing_dstar::GIndexCol,
                   mixing_dstar::TimeStamp);
+DECLARE_SOA_TABLE(HfMcParticles, "AOD", "HFMCPARTICLES",
+                  mixing_dstar::Pt,
+                  mixing_dstar::Eta,
+                  mixing_dstar::Phi,
+                  mixing_dstar::Y,
+                  mixing_dstar::CosThetaStarEP,
+                  mixing_dstar::Centrality);
 } // namespace o2::aod
 
 /// Writes the full information in an output TTree
 struct HfTreeCreatorDstarSpinAlignMixing {
   Produces<o2::aod::HfCandDstMix> rowCandidateMix;
+  Produces<o2::aod::HfMcParticles> rowGenMcParticles;
 
   Configurable<int> qVecDetector{"qVecDetector", 2, "Detector for Q vector estimation (FV0A: 0, FT0M: 1, FT0C: 2)"};
   Configurable<int> centEstimator{"centEstimator", 2, "Centrality estimator ((None: 0, FT0C: 2, FT0M: 3))"};
   Configurable<int> occEstimator{"occEstimator", 2, "If enabled, replace number of PV contributors with occupancy estimation (0: don't use, 1: ITS, 2: FT0C)"};
+  Configurable<bool> activateAxisRotation{"activateAxisRotation", false, "Activate the rotation of the quantisation axis for cos(Theta*) calculation for phi dependence studies"};
+  Configurable<float> axisRotationAngle{"axisRotationAngle", 0.f, "Rotation angle of the quantisation axis for cos(Theta*) calculation for phi dependence studies (in rad)"};
+  Configurable<bool> activatePartRecoDstar{"activatePartRecoDstar", false, "Activate the study of partly reconstructed D*+ -> D0 (-> KPiPi0) Pi decays"};
 
   using CollsWithQVecs = soa::Join<aod::Collisions, aod::EvSels, aod::QvectorFT0Cs, aod::QvectorFT0As, aod::QvectorFT0Ms, aod::QvectorFV0As, aod::QvectorBPoss, aod::QvectorBNegs, aod::QvectorBTots, aod::CentFT0Ms, aod::CentFT0Cs>;
   using TracksWithExtra = soa::Join<aod::Tracks, aod::TracksExtra>;
   using CandDstarWSelFlag = soa::Join<aod::HfCandDstars, aod::HfSelDstarToD0Pi>;
+
+  using CollisionsWithMcLabels = soa::SmallGroups<soa::Join<aod::Collisions, aod::EvSels, aod::PVMults, aod::McCollisionLabels>>;
+  using CollisionsWithMcLabelsAndCent = soa::SmallGroups<soa::Join<aod::Collisions, aod::EvSels, aod::PVMults, aod::McCollisionLabels, aod::CentFT0Ms, aod::CentFT0Cs>>;
+  using McParticlesDstarMatched = soa::Join<aod::McParticles, aod::HfCandDstarMcGen>;
+  using BCsInfo = soa::Join<aod::BCs, aod::Timestamps, aod::BcSels>;
+
   using FilteredCandDstarWSelFlagAndMl = soa::Filtered<soa::Join<CandDstarWSelFlag, aod::HfMlDstarToD0Pi>>;
+  using FilteredCandDstarWSelFlagAndMcAndMl = soa::Filtered<soa::Join<CandDstarWSelFlag, aod::HfMlDstarToD0Pi, aod::HfCandDstarMcRec>>;
 
   Filter filterSelectDstarCandidates = aod::hf_sel_candidate_dstar::isSelDstarToD0Pi == true;
 
   Preslice<FilteredCandDstarWSelFlagAndMl> dstarWithMlPerCollision = aod::hf_cand::collisionId;
+  Preslice<FilteredCandDstarWSelFlagAndMcAndMl> dstarWithMcAndMlPerCollision = aod::hf_cand::collisionId;
+
+  PresliceUnsorted<aod::McCollisionLabels> colPerMcCollision = aod::mcparticle::mcCollisionId;
 
   void init(InitContext const&)
   {
@@ -159,9 +195,72 @@ struct HfTreeCreatorDstarSpinAlignMixing {
     }
   }
 
-  template <typename CollType, typename T, typename Trk, typename BcType>
+  /// \param ptCharmHad is the pt of the candidate
+  double sampleDeltaPhi(float ptCharmHad)
+  {
+
+    double y0 = 6.49e-2;
+    double cons = 8.42e-1;
+    double coef = -3.78e-1;
+    double x0 = -6.37e-1;
+    double norm = -3.64e-1;
+    double mu = 4.68e-2;
+    double sigma = 1.513;
+
+    double sigmoid = y0 + (cons - y0) / (1 + std::exp(-coef * (ptCharmHad - x0)));
+    double peak = norm * std::exp(-(ptCharmHad - mu) * (ptCharmHad - mu) / (2 * sigma * sigma));
+
+    double v2 = sigmoid + peak;
+    const double funcMax = 1.0 + 2.0 * std::abs(v2);
+
+    while (true) {
+      double deltaPhi = gRandom->Uniform(0, constants::math::TwoPI);
+      double funcVal = gRandom->Uniform(0, funcMax);
+
+      double funcDeltaPhi = 1.0 + 2.0 * v2 * std::cos(2.0 * deltaPhi);
+
+      if (funcVal < funcDeltaPhi)
+        return deltaPhi;
+    }
+    return -1.;
+  }
+
+  /// \param daughter is the daughter particle
+  /// \param ptMin is the minimum pt
+  /// \param etaMax is the maximum eta
+  /// \return true if daughter is in acceptance
+  template <typename Part>
+  bool isDaughterInAcceptance(Part const& daughter, float ptMin, float etaMax)
+  {
+    if (daughter.pt() < ptMin) {
+      return false;
+    }
+    if (std::abs(daughter.eta()) > etaMax) {
+      return false;
+    }
+
+    return true;
+  }
+
+  template <bool DoMc, typename CollType, typename T, typename Trk, typename BcType>
   void fillCandidateTable(CollType const& collision, const T& candidate, Trk const& /*tracks*/, BcType const& /*bcs*/)
   {
+    int8_t origin{RecoDecay::OriginType::None};
+    bool partRecoDstar{false};
+    if constexpr (DoMc) {
+      partRecoDstar = std::abs(candidate.flagMcMatchRec()) == hf_decay::hf_cand_dstar::DecayChannelMain::DstarToPiKPiPi0 && std::abs(candidate.flagMcMatchRecD0()) == hf_decay::hf_cand_2prong::DecayChannelMain::D0ToPiKPi0;
+      bool const signalDstar = std::abs(candidate.flagMcMatchRec()) == hf_decay::hf_cand_dstar::DecayChannelMain::DstarToPiKPi && std::abs(candidate.flagMcMatchRecD0()) == hf_decay::hf_cand_2prong::DecayChannelMain::D0ToPiK;
+      if (!signalDstar && (!partRecoDstar || !activatePartRecoDstar)) { // this candidate is not signal and not partially reconstructed signal, skip
+        return;
+      }
+      origin = candidate.originMcRec();
+      int const pdgBhadMother = candidate.pdgBhadMotherPart();
+      // For unknown reasons there are charm hadrons coming directly from beauty diquarks without an intermediate B-hadron which have an unreasonable correlation between the pT of the charm hadron and the beauty mother. We also remove charm hadrons from quarkonia.
+      if (origin == RecoDecay::OriginType::NonPrompt && (pdgBhadMother == 5101 || pdgBhadMother == 5103 || pdgBhadMother == 5201 || pdgBhadMother == 5203 || pdgBhadMother == 5301 || pdgBhadMother == 5303 || pdgBhadMother == 5401 || pdgBhadMother == 5403 || pdgBhadMother == 5503 || pdgBhadMother == 553 || pdgBhadMother == 555 || pdgBhadMother == 557)) { // o2-linter: disable=pdg/explicit-code, magic-number (constants not in the PDG header)
+        return;
+      }
+    }
+
     const auto bc = collision.template bc_as<BcType>();
     const int64_t timeStamp = bc.timestamp();
 
@@ -186,7 +285,12 @@ struct HfTreeCreatorDstarSpinAlignMixing {
     auto trackProng1 = candidate.template prong1_as<Trk>();
     auto trackProng2 = candidate.template prongPi_as<Trk>();
     getTrackingInfos(std::vector{trackProng0, trackProng1, trackProng2}, absEtaTrackMin, numItsClsMin, numTpcClsMin);
-    std::array<float, 3> const Qvector = getQvec(collision, qVecDetector.value);
+    std::array<float, 3> Qvector = getQvec(collision, qVecDetector.value);
+    if constexpr (DoMc) {
+      double const deltaPhi = sampleDeltaPhi(candidate.pt());
+      double psi = candidate.phi() - deltaPhi;
+      Qvector = {static_cast<float>(-std::sin(psi)), static_cast<float>(std::cos(psi)), 0.f};
+    }
 
     rowCandidateMix(
       massD0,
@@ -214,6 +318,90 @@ struct HfTreeCreatorDstarSpinAlignMixing {
       timeStamp);
   }
 
+  /// \param mcParticle is the Mc particle
+  /// \param mcParticles is the table of Mc particles
+  /// \param centrality is the centrality of the collision
+  template <typename Part, typename Particles>
+  void fillGenParticleTable(Part const& mcParticle, Particles const& mcParticles, float const centrality)
+  {
+    int8_t origin{RecoDecay::OriginType::None};
+    bool partRecoDstar{false};
+    partRecoDstar = (std::abs(mcParticle.flagMcMatchGen()) == hf_decay::hf_cand_dstar::DecayChannelMain::DstarToPiKPiPi0) && (std::abs(mcParticle.flagMcMatchGenD0()) == hf_decay::hf_cand_2prong::DecayChannelMain::D0ToPiKPi0);
+    bool const signalDstar = (std::abs(mcParticle.flagMcMatchGen()) == hf_decay::hf_cand_dstar::DecayChannelMain::DstarToPiKPi) && (std::abs(mcParticle.flagMcMatchGenD0()) == hf_decay::hf_cand_2prong::DecayChannelMain::D0ToPiK);
+
+    if (!signalDstar && (!activatePartRecoDstar || !partRecoDstar)) { // this particle is not signal and not partially reconstructed signal, skip
+      return;
+    }
+    origin = mcParticle.originMcGen();
+    float massDau{0.f}, massCharmHad{0.f};
+    if (origin == RecoDecay::OriginType::NonPrompt) {
+      auto bHadMother = mcParticles.rawIteratorAt(mcParticle.idxBhadMotherPart() - mcParticles.offset());
+      int const pdgBhadMother = std::abs(bHadMother.pdgCode());
+      // For unknown reasons there are charm hadrons coming directly from beauty diquarks without an intermediate B-hadron which have an unreasonable correlation between the pT of the charm hadron and the beauty mother. We also remove charm hadrons from quarkonia.
+      if (pdgBhadMother == 5101 || pdgBhadMother == 5103 || pdgBhadMother == 5201 || pdgBhadMother == 5203 || pdgBhadMother == 5301 || pdgBhadMother == 5303 || pdgBhadMother == 5401 || pdgBhadMother == 5403 || pdgBhadMother == 5503 || pdgBhadMother == 553 || pdgBhadMother == 555 || pdgBhadMother == 557) { // o2-linter: disable=pdg/explicit-code, magic-number (constants not in the PDG header)
+        return;
+      }
+    }
+    std::vector<int> listDaughters{};
+    std::array<int, 2> const dauPdgs = {kPiPlus, o2::constants::physics::Pdg::kD0};
+    RecoDecay::getDaughters(mcParticle, &listDaughters, dauPdgs, 1);
+    massDau = MassPiPlus;
+    massCharmHad = MassDStar;
+
+    float const rapidity = mcParticle.y();
+    if (std::abs(rapidity) > 1.f) { // we do not keep particles with |y| > 1
+      return;
+    }
+
+    bool areDauInAcc{true};
+    float pxDau{-1000.f}, pyDau{-1000.f}, pzDau{-1000.f};
+    for (const auto& dauIdx : listDaughters) {
+      auto dauPart = mcParticles.rawIteratorAt(dauIdx - mcParticles.offset());
+      if (std::abs(dauPart.pdgCode()) == kPiPlus) {
+        pxDau = dauPart.px();
+        pyDau = dauPart.py();
+        pzDau = dauPart.pz();
+        if (areDauInAcc) {
+          areDauInAcc = isDaughterInAcceptance(dauPart, 0.1, 0.8);
+        }
+      } else if (areDauInAcc) { // check also D0 daughters
+        std::vector<int> listDaughtersD0{};
+        std::array<int, 2> const dauPdgsD0 = {kPiPlus, -kKPlus};
+        RecoDecay::getDaughters(mcParticle, &listDaughtersD0, dauPdgsD0, 1);
+        for (const auto& dauIdxD0 : listDaughtersD0) {
+          auto dauPartD0 = mcParticles.rawIteratorAt(dauIdxD0 - mcParticles.offset());
+          areDauInAcc = isDaughterInAcceptance(dauPartD0, 0.3, 0.8);
+        }
+      }
+    }
+
+    float const pxCharmHad = mcParticle.px();
+    float const pyCharmHad = mcParticle.py();
+    float const pzCharmHad = mcParticle.pz();
+    ROOT::Math::PxPyPzMVector const fourVecDau = ROOT::Math::PxPyPzMVector(pxDau, pyDau, pzDau, massDau);
+    ROOT::Math::PxPyPzMVector const fourVecMother = ROOT::Math::PxPyPzMVector(pxCharmHad, pyCharmHad, pzCharmHad, massCharmHad);
+    ROOT::Math::Boost const boost{fourVecMother.BoostToCM()};
+    ROOT::Math::PxPyPzMVector const fourVecDauCM = boost(fourVecDau);
+    ROOT::Math::XYZVector const threeVecDauCM = fourVecDauCM.Vect();
+
+    // EP analysis for MC
+    double const deltaPhi = sampleDeltaPhi(mcParticle.pt());
+    double psi = mcParticle.phi() - deltaPhi;
+    ROOT::Math::XYZVector qVecNorm = ROOT::Math::XYZVector(-std::sin(psi), std::cos(psi), 0.f);
+    if (activateAxisRotation) {
+      qVecNorm = ROOT::Math::RotationZ(axisRotationAngle)(qVecNorm);
+    }
+    float const cosThetaStarEP = qVecNorm.Dot(threeVecDauCM) / std::sqrt(threeVecDauCM.Mag2());
+
+    rowGenMcParticles(
+      mcParticle.pt(),
+      mcParticle.eta(),
+      mcParticle.phi(),
+      mcParticle.y(),
+      cosThetaStarEP,
+      centrality);
+  }
+
   void processData(CollsWithQVecs const& collisions,
                    FilteredCandDstarWSelFlagAndMl const& dstarCandidates,
                    TracksWithExtra const& tracks,
@@ -223,11 +411,35 @@ struct HfTreeCreatorDstarSpinAlignMixing {
       auto thisCollId = collision.globalIndex();
       auto groupedDstarCandidates = dstarCandidates.sliceBy(dstarWithMlPerCollision, thisCollId);
       for (const auto& dstarCandidate : groupedDstarCandidates) {
-        fillCandidateTable(collision, dstarCandidate, tracks, bcWithTimeStamps);
+        fillCandidateTable<false>(collision, dstarCandidate, tracks, bcWithTimeStamps);
       }
     }
   }
   PROCESS_SWITCH(HfTreeCreatorDstarSpinAlignMixing, processData, "Process data", true);
+
+  void processDstarMcWithMlInPbPb(aod::McCollisions::iterator const&,
+                                  McParticlesDstarMatched const& mcParticles,
+                                  CollisionsWithMcLabelsAndCent const& collisions, // this is grouped with SmallGroupsCollisionsWithMcLabels const& collisions,
+                                  FilteredCandDstarWSelFlagAndMcAndMl const& dstarCandidates,
+                                  TracksWithExtra const& tracks,
+                                  BCsInfo const& bcWithTimeStamps)
+  {
+    for (const auto& collision : collisions) { // loop over reco collisions associated to this gen collision
+      auto thisCollId = collision.globalIndex();
+      auto groupedDstarCandidates = dstarCandidates.sliceBy(dstarWithMcAndMlPerCollision, thisCollId);
+      for (const auto& dstarCandidate : groupedDstarCandidates) {
+        fillCandidateTable<true>(collision, dstarCandidate, tracks, bcWithTimeStamps);
+      }
+    }
+    rowGenMcParticles.reserve(mcParticles.size());
+
+    for (const auto& mcParticle : mcParticles) {
+      const auto& recoCollsPerMcColl = collisions.sliceBy(colPerMcCollision, mcParticle.mcCollision().globalIndex());
+      const auto centrality = getCentralityGenColl(recoCollsPerMcColl, centEstimator);
+      fillGenParticleTable(mcParticle, mcParticles, centrality);
+    }
+  }
+  PROCESS_SWITCH(HfTreeCreatorDstarSpinAlignMixing, processDstarMcWithMlInPbPb, "Process Dstar candidates in PbPb MC with ML", false);
 };
 
 WorkflowSpec defineDataProcessing(ConfigContext const& cfgc)