// SPDX-License-Identifier: LGPL-3.0-or-later
// Copyright (C) 2026 ePIC Collaboration

#include <edm4hep/MCParticleCollection.h>
#include <edm4hep/Vector3f.h>
#include <podio/ObjectID.h>
#include <podio/detail/Link.h>
#include <podio/detail/LinkCollectionImpl.h>
#include <cmath>
#include <memory>
#include <tuple>

#include "ClustersToParticles.h"

namespace eicrecon {

void ClustersToParticles::init() {
  const auto& particle = m_particleSvc.particle(m_cfg.pdgCode);
  m_mass               = particle.mass;
  m_charge             = particle.charge;
}

void ClustersToParticles::process(const ClustersToParticles::Input& input,
                                  const ClustersToParticles::Output& output) const {
  const auto [clusters, cluster_assocs] = input;
  auto [parts, part_links, part_assocs] = output;

  for (const auto& cluster : *clusters) {
    const auto energy = cluster.getEnergy();
    const auto pos    = cluster.getPosition();

    const auto momentum_mag =
        (energy > m_mass) ? std::sqrt(energy * energy - m_mass * m_mass) : 0.0;
    const auto pos_mag = std::sqrt(pos.x * pos.x + pos.y * pos.y + pos.z * pos.z);

    edm4hep::Vector3f momentum{0, 0, 0};
    if (pos_mag > 0) {
      momentum.x = momentum_mag * pos.x / pos_mag;
      momentum.y = momentum_mag * pos.y / pos_mag;
      momentum.z = momentum_mag * pos.z / pos_mag;
    }

    debug("Converting cluster: index={:<4} energy={:<8.3f} position=({:<8.3f}, {:<8.3f}, {:<8.3f})",
          cluster.getObjectID().index, energy, pos.x, pos.y, pos.z);

    auto rec_part = parts->create();
    rec_part.addToClusters(cluster);
    rec_part.setPDG(m_cfg.pdgCode);
    rec_part.setType(0);
    rec_part.setEnergy(energy);
    rec_part.setMomentum(momentum);
    rec_part.setCharge(static_cast<float>(m_charge));
    rec_part.setMass(m_mass);
    rec_part.setGoodnessOfPID(0); // assume no PID until proven otherwise
    rec_part.setReferencePoint({0.0f, 0.0f, 0.0f});

    // Handle associations if provided
    for (auto cluster_assoc : *cluster_assocs) {
      if (cluster_assoc.getRec() == cluster) {
        trace("Found cluster association: index={} -> index={}, weight={}",
              cluster_assoc.getRec().getObjectID().index,
              cluster_assoc.getSim().getObjectID().index, cluster_assoc.getWeight());
        auto part_link = part_links->create();
        part_link.setFrom(rec_part);
        part_link.setTo(cluster_assoc.getSim());
        part_link.setWeight(cluster_assoc.getWeight());
        auto part_assoc = part_assocs->create();
        part_assoc.setRec(rec_part);
        part_assoc.setSim(cluster_assoc.getSim());
        part_assoc.setWeight(cluster_assoc.getWeight());
      }
    }
  }
}

} // namespace eicrecon
