#include "BDSXtrackInterface.hh"
#include <cstring>
#include <cmath>
#include <BDSSamplerCustom.hh>


BDSParticleDefinition* PrepareBDSParticleDefition(long long int pdgIDIn, double momentumIn, 
                                                  double kineticEnergyIn, double ionChargeIn)
{
    G4int pdgID = (G4int) pdgIDIn;

    G4ParticleTable* particleTable = G4ParticleTable::GetParticleTable();
    G4ParticleDefinition* particleDefGeant;

    // Wrap in our class that calculates momentum and kinetic energy.
    // Requires that one of E, Ek, P be non-zero (only one).
    BDSParticleDefinition* particleDefinition = nullptr;
    BDSIonDefinition* ionDef = nullptr;

    // PDG for ions = 10LZZZAAAI
    if (pdgID > 1000000000) // is an ion
    {   

        G4IonTable* ionTable = particleTable->GetIonTable();
        particleDefGeant = ionTable->GetIon(pdgID);
        if (!particleDefGeant)
        {throw BDSException("BDSXtrackInterface> Ion \"" + std::to_string(pdgID) + "\" not found");}

        G4int ionCharge = ionChargeIn==0 ? particleDefGeant->GetAtomicNumber() : (G4int) ionChargeIn; 
        ionDef = new BDSIonDefinition(particleDefGeant->GetAtomicMass(),
                                      particleDefGeant->GetAtomicNumber(),
                                      ionCharge);

        // correct the particle mass in the case of partially-stripped ions
        G4double mass   = ionTable->GetIonMass(ionDef->Z(), ionDef->A());
        G4double charge = ionDef->Charge(); // correct even if overridden
        mass += ionDef->NElectrons()*G4Electron::Definition()->GetPDGMass();
        // The constructor with a custom name requires the bdsim name of the ion
        std::string bdsimPartName = "ion " + std::to_string(static_cast<int>(ionDef->A())) 
                                  + " " + std::to_string(static_cast<int>(ionDef->Z())) 
                                  + " " + std::to_string(static_cast<int>(charge));

        particleDefinition = new BDSParticleDefinition(bdsimPartName, mass, charge, 0, 
                                                       kineticEnergyIn, momentumIn, 1, ionDef, pdgID);
    }
    else
    {
        particleDefGeant = particleTable->FindParticle(pdgID);
        if (!particleDefGeant)
        {throw BDSException("BDSXtrackInterface> Particle \"" + std::to_string(pdgID) + "\" not found");}
        particleDefinition = new BDSParticleDefinition(particleDefGeant, 0, kineticEnergyIn, momentumIn, 1, nullptr);
    }

    return particleDefinition;
}


XtrackInterface::XtrackInterface(const  std::string& bdsimConfigFile,
                                 long long int       referencePdgIdIn,
                                 double              referenceEkIn,
                                 double              relativeEnergyCutIn,
                                 int                 seedIn,
                                 int                 referenceIonChargeIn,
                                 bool                batchMode):
        relativeEnergyCut(relativeEnergyCutIn),
        seed(seedIn)
{
    stp = new BDSBunchSixTrackLink();
    bds = new BDSIMLink(stp);

    std::string seedStr = std::to_string(seed);
    std::vector<std::string> arguments = {"--file=" + bdsimConfigFile,
                                          "--file=" + bdsimConfigFile,
                                          //"--vis_debug",
                                          "--output=none",
                                          "--seed=" + seedStr,
                                          "--outfile=output_" + seedStr};

    for(auto & argument : arguments)
    {
        argv.push_back(strdup(argument.c_str()));
    }
    if (batchMode)
    {
        std::string batch_flag = "--batch";
        argv.push_back(strdup(batch_flag.c_str()));
    }
    argv.push_back(nullptr);

    double referenceEk = referenceEkIn * CLHEP::GeV;

    double relEKCut = relativeEnergyCut;
    if (relEKCut < 1e-6) // defaults to 0 which means 0eV cut which is bad
    { relEKCut = 1.0; }

    double minimumEK = relEKCut * (referenceEk);

    G4cout << "Minimum kinetic energy " << minimumEK << " MeV" << G4endl;
    auto data = argv.data();
    try
    { bds->Initialise(argv.size() - 1, &argv[0], true, minimumEK / CLHEP::GeV, false); } // minimumEk in GeV
    catch (const std::exception &e)
    {
        std::cout << e.what() << std::endl;
        exit(1);
    }

    G4double ionCharge = (G4double) referenceIonChargeIn;
    refParticleDefinition = PrepareBDSParticleDefition(referencePdgIdIn, 0, referenceEk, ionCharge);
}


XtrackInterface::~XtrackInterface()
{
    delete bds;
    delete stp;

    delete refParticleDefinition;
}


void XtrackInterface::addCollimator(const std::string&   name,
                                    const std::string&   material,
                                    double lengthIn,
                                    double apertureLeftIn,
                                    double apertureRightIn,
                                    double rotationIn,
                                    double xOffsetIn,
                                    double yOffsetIn,
                                    double jawTiltLeft,
                                    double jawTiltRight,
                                    int    side)
    {

        bool buildLeft  = side == 0 || side == 1;
        bool buildRight = side == 0 || side == 2;

        bool isACrystal = false;

        bds->AddLinkCollimatorJaw(name,
                                  material,
                                  lengthIn * CLHEP::m,
                                  apertureLeftIn * CLHEP::m,
                                  apertureRightIn * CLHEP::m,
                                  rotationIn * CLHEP::rad,
                                  xOffsetIn * CLHEP::m,
                                  yOffsetIn * CLHEP::m,
                                  jawTiltLeft * CLHEP::rad,
                                  jawTiltRight * CLHEP::rad,
                                  buildLeft,
                                  buildRight,
                                  isACrystal,
                                  0);
    }


void XtrackInterface::addParticle(double xIn,
                                  double yIn,
                                  double pxIn,
                                  double pyIn,
                                  double ctIn,
                                  double deltapIn,
                                  double chiIn,
                                  double chargeRatioIn,
                                  double sIn,
                                  int64_t pdgIDIn,
                                  int64_t trackidIn
                                  )
{

    auto x  = (G4double) xIn;
    auto y  = (G4double) yIn;
    auto px  = (G4double) pxIn;
    auto py  = (G4double) pyIn;
    auto ct  = (G4double) ctIn;
    auto deltap = (G4double) deltapIn;
    auto chi = (G4double) chiIn;
    auto charge_ratio = (G4double) chargeRatioIn;
    auto s = (G4double) sIn;
    auto trackid = (G4int) trackidIn;
    auto pdgid = (long long int) pdgIDIn;

    //G4double mass_ratio = charge_ratio / chi;
    G4double q = charge_ratio * refParticleDefinition->Charge();
    G4double mass_ratio = charge_ratio / chi;
    G4double p = refParticleDefinition->Momentum() * (deltap + 1) * mass_ratio;

    // If the pdg id for the particle is 0 (default), take the reference pdg id
    long long int pdgidPart = (pdgid==0) ? refParticleDefinition->PDGID() : pdgid; 
    auto partDef = PrepareBDSParticleDefition(pdgidPart, p, 0, q);

    G4double t = - ct * CLHEP::m / (refParticleDefinition->Beta() * CLHEP::c_light); // this is time difference in ns

    G4double oneplusdelta = (1 + deltap) * mass_ratio;
    // G4double pz = std::sqrt(oneplusdelta*oneplusdelta - px*px - py*py);
    G4double xp = px / oneplusdelta;
    G4double yp = py / oneplusdelta;

    // Zp0 is 1 as here we assume no back-scatterd particles, e.g p>0
    G4double zp = BDSBunch::CalculateZp(xp, yp, 1);

    BDSParticleCoordsFull coords = BDSParticleCoordsFull(x * CLHEP::m,
                                                         y * CLHEP::m,
                                                         0,
                                                         xp,
                                                         yp,
                                                         zp,
                                                         t,
                                                         0,
                                                         partDef->TotalEnergy(),
                                                         1);

    stp->AddParticle(partDef, coords, trackid, trackid);
}


void XtrackInterface::addParticles(const py::list& coordinates)
{   
    //TODO get the charge and mass ratios
    // Obtain the arrays from the list and cast them to the correct array type
    py::array_t<double> x = py::cast<py::array>(coordinates[0]);
    py::array_t<double> y = py::cast<py::array>(coordinates[1]);
    py::array_t<double> px = py::cast<py::array>(coordinates[2]);
    py::array_t<double> py = py::cast<py::array>(coordinates[3]);
    py::array_t<double> zeta = py::cast<py::array>(coordinates[4]);
    py::array_t<double> delta = py::cast<py::array>(coordinates[5]);
    py::array_t<double> chi = py::cast<py::array>(coordinates[6]);
    py::array_t<double> charge_ratio = py::cast<py::array>(coordinates[7]);
    py::array_t<double> s = py::cast<py::array>(coordinates[8]);
    py::array_t<int64_t> pdgid = py::cast<py::array>(coordinates[9]);
    py::array_t<int64_t> trackid = py::cast<py::array>(coordinates[10]);
    py::array_t<int64_t> state = py::cast<py::array>(coordinates[11]);
    py::array_t<int64_t> at_element = py::cast<py::array>(coordinates[12]);
    py::array_t<int64_t> at_turn = py::cast<py::array>(coordinates[13]);


    // Obtain the buffers
    py::buffer_info x_buff = x.request();
    py::buffer_info y_buff = y.request();
    py::buffer_info px_buff = px.request();
    py::buffer_info py_buff = py.request();
    py::buffer_info zeta_buff = zeta.request();
    py::buffer_info delta_buff = delta.request();
    py::buffer_info chi_buff = chi.request();
    py::buffer_info charge_ratio_buff = charge_ratio.request();
    py::buffer_info s_buff = s.request();
    py::buffer_info pdgid_buff = pdgid.request();
    py::buffer_info id_buff = trackid.request();
    py::buffer_info state_buff = state.request();
    py::buffer_info at_element_buff = at_element.request();
    py::buffer_info at_turn_buff = at_turn.request();

    // Get the pointers for iteration
    auto x_ptr = static_cast<double *>(x_buff.ptr);
    auto y_ptr = static_cast<double *>(y_buff.ptr);
    auto px_ptr = static_cast<double *>(px_buff.ptr);
    auto py_ptr = static_cast<double *>(py_buff.ptr);
    auto zeta_ptr = static_cast<double *>(zeta_buff.ptr);
    auto delta_ptr = static_cast<double *>(delta_buff.ptr);
    auto chi_ptr = static_cast<double *>(chi_buff.ptr);
    auto charge_ratio_ptr = static_cast<double *>(charge_ratio_buff.ptr);
    auto s_ptr = static_cast<double *>(s_buff.ptr);
    auto pdgid_ptr = static_cast<int64_t *>(pdgid_buff.ptr);
    auto trackid_ptr = static_cast<int64_t *>(id_buff.ptr);
    auto state_ptr = static_cast<int64_t *>(state_buff.ptr);
    auto at_element_ptr = static_cast<int64_t *>(at_element_buff.ptr);
    auto at_turn_ptr = static_cast<int64_t *>(at_turn_buff.ptr);

    long n = 1; // Get the number of elements in the array, assume all arrays have the same number of elements
    for (auto r: x_buff.shape) {
        n *= r;
    }

    for(int i=0; i < n; i++)
    {
        auto x_part = x_ptr[i];
        auto y_part = y_ptr[i];
        auto px_part = px_ptr[i];
        auto py_part = py_ptr[i];
        auto zeta_part = zeta_ptr[i];
        auto delta_part = delta_ptr[i];
        auto chi_part = chi_ptr[i];
        auto charge_ratio_part = charge_ratio_ptr[i];
        auto s_part = s_ptr[i];
        auto pdgid_part = pdgid_ptr[i];
        auto trackid_part = trackid_ptr[i];
        auto state_part = state_ptr[i];
        auto at_element_part = at_element_ptr[i];
        auto at_turn_part = at_turn_ptr[i];

        // The particles with state=0 are inactive and should not be processed at all, but need to keep track of them
        // The internal processing in BDSIM does not feature an active state check, so this must be done at a higher
        // level here

        if (state_part == 1) // State == 1 means that the particle is active
        {
            particleActiveState.push_back(true);
            addParticle(x_part, y_part, px_part, py_part,
                        zeta_part, delta_part, chi_part,
                        charge_ratio_part, s_part, pdgid_part, trackid_part);

            maxParticleID = std::max(maxParticleID, trackid_part);
            bds->SetCurrentMaximumExternalParticleID(maxParticleID);
        }
        else
        {
            if (state_part == -999999999) // This is reserve memory for new particles, no need to process
            {
                break;
            }
            particleActiveState.push_back(false);
        }

        auto particle_coords = new XtrackCoordinates{x_part, y_part, px_part, py_part, zeta_part,
                                                     delta_part, chi_part, charge_ratio_part, s_part,
                                                     pdgid_part, trackid_part, state_part,
                                                     at_element_part, at_turn_part};
        xtrackParticles.push_back(particle_coords);
    }
}


void XtrackInterface::collimate()
{
    bds->BeamOn((G4int)stp->Size());
}


void XtrackInterface::selectCollimator(const std::string& collimatorName)
{
    currentCollimatorName = collimatorName;
    // This doesn't throw an error if the element doesn't exist
    bds->SelectLinkElement(collimatorName);

    // Check if the element exists by querying the index: -1 means it doesn't exist
    if (bds->GetLinkIndex(collimatorName) == -1)
        {throw std::runtime_error("Element not found " + collimatorName);}
}


void XtrackInterface::clearData()
{
    bds->ClearSamplerHits();
    // A malloc error about freeing a pointer not allocated is thrown if the run is terminated after
    // the bunch is manually cleared. Consider using an alternative to vector.clear() in the BDSIM bunch class:
    // vector<T>().swap(x);   // clear x reallocating  (https://www.cplusplus.com/reference/vector/vector/clear/)
    // to ensure the pointer is still allocated after clearing
    stp->ClearParticles();

    for (auto part : xtrackParticles)
    {
       delete part;
    }

    std::vector<XtrackCoordinates*>().swap(xtrackParticles);
    std::vector<bool>().swap(particleActiveState);
    currentCollimatorName.clear();

    maxParticleID = 0;

}


py::dict XtrackInterface::collimateReturn(const py::list& coordinates)
{
    // Prepare the buffers for modifying the primary particle coordinates in place

    // Obtain the arrays from the list and cast them to the correct array type
    py::array_t<double> x = py::cast<py::array>(coordinates[0]);
    py::array_t<double> y = py::cast<py::array>(coordinates[1]);
    py::array_t<double> px = py::cast<py::array>(coordinates[2]);
    py::array_t<double> py = py::cast<py::array>(coordinates[3]);
    py::array_t<double> zeta = py::cast<py::array>(coordinates[4]);
    py::array_t<double> delta = py::cast<py::array>(coordinates[5]);
    py::array_t<double> chi = py::cast<py::array>(coordinates[6]);
    py::array_t<double> charge_ratio = py::cast<py::array>(coordinates[7]);
    py::array_t<double> s = py::cast<py::array>(coordinates[8]);
    py::array_t<int64_t> pdgid = py::cast<py::array>(coordinates[9]);
    py::array_t<int64_t> trackid = py::cast<py::array>(coordinates[10]);
    py::array_t<int64_t> state = py::cast<py::array>(coordinates[11]);
    py::array_t<int64_t> at_element = py::cast<py::array>(coordinates[12]);
    py::array_t<int64_t> at_turn = py::cast<py::array>(coordinates[13]);

    // Obtain the buffers
    py::buffer_info x_buff = x.request();
    py::buffer_info y_buff = y.request();
    py::buffer_info px_buff = px.request();
    py::buffer_info py_buff = py.request();
    py::buffer_info zeta_buff = zeta.request();
    py::buffer_info delta_buff = delta.request();
    py::buffer_info chi_buff = chi.request();
    py::buffer_info charge_ratio_buff = charge_ratio.request();
    py::buffer_info s_buff = s.request();
    py::buffer_info pdgid_buff = pdgid.request();
    py::buffer_info id_buff = trackid.request();
    py::buffer_info state_buff = state.request();
    py::buffer_info at_element_buff = at_element.request();
    py::buffer_info at_turn_buff = at_turn.request();

    // Get the pointers for iteration
    auto x_ptr = static_cast<double *>(x_buff.ptr);
    auto y_ptr = static_cast<double *>(y_buff.ptr);
    auto px_ptr = static_cast<double *>(px_buff.ptr);
    auto py_ptr = static_cast<double *>(py_buff.ptr);
    auto zeta_ptr = static_cast<double *>(zeta_buff.ptr);
    auto delta_ptr = static_cast<double *>(delta_buff.ptr);
    auto chi_ptr = static_cast<double *>(chi_buff.ptr);
    auto charge_ratio_ptr = static_cast<double *>(charge_ratio_buff.ptr);
    auto s_ptr = static_cast<double *>(s_buff.ptr);
    auto pdgid_ptr = static_cast<int64_t *>(pdgid_buff.ptr);
    auto trackid_ptr = static_cast<int64_t *>(id_buff.ptr);
    auto state_ptr = static_cast<int64_t *>(state_buff.ptr);
    auto at_element_ptr = static_cast<int64_t *>(at_element_buff.ptr);
    auto at_turn_ptr = static_cast<int64_t *>(at_turn_buff.ptr);

    // N.B book keeping needed
    // Note: Load and and transfer for surviving primaries
    // Secondaries inherit the element and turn of parent particle
    //auto at_element = static_cast<int64_t *>(at_element.ptr);
    //auto turn = static_cast<int64_t *>(turn.ptr);


    // Prepare the buffers for writing out the products
    // New numpy arrays are allocated for writing the products

    // Access the sampler hits - particles reaching the planes for transport back
    const BDSHitsCollectionSamplerLink* hits = bds->SamplerHits();

    //size_t hitsCount = hits ? hits->GetSize() : 0;

    size_t hitsCount = 0;
    if (hits)
    {
        hitsCount = hits->GetSize();
    }
    else
    {
        // There were no hits - check if there were any active particles at all coming in
        if (!stp->Size())
        {
            // A particle needs to be added to the bunch and the GetNextParticleLocal method
            // must be called to initialise the variables and ensure a safe deletion of the
            // underlying BDSBunch object. This will not be needed in the next release of BDSIM
            addParticle(0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0);
            // The dummy particle added is never used, as the collimateReturn method terminates
            // the processing
            stp->GetNextParticleLocal();
        }
    }


    // Count the number of secondary particles
    size_t secondaryCount = 0;
    for (size_t i = 0; i < hitsCount; i++)
    {
        auto hit = (*hits)[i];
        if (hit->externalParticleID != hit->externalParentID) { secondaryCount++; }
    }

    // The output arrays has slots for all primary particles, regardless if lost or not, and for secondary particles
    size_t output_size = secondaryCount;

    // TODO: this is mostly default buffers, so there are probably simpler constructors to use
    // Prepare the numpy array that will be returned
    auto x_out = py::array(py::buffer_info(
            nullptr,            /* Pointer to data (nullptr -> ask NumPy to allocate!) */
            sizeof(double),     /* Size of one item */
            py::format_descriptor<double>::value, /* Buffer format */
            1,          /* How many dimensions? */
            { output_size },  /* Number of elements for each dimension */
            { sizeof(double) }  /* Strides for each dimension */
    ));

    auto y_out = py::array(py::buffer_info(
            nullptr,
            sizeof(double),
            py::format_descriptor<double>::value,
            1,
            { output_size },
            { sizeof(double) }
    ));

    auto px_out = py::array(py::buffer_info(
            nullptr,
            sizeof(double),
            py::format_descriptor<double>::value,
            1,
            { output_size },
            { sizeof(double) }
    ));

    auto py_out = py::array(py::buffer_info(
            nullptr,
            sizeof(double),
            py::format_descriptor<double>::value,
            1,
            { output_size },
            { sizeof(double) }
    ));

    auto zeta_out = py::array(py::buffer_info(
            nullptr,
            sizeof(double),
            py::format_descriptor<double>::value,
            1,
            { output_size },
            { sizeof(double) }
    ));

    auto delta_out = py::array(py::buffer_info(
            nullptr,
            sizeof(double),
            py::format_descriptor<double>::value,
            1,
            { output_size },
            { sizeof(double) }
    ));

    auto s_out = py::array(py::buffer_info(
            nullptr,
            sizeof(double),
            py::format_descriptor<double>::value,
            1,
            { output_size },
            { sizeof(double) }
    ));

    auto pdgid_out = py::array(py::buffer_info(
            nullptr,
            sizeof(int64_t),
            py::format_descriptor<int64_t>::value,
            1,
            { output_size },
            { sizeof(int64_t) }
    ));

    auto trackid_out = py::array(py::buffer_info(
            nullptr,
            sizeof(int64_t),
            py::format_descriptor<int64_t>::value,
            1,
            { output_size },
            { sizeof(int64_t) }
    ));

    auto state_out = py::array(py::buffer_info(
            nullptr,
            sizeof(int64_t),
            py::format_descriptor<int64_t>::value,
            1,
            { output_size },
            { sizeof(int64_t) }
    ));

    auto at_element_out = py::array(py::buffer_info(
            nullptr,
            sizeof(int64_t),
            py::format_descriptor<int64_t>::value,
            1,
            { output_size },
            { sizeof(int64_t) }
    ));

    auto at_turn_out = py::array(py::buffer_info(
            nullptr,
            sizeof(int64_t),
            py::format_descriptor<int64_t>::value,
            1,
            { output_size },
            { sizeof(int64_t) }
    ));

    auto massratio_out = py::array(py::buffer_info(
            nullptr,
            sizeof(double),
            py::format_descriptor<double>::value,
            1,
            { output_size },
            { sizeof(double) }
    ));

    auto chargeratio_out = py::array(py::buffer_info(
            nullptr,
            sizeof(double),
            py::format_descriptor<double>::value,
            1,
            { output_size },
            { sizeof(double) }
    ));

    auto x_prod_buf = x_out.request();
    auto y_prod_buf = y_out.request();
    auto px_prod_buf = px_out.request();
    auto py_prod_buf = py_out.request();
    auto zeta_prod_buf = zeta_out.request();
    auto delta_prod_buf = delta_out.request();
    auto s_prod_buf = s_out.request();
    auto pdgid_prod_buf = pdgid_out.request();
    auto trackid_prod_buf = trackid_out.request();
    auto state_prod_buf = state_out.request();
    auto at_element_prod_buf = at_element_out.request();
    auto at_turn_prod_buf = at_turn_out.request();
    auto massratio_prod_buf = massratio_out.request();
    auto chargeratio_prod_buf = chargeratio_out.request();

    auto *x_prod_ptr = (double *) x_prod_buf.ptr;
    auto *y_prod_ptr = (double *) y_prod_buf.ptr;
    auto *px_prod_ptr = (double *) px_prod_buf.ptr;
    auto *py_prod_ptr = (double *) py_prod_buf.ptr;
    auto *zeta_prod_ptr = (double *) zeta_prod_buf.ptr;
    auto *delta_prod_ptr = (double *) delta_prod_buf.ptr;
    auto *s_prod_ptr = (double *) s_prod_buf.ptr;
    auto *pdgid_prod_ptr = (int64_t *) pdgid_prod_buf.ptr;
    auto *trackid_prod_ptr = (int64_t *) trackid_prod_buf.ptr;
    auto *state_prod_ptr = (int64_t *) state_prod_buf.ptr;
    auto *at_element_prod_ptr = (int64_t *) at_element_prod_buf.ptr;
    auto *at_turn_prod_ptr = (int64_t *) at_turn_prod_buf.ptr;
    auto *massratio_prod_ptr = (double *) massratio_prod_buf.ptr;
    auto *chargeratio_prod_ptr = (double *) chargeratio_prod_buf.ptr;

    // Loop through the particles in the *original* bunch - the primaries
    size_t hits_index = 0;
    bool prim_survied = false;
    // double sum_deltaplusone_sec = 0.0;
    double sum_secondary_energy = 0.0;

    size_t prod_write_index = 0;

    for (size_t i=0; i < particleActiveState.size(); i++){
        if (!particleActiveState.at(i)){
            continue; // This was an inactive particle that hasn't been processed, do not change it
        }

        auto xtrack_part = xtrackParticles.at(i); // Get the cached coordinates of the original xtrack particle

        auto part = stp->GetNextParticle(); // Advance through the bunch
        auto prim_part_id = stp->CurrentExternalParticleID(); // Get the ID of the primary particle

        // Now start looping over the hits - the particles to be returned to the tracker
        // These can be primary or secondary particles. Each primary can produce 0, 1, or 2+ products
        // The products need to be sorted to keep the array order - surviving primary particles are all
        // filled in first. If a primary didn't survive, keep the original coordinates and make it inactive.
        // The hits are ordered by primary event, so just need one loop.
        while (hits_index < hitsCount)
        {
            BDSHitSamplerLink* hit = (*hits)[hits_index];

            if (hit->externalParentID != prim_part_id) { // The hits corresponding to the current primary are exhausted
                break;
            }

            const BDSParticleCoordsFull &coords = hit->coords;

            double qratio = hit->charge / refParticleDefinition->Charge();
            double mratio = hit->mass / refParticleDefinition->Mass();

            double dp = (hit->momentum / mratio - refParticleDefinition->Momentum()) / refParticleDefinition->Momentum();

            double collLength = bds->GetArcLengthOfLinkElement(currentCollimatorName);
            /// Need to compensate for the geometry construction in BDSIM
            /// There is a safety margin that is added to the collimator legnth
            double collMargin = 2.5 * BDSSamplerCustom::ChordLength();
            double zt = refParticleDefinition->Beta() * CLHEP::c_light * ((collLength + collMargin) / (CLHEP::c_light * refParticleDefinition->Beta()) - coords.T);

            double oneplusdelta = (1 + dp) * mratio;

            auto track_id = hit->externalParticleID;
            auto parent_id = hit->externalParentID;
            auto pdg_id = hit->pdgID;

            if (track_id == parent_id){
                // This is a primary particle as its parent is itself
                prim_survied = true;

                x_ptr[i] = coords.x / CLHEP::m;
                y_ptr[i] = coords.y / CLHEP::m;
                px_ptr[i] = coords.xp * oneplusdelta; // convert back to px proper
                py_ptr[i] = coords.yp * oneplusdelta;
                zeta_ptr[i] = zt / CLHEP::m;
                delta_ptr[i] = dp;
                s_ptr[i] = xtrack_part->s + collLength / CLHEP::m;
                //trackid_ptr[i] = track_id; // Don't touch the primary particle id
                state_ptr[i] = 1; // active
            }
            else
            {
                // Secondary particles are populated in newly allocated arrays
                x_prod_ptr[prod_write_index] = coords.x / CLHEP::m;
                y_prod_ptr[prod_write_index] = coords.y / CLHEP::m;
                px_prod_ptr[prod_write_index] = coords.xp * oneplusdelta; // convert back to px proper
                py_prod_ptr[prod_write_index] = coords.yp * oneplusdelta; // convert back to py proper;
                zeta_prod_ptr[prod_write_index] = zt / CLHEP::m;
                delta_prod_ptr[prod_write_index] = dp;
                s_prod_ptr[prod_write_index] = xtrack_part->s + collLength / CLHEP::m;
                pdgid_prod_ptr[prod_write_index] = pdg_id;
                trackid_prod_ptr[prod_write_index] = parent_id;
                state_prod_ptr[prod_write_index] = 1; // active
                at_element_prod_ptr[prod_write_index] = xtrack_part->at_element; // active
                at_turn_prod_ptr[prod_write_index] = xtrack_part->at_turn; // active
                massratio_prod_ptr[prod_write_index] = mratio;
                chargeratio_prod_ptr[prod_write_index] = qratio;

                sum_secondary_energy += std::sqrt(std::pow(hit->momentum,2) + std::pow(hit->mass,2));
                prod_write_index++;
            }

            hits_index++;
        }

        if (!prim_survied) // Primary didn't survive - set inactive
        {
            state_ptr[i] = -333; // inactive

            // Correct the energy of the lost primary particle to account for the production of secondaries
            // The effective delta is such that the lost particle has the effective delta
            // which corresponds to the energy in - energy out for this primary

            // reconstruct the incoming primary particle energy
            G4double qprim = charge_ratio_ptr[i] * refParticleDefinition->Charge();
            G4double mass_ratio_prim = charge_ratio_ptr[i] / chi_ptr[i];
            G4double p_prim = refParticleDefinition->Momentum() * (delta_ptr[i] + 1) * mass_ratio_prim;
            G4double mass_prim = mass_ratio_prim * refParticleDefinition->Mass();
            G4double energy_prim = std::sqrt(std::pow(p_prim, 2) + std::pow(mass_prim, 2));

            // compute the effective delta
            G4double energy_diff = energy_prim - sum_secondary_energy;
            G4double p_eff;
            G4double squared_energy_mass_diff = std::pow(energy_diff, 2) - std::pow(mass_prim, 2);
            if (squared_energy_mass_diff < 0){
                // This means that the total energy escaping includes part of the rest mass of 
                // the primary. Tolerate the error for now, as otherwise need to adjust also the 
                // mass and PDG id of the lost primary particle
                p_eff = 0;
            }
            else
            {
                p_eff = std::sqrt(squared_energy_mass_diff);
            }
            G4double delta_eff = (p_eff / mass_ratio_prim - refParticleDefinition->Momentum()) / refParticleDefinition->Momentum();
            delta_ptr[i] = delta_eff;
        }
        prim_survied = false; // reset for next particle
        sum_secondary_energy = 0.0;
    }

    auto result = py::dict();

    result["s"] = s_out;
    result["x"] = x_out;
    result["px"] = px_out;
    result["y"] = y_out;
    result["py"] = py_out;
    result["zeta"] = zeta_out;
    result["delta"] = delta_out;
    result["pdg_id"] = pdgid_out;
    result["at_element"] = at_element_out;
    result["at_turn"] = at_turn_out;
    result["mass_ratio"] = massratio_out;
    result["charge_ratio"] = chargeratio_out;
    result["parent_particle_id"] = trackid_out;

    return result;
}
