diff --git a/PWGJE/Core/JetTaggingUtilities.h b/PWGJE/Core/JetTaggingUtilities.h
index d6a4ee8f32b..7ed9d0b2db7 100644
--- a/PWGJE/Core/JetTaggingUtilities.h
+++ b/PWGJE/Core/JetTaggingUtilities.h
@@ -626,6 +626,182 @@ bool isTaggedJetSV(T const jet, U const& /*prongs*/, float const& prongChi2PCAMi
   return true;
 }
 
+/**
+ * Clusters jet constituent tracks into groups of tracks originating from same mcParticle position (trkVtxIndex), and finds each track origin (trkOrigin). (for GNN b-jet tagging)
+ * @param trkLabels Track labels for GNN vertex and track origin predictions. trkVtxIndex: The index value of each vertex (cluster) which is determined by the function. trkOrigin: The category of the track origin (0: not physical primary, 1: charm, 2: beauty, 3: primary vertex, 4: other secondary vertex).
+ * @param vtxResParam Vertex resolution parameter which determines the cluster size. (cm)
+ * @param trackPtMin Minimum value of track pT.
+ * @return The number of vertices (clusters) in the jet.
+ */
+template <typename AnyCollision, typename AnalysisJet, typename AnyTracks, typename AnyParticles, typename AnyOriginalParticles>
+int vertexClustering(AnyCollision const& collision, AnalysisJet const& jet, AnyTracks const&, AnyParticles const& particles, AnyOriginalParticles const&, std::unordered_map<std::string, std::vector<int>>& trkLabels, float vtxResParam = 0.01 /* 0.01cm = 100um */, float trackPtMin = 0.5)
+{
+  const auto& tracks = jet.template tracks_as<AnyTracks>();
+  const int n_trks = tracks.size();
+
+  // trkVtxIndex
+
+  std::vector<int> tempTrkVtxIndex;
+
+  int i=0;
+  for (const auto& constituent : tracks) {
+    if (!constituent.has_mcParticle() || !constituent.template mcParticle_as<AnyParticles>().isPhysicalPrimary() || constituent.pt() < trackPtMin)
+      tempTrkVtxIndex.push_back(-1);
+    else
+      tempTrkVtxIndex.push_back(i++);
+  }
+  tempTrkVtxIndex.push_back(i); // temporary index for PV
+  if (n_trks < 1) { // the process should be done for n_trks == 1 as well
+    trkLabels["trkVtxIndex"] = tempTrkVtxIndex;
+    return n_trks;
+  }
+
+  int n_pos = n_trks + 1;
+  std::vector<float> dists(n_pos * (n_pos - 1) / 2);
+  auto trk_pair_idx = [n_pos](int ti, int tj) {
+    if (ti==tj || ti>=n_pos || tj>=n_pos || ti<0 || tj<0) {
+      LOGF(info,"Track pair index out of range");
+      return -1;
+    }
+    else
+      return (ti < tj) ? (ti * n_pos - (ti * (ti + 1)) / 2 + tj - ti - 1) : (tj * n_pos - (tj * (tj + 1)) / 2 + ti - tj - 1);
+  }; // index n_trks is for PV
+
+  for (int ti=0; ti<n_pos-1; ti++)
+    for (int tj=ti+1; tj<n_pos; tj++) {
+      std::array<float, 3> posi, posj;
+      
+      if (tj < n_trks) {
+        if (tracks[tj].has_mcParticle()) {
+          const auto& pj = tracks[tj].template mcParticle_as<AnyParticles>().template mcParticle_as<AnyOriginalParticles>();
+          posj = std::array<float, 3>{pj.vx(), pj.vy(), pj.vz()};
+        }
+        else {
+          dists[trk_pair_idx(ti, tj)] = std::numeric_limits<float>::max();
+          continue;
+        }
+      }
+      else {
+        posj = std::array<float, 3>{collision.posX(), collision.posY(), collision.posZ()};
+      }
+
+      if (tracks[ti].has_mcParticle()) {
+        const auto& pi = tracks[ti].template mcParticle_as<AnyParticles>().template mcParticle_as<AnyOriginalParticles>();
+        posi = std::array<float, 3>{pi.vx(), pi.vy(), pi.vz()};
+      }
+      else {
+        dists[trk_pair_idx(ti, tj)] = std::numeric_limits<float>::max();
+        continue;
+      }
+
+      dists[trk_pair_idx(ti, tj)] = RecoDecay::distance(posi, posj);
+    }
+
+  int clusteri = -1, clusterj = -1;
+  float min_min_dist = -1.f; // If there is an not-merge-able min_dist pair, check the 2nd-min_dist pair. 
+  while (true) {
+
+    float min_dist = -1.f; // Get min_dist pair
+    for (int ti=0; ti<n_pos-1; ti++)
+      for (int tj=ti+1; tj<n_pos; tj++)
+        if (tempTrkVtxIndex[ti] != tempTrkVtxIndex[tj] && tempTrkVtxIndex[ti]>=0 && tempTrkVtxIndex[tj]>=0) {
+          float dist = dists[trk_pair_idx(ti, tj)];
+          if ((dist < min_dist || min_dist < 0.f) && dist > min_min_dist) {
+            min_dist = dist;
+            clusteri = ti;
+            clusterj = tj;
+          }
+        }
+    if (clusteri < 0 || clusterj < 0)
+      break;
+    
+    bool mrg = true; // Merge-ability check
+    for (int ti=0; ti<n_pos && mrg; ti++)
+      if (tempTrkVtxIndex[ti] == tempTrkVtxIndex[clusteri] && tempTrkVtxIndex[ti]>=0)
+        for (int tj=0; tj<n_pos && mrg; tj++)
+          if (tj != ti && tempTrkVtxIndex[tj] == tempTrkVtxIndex[clusterj] && tempTrkVtxIndex[tj]>=0)
+            if (dists[trk_pair_idx(ti, tj)] > vtxResParam) { // If there is more distant pair compared to vtx_res between two clusters, they cannot be merged.
+              mrg = false;
+              min_min_dist = min_dist;
+            }
+    if (min_dist > vtxResParam || min_dist < 0.f)
+      break;
+    
+    if (mrg) { // Merge two clusters
+      int old_index = tempTrkVtxIndex[clusterj];
+      for (int t=0; t<n_pos; t++)
+        if (tempTrkVtxIndex[t] == old_index)
+          tempTrkVtxIndex[t] = tempTrkVtxIndex[clusteri];
+    }
+  }
+
+  int n_vertices = 0;
+
+  // Sort the indices from PV (as 0) to the most distant SV (as 1~).
+  int idxPV = tempTrkVtxIndex[n_trks];
+  for (int t=0; t<n_trks; t++)
+    if (tempTrkVtxIndex[t] == idxPV) {
+      tempTrkVtxIndex[t] = -2;
+      n_vertices = 1; // There is a track originating from PV
+    }
+  
+  std::unordered_map<int, float> avgDistances;
+  std::unordered_map<int, int> count;
+  for (int t=0; t<n_trks; t++) {
+    if (tempTrkVtxIndex[t] >= 0) {
+      avgDistances[tempTrkVtxIndex[t]] += dists[trk_pair_idx(t, n_trks)];
+      count[tempTrkVtxIndex[t]]++;
+    }
+  }
+  
+  trkLabels["trkVtxIndex"] = std::vector<int>(n_trks, -1);
+  if (count.size() != 0) { // If there is any SV cluster not only PV cluster
+    for (auto& [idx, avgDistance] : avgDistances)
+      avgDistance /= count[idx];
+
+    n_vertices += avgDistances.size();
+    
+    std::vector<std::pair<int, float>> sortedIndices(avgDistances.begin(), avgDistances.end());
+    std::sort(sortedIndices.begin(), sortedIndices.end(), [](const auto& a, const auto& b) { return a.second < b.second; });
+    int rank = 1;
+    for (const auto& [idx, avgDistance] : sortedIndices) {
+      bool found = false;
+      for (int t=0; t<n_trks; t++)
+        if (tempTrkVtxIndex[t] == idx) {
+          trkLabels["trkVtxIndex"][t] = rank;
+          found = true;
+        }
+      rank += found;
+    }
+  }
+  
+  for (int t=0; t<n_trks; t++)
+    if (tempTrkVtxIndex[t] == -2)
+      trkLabels["trkVtxIndex"][t] = 0;
+
+  // trkOrigin
+
+  int trkIdx = 0;
+  for (auto &constituent : jet.template tracks_as<AnyTracks>())
+  {
+    if (!constituent.has_mcParticle() || !constituent.template mcParticle_as<AnyParticles>().isPhysicalPrimary() || constituent.pt() < trackPtMin)
+    {
+      trkLabels["trkOrigin"].push_back(0);
+    }
+    else
+    {
+      const auto &particle = constituent.template mcParticle_as<AnyParticles>();
+      int orig = RecoDecay::getParticleOrigin(particles, particle, true);
+      trkLabels["trkOrigin"].push_back((orig > 0) ? orig :
+                                       (trkLabels["trkVtxIndex"][trkIdx] == 0) ? 3 : 4);
+    }
+
+    trkIdx++;
+  }
+
+  return n_vertices;
+}
+
 }; // namespace jettaggingutilities
 
 #endif // PWGJE_CORE_JETTAGGINGUTILITIES_H_
diff --git a/PWGJE/Tasks/bjetTreeCreator.cxx b/PWGJE/Tasks/bjetTreeCreator.cxx
index 2d60344085a..a415c92ebfb 100644
--- a/PWGJE/Tasks/bjetTreeCreator.cxx
+++ b/PWGJE/Tasks/bjetTreeCreator.cxx
@@ -81,6 +81,15 @@ DECLARE_SOA_COLUMN(SignedIP3D, ip3d, float);                           //! The t
 DECLARE_SOA_COLUMN(SignedIP3DSign, ip3dsigma, float);                  //! The track signed 3D IP significance
 DECLARE_SOA_COLUMN(MomFraction, momfraction, float);                   //! The track momentum fraction of the jets
 DECLARE_SOA_COLUMN(DeltaRTrackVertex, rtrackvertex, float);            //! DR between the track and the closest SV, to be decided whether to add to or not
+DECLARE_SOA_COLUMN(TrackPhi, trackphi, float);                         //! The track phi
+DECLARE_SOA_COLUMN(TrackCharge, trackcharge, float);                   //! The track sign (charge)
+DECLARE_SOA_COLUMN(TrackITSChi2NCl, trackitschi2ncl, float);           //! The track ITS Chi2NCl
+DECLARE_SOA_COLUMN(TrackTPCChi2NCl, tracktpcchi2ncl, float);           //! The track TPC Chi2NCl
+DECLARE_SOA_COLUMN(TrackITSNCls, trackitsncls, float);                 //! The track ITS NCls
+DECLARE_SOA_COLUMN(TrackTPCNCls, tracktpcncls, float);                 //! The track TPC NCls (Found)
+DECLARE_SOA_COLUMN(TrackTPCNCrossedRows, tracktpcncrossedrows, float); //! The track TPC NCrossedRows
+DECLARE_SOA_COLUMN(TrackOrigin, trk_origin, int);                      //! The track origin label for GNN track origin predictions
+DECLARE_SOA_COLUMN(TrackVtxIndex, trk_vtx_index, int);                 //! The track vertex index for GNN vertex predictions
 // DECLARE_SOA_COLUMN(DCATrackJet, dcatrackjet, float);                              //! The distance between track and jet, unfortunately it cannot be calculated in O2
 } // namespace trackInfo
 
@@ -101,6 +110,20 @@ DECLARE_SOA_TABLE(bjetTracksParams, "AOD", "BJETTRACKSPARAM",
 
 using bjetTracksParam = bjetTracksParams::iterator;
 
+DECLARE_SOA_TABLE(bjetTracksParamsExtra, "AOD", "BJETTRACKSEXTRA",
+                  // o2::soa::Index<>,
+                  trackInfo::TrackPhi,
+                  trackInfo::TrackCharge,
+                  trackInfo::TrackITSChi2NCl,
+                  trackInfo::TrackTPCChi2NCl,
+                  trackInfo::TrackITSNCls,
+                  trackInfo::TrackTPCNCls,
+                  trackInfo::TrackTPCNCrossedRows,
+                  trackInfo::TrackOrigin,
+                  trackInfo::TrackVtxIndex);
+
+using bjetTracksParamExtra = bjetTracksParamsExtra::iterator;
+
 namespace SVInfo
 {
 DECLARE_SOA_INDEX_COLUMN(bjetParam, jetindex);            //! The jet index
@@ -154,6 +177,7 @@ struct BJetTreeCreator {
 
   Produces<aod::bjetParams> bjetParamsTable;
   Produces<aod::bjetTracksParams> bjetTracksParamsTable;
+  Produces<aod::bjetTracksParamsExtra> bjetTracksExtraTable;
   Produces<aod::bjetSVParams> bjetSVParamsTable;
   Produces<aod::bjetConstituents> bjetConstituentsTable;
 
@@ -196,6 +220,8 @@ struct BJetTreeCreator {
 
   Configurable<bool> produceTree{"produceTree", true, "produce the jet TTree"};
 
+  Configurable<float> vtxRes{"vtxRes", 0.01, "Vertex position resolution (cluster size) for GNN vertex predictions (cm)"};
+
   int eventSelection = -1;
 
   std::vector<double> jetRadiiValues;
@@ -225,7 +251,7 @@ struct BJetTreeCreator {
     registry.add("h2_jetMass_jetpT", "Jet mass;#it{p}_{T,jet} (GeV/#it{c});#it{m}_{jet} (GeV/#it{c}^{2})", {HistType::kTH2F, {{200, 0., 200.}, {50, 0, 50.0}}});
     registry.add("h2_SVMass_jetpT", "Secondary vertex mass;#it{p}_{T,jet} (GeV/#it{c});#it{m}_{SV} (GeV/#it{c}^{2})", {HistType::kTH2F, {{200, 0., 200.}, {50, 0, 10}}});
 
-    if (doprocessMCJets) {
+    if (doprocessMCJets || doprocessMCJetsForGNN) {
       registry.add("h2_SIPs2D_jetpT_bjet", "2D IP significance b-jets;#it{p}_{T,jet} (GeV/#it{c});IPs", {HistType::kTH2F, {{200, 0., 200.}, {100, -50.0, 50.0}}});
       registry.add("h2_SIPs3D_jetpT_bjet", "3D IP significance b-jets;#it{p}_{T,jet} (GeV/#it{c});IPs", {HistType::kTH2F, {{200, 0., 200.}, {100, -50.0, 50.0}}});
       registry.add("h2_LxyS_jetpT_bjet", "Decay length in XY b-jets;#it{p}_{T,jet} (GeV/#it{c});S#it{L}_{xy}", {HistType::kTH2F, {{200, 0., 200.}, {100, 0., 100.0}}});
@@ -259,6 +285,47 @@ struct BJetTreeCreator {
       registry.add("h2_Response_DetjetpT_PartjetpT_cjet", "Response matrix c-jets;#it{p}_{T,jet}^{det} (GeV/#it{c});#it{p}_{T,jet}^{part} (GeV/#it{c})", {HistType::kTH2F, {{200, 0., 200.}, {200, 0., 200.}}});
       registry.add("h2_Response_DetjetpT_PartjetpT_lfjet", "Response matrix lf-jet;#it{p}_{T,jet}^{det} (GeV/#it{c});#it{p}_{T,jet}^{part} (GeV/#it{c})", {HistType::kTH2F, {{200, 0., 200.}, {200, 0., 200.}}});
     }
+
+    if (doprocessMCJetsForGNN)
+    {
+      //+jet
+      registry.add("h_jet_pt", "jet_pt;#it{p}_{T}^{ch jet} (GeV/#it{c});Entries", {HistType::kTH1F, {{100, 0., 200.}}});
+      registry.add("h_jet_eta", "jet_eta;#it{#eta}_{ch jet};Entries", {HistType::kTH1F, {{100, -2., 2.}}});
+      registry.add("h_jet_phi", "jet_phi;#it{#phi}_{ch jet};Entries", {HistType::kTH1F, {{100, 0., 2. * M_PI}}});
+      registry.add("h_jet_flav", "jet_flav;jet flavor;Entries", {HistType::kTH1F, {{4, 0., 4.}}});
+      registry.add("h_n_trks", "n_trks;#it{n}_{tracks};Entries", {HistType::kTH1F, {{50, 0., 50.}}});
+      registry.add("h_jet_mass", "jet_mass;#it{m}_{jet} (GeV/#it{c}^2);Entries", {HistType::kTH1F, {{100, 0., 50.}}});
+      auto h_jet_flav = registry.get<TH1>(HIST("h_jet_flav"));
+      h_jet_flav->GetXaxis()->SetBinLabel(1, "no mcparticle"); // 0
+      h_jet_flav->GetXaxis()->SetBinLabel(2, "c-jet");         // 1
+      h_jet_flav->GetXaxis()->SetBinLabel(3, "b-jet");         // 2
+      h_jet_flav->GetXaxis()->SetBinLabel(4, "lf-jet");        // 3
+      registry.add("h_n_vertices", "n_vertices;#it{n}_{vertex};Entries", {HistType::kTH1F, {{50, 0., 50.}}});
+      //+trk
+      registry.add("h_trk_pt", "trk_pt;#it{p}_{T} (GeV/#it{c});Entries", {HistType::kTH1F, {{100, 0., 100.}}});
+      registry.add("h_trk_eta", "trk_eta;#it{#eta};Entries", {HistType::kTH1F, {{100, -2., 2.}}});
+      registry.add("h_trk_phi", "trk_phi;#it{#phi};Entries", {HistType::kTH1F, {{100, 0., 2. * M_PI}}});
+      registry.add("h_trk_charge", "trk_charge;#it{q};Entries", {HistType::kTH1F, {{3, -1.5, 1.5}}});
+      registry.add("h_trk_dcaxy", "trk_dcaxy;#it{DCA}_{xy} (cm);Entries", {HistType::kTH1F, {{100, -0.1, 0.1}}});
+      registry.add("h_trk_dcaxyz", "trk_dcaxyz;#it{DCA}_{xyz} (cm);Entries", {HistType::kTH1F, {{100, -0.1, 0.1}}});
+      registry.add("h_trk_sigmadcaxy", "trk_sigmadcaxy;#it{#sigma}_{#it{DCA}_{xy}} (cm);Entries", {HistType::kTH1F, {{100, 0., 0.1}}});
+      registry.add("h_trk_sigmadcaxyz", "trk_sigmadcaxyz;#it{#sigma}_{#it{DCA}_{xyz}} (cm);Entries", {HistType::kTH1F, {{100, 0., 0.1}}});
+      registry.add("h_trk_itsncls", "trk_itsncls;ITS NCls;Entries", {HistType::kTH1F, {{10, 0., 10.}}});
+      registry.add("h_trk_tpcncls", "trk_tpcncls;TPC NCls (Found);Entries", {HistType::kTH1F, {{200, 0., 200.}}});
+      registry.add("h_trk_tpcncrs", "trk_tpcncrs;TPC NCrossedRows;Entries", {HistType::kTH1F, {{200, 0., 200.}}});
+      registry.add("h_trk_itschi2ncl", "trk_itschi2ncl;ITS #it{#chi}^{2}/ndf;Entries", {HistType::kTH1F, {{100, 0., 20.}}});
+      registry.add("h_trk_tpcchi2ncl", "trk_tpcchi2ncl;TPC #it{#chi}^{2}/ndf;Entries", {HistType::kTH1F, {{100, 0., 10.}}});
+      registry.add("h_jtrack_counter", "jtrack counter", {HistType::kTH1F, {{1, 0., 1.}}});
+      registry.add("h2_trk_jtrackpt_vs_origtrackpt", "JTracks::pt vs Tracks::pt", {HistType::kTH2F, {{100, 0., 100.}, {100, 0., 100.}}});
+      registry.add("h_trk_vtx_index", "trk_vtx_index;Vertex index;Entries", {HistType::kTH1F, {{20, 0., 20.}}});
+      registry.add("h_trk_origin", "trk_origin;Track origin;Entries", {HistType::kTH1F, {{5, 0., 5.}}});
+      auto h_trk_origin = registry.get<TH1>(HIST("h_trk_origin"));
+      h_trk_origin->GetXaxis()->SetBinLabel(1, "NotPhysPrim");
+      h_trk_origin->GetXaxis()->SetBinLabel(2, "Charm");
+      h_trk_origin->GetXaxis()->SetBinLabel(3, "Beauty");
+      h_trk_origin->GetXaxis()->SetBinLabel(4, "Primary");
+      h_trk_origin->GetXaxis()->SetBinLabel(5, "OtherSecondary");
+    }
   }
 
   // FIXME filtering only works when you loop directly over the list, but if you loop over it as a constituent they will not be filtered
@@ -272,6 +339,8 @@ struct BJetTreeCreator {
   using JetTracksMCDwID = soa::Filtered<soa::Join<aod::JetTracksMCD, aod::JTrackExtras, aod::JTrackPIs>>;
   using DataJets = soa::Filtered<soa::Join<aod::ChargedJets, aod::ChargedJetConstituents, aod::DataSecondaryVertex3ProngIndices>>;
 
+  using OriginalTracks = soa::Join<aod::Tracks, aod::TracksCov, aod::TrackSelection, aod::TracksDCA, aod::TracksDCACov, aod::TracksExtra>;
+
   // Function to get the reduction factor based on jet pT
   double getReductionFactor(double jetPT)
   {
@@ -347,6 +416,8 @@ struct BJetTreeCreator {
     }
   }
 
+  using TrackLabelMap = std::unordered_map<std::string, std::vector<int>>;
+
   template <typename AnyCollision, typename AnalysisJet, typename AnyTracks, typename SecondaryVertices>
   void analyzeJetTrackInfo(AnyCollision const& /*collision*/, AnalysisJet const& analysisJet, AnyTracks const& /*allTracks*/, SecondaryVertices const& /*allSVs*/, std::vector<int>& trackIndices, int jetFlavor = 0, double eventweight = 1.0)
   {
@@ -392,6 +463,77 @@ struct BJetTreeCreator {
     }
   }
 
+  template <typename AnyCollision, typename AnalysisJet, typename AnyTracks, typename AnyOriginalTracks>
+  void analyzeJetTrackInfoForGNN(AnyCollision const& /*collision*/, AnalysisJet const& analysisJet, AnyTracks const& /*allTracks*/, AnyOriginalTracks const&, std::vector<int>& trackIndices, int jetFlavor = 0, double eventweight = 1.0, TrackLabelMap *trkLabels = nullptr)
+  {
+    int trkIdx = -1;
+    for (auto& constituent : analysisJet.template tracks_as<AnyTracks>()) {
+
+      trkIdx++;
+
+      if (constituent.pt() < trackPtMin) {
+        continue;
+      }
+
+      double deltaRJetTrack = jetutilities::deltaR(analysisJet, constituent);
+      double dotProduct = RecoDecay::dotProd(std::array<float, 3>{analysisJet.px(), analysisJet.py(), analysisJet.pz()}, std::array<float, 3>{constituent.px(), constituent.py(), constituent.pz()});
+      int sign = jettaggingutilities::getGeoSign(analysisJet, constituent);
+
+      registry.fill(HIST("h2_SIPs2D_jetpT"), analysisJet.pt(), sign * std::abs(constituent.dcaXY()) / constituent.sigmadcaXY(), eventweight);
+      registry.fill(HIST("h2_SIPs3D_jetpT"), analysisJet.pt(), sign * std::abs(constituent.dcaXYZ()) / constituent.sigmadcaXYZ(), eventweight);
+
+      if (doprocessMCJetsForGNN) {
+        if (jetFlavor == 2) {
+          registry.fill(HIST("h2_SIPs2D_jetpT_bjet"), analysisJet.pt(), sign * std::abs(constituent.dcaXY()) / constituent.sigmadcaXY(), eventweight);
+          registry.fill(HIST("h2_SIPs3D_jetpT_bjet"), analysisJet.pt(), sign * std::abs(constituent.dcaXYZ()) / constituent.sigmadcaXYZ(), eventweight);
+        } else if (jetFlavor == 1) {
+          registry.fill(HIST("h2_SIPs2D_jetpT_cjet"), analysisJet.pt(), sign * std::abs(constituent.dcaXY()) / constituent.sigmadcaXY(), eventweight);
+          registry.fill(HIST("h2_SIPs3D_jetpT_cjet"), analysisJet.pt(), sign * std::abs(constituent.dcaXYZ()) / constituent.sigmadcaXYZ(), eventweight);
+        } else {
+          registry.fill(HIST("h2_SIPs2D_jetpT_lfjet"), analysisJet.pt(), sign * std::abs(constituent.dcaXY()) / constituent.sigmadcaXY(), eventweight);
+          registry.fill(HIST("h2_SIPs3D_jetpT_lfjet"), analysisJet.pt(), sign * std::abs(constituent.dcaXYZ()) / constituent.sigmadcaXYZ(), eventweight);
+        }
+
+        auto origConstit = constituent.template track_as<AnyOriginalTracks>();
+
+        //+
+        int trkVtxIndex = 0;
+        int trkOrigin = 0;
+        if (trkLabels != nullptr) {
+          trkVtxIndex = (*trkLabels)["trkVtxIndex"][trkIdx];
+          trkOrigin = (*trkLabels)["trkOrigin"][trkIdx];
+        }
+
+        //+trk
+        registry.fill(HIST("h_trk_pt"), constituent.pt(), eventweight);
+        registry.fill(HIST("h_trk_eta"), constituent.eta(), eventweight);
+        registry.fill(HIST("h_trk_phi"), origConstit.phi(), eventweight);
+        registry.fill(HIST("h_trk_charge"), constituent.sign(), eventweight);
+        registry.fill(HIST("h_trk_dcaxy"), std::abs(constituent.dcaXY()) * sign, eventweight);
+        registry.fill(HIST("h_trk_dcaxyz"), std::abs(constituent.dcaXYZ()) * sign, eventweight);
+        registry.fill(HIST("h_trk_sigmadcaxy"), constituent.sigmadcaXY(), eventweight);
+        registry.fill(HIST("h_trk_sigmadcaxyz"), constituent.sigmadcaXYZ(), eventweight);
+        registry.fill(HIST("h_trk_itsncls"), origConstit.itsNCls(), eventweight);
+        registry.fill(HIST("h_trk_tpcncls"), origConstit.tpcNClsFound(), eventweight);
+        registry.fill(HIST("h_trk_tpcncrs"), origConstit.tpcNClsCrossedRows(), eventweight);
+        registry.fill(HIST("h_trk_itschi2ncl"), origConstit.itsChi2NCl(), eventweight);
+        registry.fill(HIST("h_trk_tpcchi2ncl"), origConstit.tpcChi2NCl(), eventweight);
+        registry.fill(HIST("h2_trk_jtrackpt_vs_origtrackpt"), constituent.pt(), origConstit.pt(), eventweight); // jtrack & new extra table are well joined (linear correlation)
+        registry.fill(HIST("h_trk_vtx_index"), trkVtxIndex);
+        registry.fill(HIST("h_trk_origin"), trkOrigin);
+
+        if (produceTree) {
+          bjetTracksExtraTable(/*bjetParamsTable.lastIndex() + 1, */origConstit.phi(), constituent.sign(), origConstit.itsChi2NCl(), origConstit.tpcChi2NCl(), origConstit.itsNCls(), origConstit.tpcNClsFound(), origConstit.tpcNClsCrossedRows(), trkOrigin, trkVtxIndex); //+
+        }
+      }
+
+      if (produceTree) {
+        bjetTracksParamsTable(bjetParamsTable.lastIndex() + 1, constituent.pt(), constituent.eta(), dotProduct, dotProduct / analysisJet.p(), deltaRJetTrack, std::abs(constituent.dcaXY()) * sign, constituent.sigmadcaXY(), std::abs(constituent.dcaXYZ()) * sign, constituent.sigmadcaXYZ(), constituent.p() / analysisJet.p(), 0.);
+      }
+      trackIndices.push_back(bjetTracksParamsTable.lastIndex());
+    }
+  }
+
   void processDummy(FilteredCollision::iterator const& /*collision*/)
   {
   }
@@ -548,6 +690,85 @@ struct BJetTreeCreator {
   }
   PROCESS_SWITCH(BJetTreeCreator, processMCJets, "jet information in MC", false);
 
+  using MCDJetTableNoSV = soa::Filtered<soa::Join<aod::ChargedMCDetectorLevelJets, aod::ChargedMCDetectorLevelJetConstituents, aod::ChargedMCDetectorLevelJetsMatchedToChargedMCParticleLevelJets, aod::ChargedMCDetectorLevelJetEventWeights>>;
+  using JetParticleswID = soa::Join<aod::JetParticles, aod::JMcParticlePIs>;
+
+  void processMCJetsForGNN(FilteredCollisionMCD::iterator const &collision, aod::JMcCollisions const &, MCDJetTableNoSV const &MCDjets, MCPJetTable const &MCPjets, JetTracksMCDwID const &allTracks, JetParticleswID const &MCParticles, OriginalTracks const& origTracks, aod::McParticles const& origParticles)
+  {
+    if (!jetderiveddatautilities::selectCollision(collision, eventSelection) || (static_cast<double>(std::rand()) / RAND_MAX < eventReductionFactor)) {
+      return;
+    }
+
+    registry.fill(HIST("h_vertexZ"), collision.posZ());
+
+    auto const mcParticlesPerColl = MCParticles.sliceBy(McParticlesPerCollision, collision.mcCollisionId());
+    auto const mcPJetsPerColl = MCPjets.sliceBy(McPJetsPerCollision, collision.mcCollisionId());
+
+    for (const auto& analysisJet : MCDjets) {
+
+      bool jetIncluded = false;
+      for (auto jetR : jetRadiiValues) {
+        if (analysisJet.r() == static_cast<int>(jetR * 100)) {
+          jetIncluded = true;
+          break;
+        }
+      }
+
+      if (!jetIncluded) {
+        continue;
+      }
+
+      std::vector<int> tracksIndices;
+      std::vector<int> SVsIndices;
+
+      int16_t jetFlavor = 0;
+
+      for (auto& mcpjet : analysisJet.template matchedJetGeo_as<MCPJetTable>()) {
+        if (useQuarkDef) {
+          jetFlavor = jettaggingutilities::getJetFlavor(mcpjet, mcParticlesPerColl);
+        } else {
+          jetFlavor = jettaggingutilities::getJetFlavorHadron(mcpjet, mcParticlesPerColl);
+        }
+      }
+
+      float eventWeight = analysisJet.eventWeight();
+
+      //+
+      TrackLabelMap trkLabels{{"trkVtxIndex", {}}, {"trkOrigin", {}}};
+      int nVertices = jettaggingutilities::vertexClustering(collision.template mcCollision_as<aod::JMcCollisions>(), analysisJet, allTracks, MCParticles, origParticles, trkLabels, vtxRes, trackPtMin);
+      analyzeJetTrackInfoForGNN(collision, analysisJet, allTracks, origTracks, tracksIndices, jetFlavor, eventWeight, &trkLabels);
+
+      registry.fill(HIST("h2_jetMass_jetpT"), analysisJet.pt(), analysisJet.mass(), eventWeight);
+      registry.fill(HIST("h2_nTracks_jetpT"), analysisJet.pt(), tracksIndices.size());
+
+      //+jet
+      registry.fill(HIST("h_jet_pt"), analysisJet.pt());
+      registry.fill(HIST("h_jet_eta"), analysisJet.eta());
+      registry.fill(HIST("h_jet_phi"), analysisJet.phi());
+      registry.fill(HIST("h_jet_flav"), jetFlavor);
+      registry.fill(HIST("h_n_trks"), tracksIndices.size());
+      registry.fill(HIST("h_jet_mass"), analysisJet.mass());
+      registry.fill(HIST("h_n_vertices"), nVertices);
+
+      if (jetFlavor == 2) {
+        registry.fill(HIST("h2_jetMass_jetpT_bjet"), analysisJet.pt(), analysisJet.mass(), eventWeight);
+        registry.fill(HIST("h_jetpT_detector_bjet"), analysisJet.pt(), eventWeight);
+      } else if (jetFlavor == 1) {
+        registry.fill(HIST("h2_jetMass_jetpT_cjet"), analysisJet.pt(), analysisJet.mass(), eventWeight);
+        registry.fill(HIST("h_jetpT_detector_cjet"), analysisJet.pt(), eventWeight);
+      } else {
+        registry.fill(HIST("h2_jetMass_jetpT_lfjet"), analysisJet.pt(), analysisJet.mass(), eventWeight);
+        registry.fill(HIST("h_jetpT_detector_lfjet"), analysisJet.pt(), eventWeight);
+      }
+
+      if (produceTree) {
+        bjetConstituentsTable(bjetParamsTable.lastIndex() + 1, tracksIndices, SVsIndices);
+        bjetParamsTable(analysisJet.pt(), analysisJet.eta(), analysisJet.phi(), tracksIndices.size(), nVertices, analysisJet.mass(), jetFlavor, analysisJet.r());
+      }
+    }
+  }
+  PROCESS_SWITCH(BJetTreeCreator, processMCJetsForGNN, "jet information in MC for GNN", false);
+
   Filter mccollisionFilter = nabs(aod::jmccollision::posZ) < vertexZCut;
   using FilteredCollisionMCP = soa::Filtered<aod::JMcCollisions>;