//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Resample/Refine/ParticlesInSlice.cpp
//! @brief     Implements function Refine::particlesInSlice.
//!
//! @homepage  http://www.bornagainproject.org
//! @license   GNU General Public License v3 or higher (see COPYING)
//! @copyright Forschungszentrum Jülich GmbH 2018
//! @authors   Scientific Computing Group at MLZ (see CITATION, AUTHORS)
//
//  ************************************************************************************************

#include "Resample/Refine/ParticlesInSlice.h"
#include "Base/Math/Functions.h"
#include "Base/Math/Numeric.h"
#include "Base/Type/Span.h"
#include "Base/Util/Assert.h"
#include "Base/Vector/RotMatrix.h"
#include "Sample/HardParticle/FormfactorPolyhedron.h"
#include "Sample/HardParticle/HardParticles.h"
#include "Sample/Lattice/Lattice3D.h"
#include "Sample/Material/Material.h"
#include "Sample/Material/Vacuum.h"
#include "Sample/Multilayer/Slice.h"
#include "Sample/Option/SimulationOptions.h"
#include "Sample/Particle/Compound.h"
#include "Sample/Particle/CoreAndShell.h"
#include "Sample/Particle/Crystal.h"
#include "Sample/Particle/HardFormfactor.h"
#include "Sample/Particle/Mesocrystal.h"
#include "Sample/Particle/Mixture.h"
#include "Sample/Particle/Particle.h"
#include "Sample/Particle/Rotations.h"
#include <algorithm>
#include <cmath>
#include <ff/Polyhedron.h>
#include <limits>
#include <numbers>

using std::numbers::pi;

namespace {

//! Clears all symmetry flags. They are validated against the origin and break once the polyhedron
//! is shifted to a bottom origin or rotated; clearing them selects the general form factor path.
ff::Topology desymmetrized(ff::Topology topology)
{
    topology.symmetry_Ci = false;
    for (ff::FacialTopology& face : topology.faces)
        face.symmetry_S2 = false;
    return topology;
}

//! Internal form factor wrapping an arbitrary polyhedron (topology + bottom-origin vertices), the
//! result of generic slicing. Symmetry flags are cleared; not exposed to users.
class FormfactorClippedPolyhedron : public FormfactorPolyhedron {
public:
    FormfactorClippedPolyhedron(const ff::Topology& topology, const std::vector<R3>& vertices)
    {
        setPolyhedron(desymmetrized(topology), 0., vertices); // vertices already bottom-origin
        validate();
    }
    FormfactorClippedPolyhedron* clone() const override
    {
        const auto poly = getPolyhedron();
        return new FormfactorClippedPolyhedron(poly->topology(), poly->vertices());
    }
    std::string className() const final { return "ClippedPolyhedron"; }
};

//! Slices a rotated polyhedron: rotate into the slice frame, clip by the horizontal slice planes,
//! rotate back. The particle keeps its rotation; its bottom-origin position is returned via
//! `position`.
std::unique_ptr<Formfactor> clipRotatedPolyhedron(const FormfactorPolyhedron* ff,
                                                  const Rotation* rotation, R3& position,
                                                  const Span& limits)
{
    const std::shared_ptr<ff::Polyhedron> src = ff->getPolyhedron();
    ASSERT(src);
    const RotMatrix R = rotation->rotMatrix();
    const R3 x0 = position;
    const double zb = src->z_bottom();
    // rotate into the slice frame, kept relative to x0 (small coordinates ease libformfactor)
    std::vector<R3> rel;
    rel.reserve(src->vertices().size());
    double z_lo = std::numeric_limits<double>::max();
    double z_hi = std::numeric_limits<double>::lowest();
    for (const R3& v : src->vertices()) {
        const R3 Y = R.transformed(R3(v.x(), v.y(), v.z() - zb));
        rel.push_back(Y);
        z_lo = std::min(z_lo, Y.z());
        z_hi = std::max(z_hi, Y.z());
    }
    const ff::Polyhedron poly(desymmetrized(src->topology()), rel);
    // clip planes in the x0-relative frame; an infinite limit means no cut on that side.
    // Guard: if the requested boundary is within a small tolerance of the computed z-extent,
    // extend past the extent rather than clipping.  Two independent FP paths (PolyhedralUtil
    // and this vertex loop) compute the same z-coordinate via different formula and can
    // disagree by a few ULP; attempting to clip at such a near-coincident plane triggers a
    // "vanishing edge" error in libformfactor.
    const double zlo_raw = std::isinf(limits.low()) ? z_lo - 1 : limits.low() - x0.z();
    const double zhi_raw = std::isinf(limits.hig()) ? z_hi + 1 : limits.hig() - x0.z();
    const double tol = 1e-10 * (z_hi - z_lo);
    const double zlo = zlo_raw - z_lo < tol ? z_lo - 1 : zlo_raw;
    const double zhi = z_hi - zhi_raw < tol ? z_hi + 1 : zhi_raw;
    try {
        const auto clipped = std::unique_ptr<ff::Polyhedron>(poly.clipped(zlo, zhi));
        // rotate back into the particle frame and locate the new bottom
        const RotMatrix Rinv = R.Inverse();
        std::vector<R3> local;
        local.reserve(clipped->vertices().size());
        double zmin = std::numeric_limits<double>::max();
        for (const R3& c : clipped->vertices()) {
            const R3 l = Rinv.transformed(c);
            local.push_back(l);
            zmin = std::min(zmin, l.z());
        }
        for (R3& l : local)
            l = R3(l.x(), l.y(), l.z() - zmin);
        position = x0 + R.transformed(R3(0, 0, zmin)); // compensate the bottom-origin shift
        return std::make_unique<FormfactorClippedPolyhedron>(clipped->topology(), local);
    } catch (const std::exception& e) {
        throw std::runtime_error(
            "Cannot slice " + ff->className() + ": libformfactor clipping failed (" + e.what()
            + "). A slice plane may nearly coincide with a vertex; try shifting "
              "the particle position or the slice or layer boundaries.");
    }
}

//! Slices a polyhedron lacking an analytic slicer, for an unrotated or z-invariant particle: the
//! identity-rotation case of clipRotatedPolyhedron (the caller anchors the fragment via the
//! limits).
Formfactor* clipPolyhedronAxisAligned(const FormfactorPolyhedron* ff, double dz_bottom,
                                      double dz_top)
{
    const double height = ff->zTop(nullptr); // bottom origin at z=0
    const Span limits(dz_bottom > 0 ? dz_bottom : -Span::inf,
                      dz_top > 0 ? height - dz_top : Span::inf);
    const IdentityRotation identity;
    R3 discarded_position;
    return clipRotatedPolyhedron(ff, &identity, discarded_position, limits).release();
}

//! Slices a formfactor by removing material from top and/or bottom.
//! Uses dynamic_cast to dispatch to geometry-specific slicing algorithms. Each formfactor type
//! requires its own slicing logic based on its geometric properties. A virtual method approach
//! would be possible but would spread the slicing logic across many formfactor classes.
Formfactor* doSlice(const Formfactor* ff, double dz_bottom, double dz_top)
{
    if (const auto* f = dynamic_cast<const Pyramid2*>(ff)) {
        double dbase_edge = 2 * dz_bottom * Math::cot(f->alpha());
        return new Pyramid2(f->length() - dbase_edge, f->width() - dbase_edge,
                            f->height() - dz_bottom - dz_top, f->alpha());
    }
    if (const auto* f = dynamic_cast<const Box*>(ff))
        return new Box(f->length(), f->width(), f->height() - dz_bottom - dz_top);
    if (const auto* f = dynamic_cast<const Cone*>(ff)) {
        double dradius = dz_bottom * Math::cot(f->alpha());
        return new Cone(f->radius() - dradius, f->height() - dz_bottom - dz_top, f->alpha());
    }
    if (const auto* f = dynamic_cast<const Pyramid6*>(ff)) {
        double dbase_edge = 2 / sqrt(3) * dz_bottom * Math::cot(f->alpha());
        return new Pyramid6(f->baseEdge() - dbase_edge, f->height() - dz_bottom - dz_top,
                            f->alpha());
    }
    // Bipyramid4 is a polyhedron; it has no correct analytic slicer (a top cut into the lower
    // frustum yields an inverted pyramid Pyramid4 cannot represent), so it falls through to the
    // generic clipper below.
    if (const auto* f = dynamic_cast<const Cylinder*>(ff))
        return new Cylinder(f->radius(), f->height() - dz_bottom - dz_top);
    if (const auto* f = dynamic_cast<const EllipsoidalCylinder*>(ff)) {
        return new EllipsoidalCylinder(f->radiusX(), f->radiusY(),
                                       f->height() - dz_bottom - dz_top);
    }
    if (const auto* f = dynamic_cast<const HorizontalCylinder*>(ff)) {
        return new HorizontalCylinder(f->radius(), f->length(), f->cutFromTop() + dz_top,
                                      f->cutFromBottom() + dz_bottom);
    }
    if (const auto* f = dynamic_cast<const Sphere*>(ff))
        return new SphericalSegment(f->radius(), dz_top, dz_bottom);
    if (const auto* f = dynamic_cast<const Spheroid*>(ff))
        return new SpheroidalSegment(f->radiusXY(), f->radiusZ(), dz_top, dz_bottom);
    if (const auto* f = dynamic_cast<const Ellipsoid*>(ff))
        return new EllipsoidalSegment(f->radiusX(), f->radiusY(), f->radiusZ(), dz_top, dz_bottom);
    if (const auto* f = dynamic_cast<const LongBoxGauss*>(ff))
        return new LongBoxGauss(f->length(), f->width(), f->height() - dz_bottom - dz_top);
    if (const auto* f = dynamic_cast<const LongBoxLorentz*>(ff))
        return new LongBoxLorentz(f->length(), f->width(), f->height() - dz_bottom - dz_top);
    if (const auto* f = dynamic_cast<const Prism3*>(ff))
        return new Prism3(f->baseEdge(), f->height() - dz_bottom - dz_top);
    if (const auto* f = dynamic_cast<const Prism6*>(ff))
        return new Prism6(f->baseEdge(), f->height() - dz_bottom - dz_top);
    if (const auto* f = dynamic_cast<const Pyramid4*>(ff)) {
        double dbaseEdge = 2 * dz_bottom * Math::cot(f->alpha());
        return new Pyramid4(f->baseEdge() - dbaseEdge, f->height() - dz_bottom - dz_top,
                            f->alpha());
    }
    if (const auto* f = dynamic_cast<const Pyramid3*>(ff)) {
        double dbaseEdge = 2 * sqrt(3) * dz_bottom * Math::cot(f->alpha());
        return new Pyramid3(f->baseEdge() - dbaseEdge, f->height() - dz_bottom - dz_top,
                            f->alpha());
    }
    if (const auto* f = dynamic_cast<const SphericalSegment*>(ff)) {
        return new SphericalSegment(f->radius(), f->cutFromTop() + dz_top,
                                    f->cutFromBottom() + dz_bottom);
    }
    if (const auto* f = dynamic_cast<const SpheroidalSegment*>(ff)) {
        return new SpheroidalSegment(f->radiusXY(), f->radiusZ(), f->cutFromTop() + dz_top,
                                     f->cutFromBottom() + dz_bottom);
    }
    if (const auto* f = dynamic_cast<const EllipsoidalSegment*>(ff)) {
        return new EllipsoidalSegment(f->radiusX(), f->radiusY(), f->radiusZ(),
                                      f->cutFromTop() + dz_top, f->cutFromBottom() + dz_bottom);
    }
    // rest of polyhedra: no shape-specific slicer, fall back to generic clipping
    if (const auto* f = dynamic_cast<const FormfactorPolyhedron*>(ff))
        return clipPolyhedronAxisAligned(f, dz_bottom, dz_top);
    throw std::runtime_error("Slicing of " + ff->className() + " not supported");
}

std::unique_ptr<Particle> createParticleSlice(const Particle* p, const Formfactor* ff, Span limits,
                                              const Slice* slice)
{
    // Input coordinates:
    // - p->particlePosition(): particle position in slice-local coordinates (slice.hig() = 0)
    // - ff: formfactor in particle-local coords (z typically in [0, height])
    // - limits: slice z-boundaries in slice-local coordinates
    R3 position = p->particlePosition(); // slice-local
    const Rotation* rot = p->particleRotation();

    // particle z-extent in slice-local coordinates
    const double span_low = ff->zBottom(rot) + position.z(); // slice-local
    const double span_hig = ff->zTop(rot) + position.z();    // slice-local
    if (span_hig <= limits.low() || span_low >= limits.hig())
        return nullptr;

    std::unique_ptr<Formfactor> slicedff;
    if (limits.low() <= span_low && span_hig <= limits.hig()) {
        // Particle fully contained in slice - no slicing needed
        slicedff.reset(ff->clone());

    } else if (ff->canSliceAnalytically(rot)) {
        // No rotation, or z-invariant rotation: slice planes stay horizontal in the particle
        // frame. All z values below are in slice-local coordinates.
        const double height = span_hig - span_low; // particle-local (size)
        double z_bottom = position.z();            // slice-local
        double z_top = position.z() + height;      // slice-local
        // dz_top/dz_bottom: amount to cut from top/bottom (particle-local sizes)
        double dz_top = std::isinf(limits.hig()) ? -1 : z_top - limits.hig(); // particle-local
        double dz_bottom =
            std::isinf(limits.low()) ? -1 : limits.low() - z_bottom; // particle-local
        ASSERT(dz_top >= 0 || dz_bottom >= 0);
        ASSERT(dz_bottom <= height);
        ASSERT(dz_top <= height);
        if (dz_bottom < 0)
            dz_bottom = 0;
        if (dz_top < 0)
            dz_top = 0;
        if (dz_bottom > 0)
            position.setZ(limits.low()); // adjust particle position to slice bottom

        // slicedff: new formfactor in slice-local coords (z in [0, new_height])
        slicedff.reset(doSlice(ff, dz_bottom, dz_top));

    } else if (const auto* poly = dynamic_cast<const FormfactorPolyhedron*>(ff)) {
        // Tilted polyhedron: slice planes are oblique in the particle frame, so clip in slice/lab
        // coordinates. Updates `position` to the bottom-origin-adjusted particle position.
        slicedff = clipRotatedPolyhedron(poly, rot, position, limits);

    } else
        throw std::runtime_error("Slicing of " + ff->className()
                                 + " not supported for the given rotation");

    // Create a new Particle with the sliced formfactor
    auto result = std::unique_ptr<Particle>(p->clone());
    result->setFormfactor(slicedff.release());
    // Reset position to the adjusted one (for sliced particles)
    // Note: We need to use translate to adjust position since clone preserves original position
    R3 pos_delta = position - p->particlePosition();
    if (pos_delta != R3())
        result->translate(pos_delta);
    result->setSliceNumber(slice->number());
    return result;
}

//! Returns the intra-cell residual of the slice-induced local shift.
R3 mesocrystalResidualShift(const Mesocrystal& source, const Particle& sliced_shape)
{
    R3 local_shift = source.particlePosition() - sliced_shape.particlePosition();
    if (const auto* rot = source.particleRotation())
        local_shift = rot->rotMatrix().Inverse().transformed(local_shift);

    const auto* lattice = source.particleStructure().lattice();
    ASSERT(lattice);

    R3 ra, rb, rc;
    lattice->reciprocalLatticeBasis(ra, rb, rc);

    // Keep the shifted basis within one unit cell of the sliced mesocrystal.
    const I3 lattice_shift(static_cast<int>(std::lround(local_shift.dot(ra) / (2 * pi))),
                           static_cast<int>(std::lround(local_shift.dot(rb) / (2 * pi))),
                           static_cast<int>(std::lround(local_shift.dot(rc) / (2 * pi))));

    return local_shift - lattice_shift.x() * lattice->basisVectorA()
           - lattice_shift.y() * lattice->basisVectorB()
           - lattice_shift.z() * lattice->basisVectorC();
}

} // namespace

//! Slices a particle to fit within the given z limits.
//! Uses dynamic_cast to handle each leaf IParticle subclass differently:
//! - Particle: slices formfactor, sets up material
//! - CoreAndShell: recursively slices shell and core, handles overlapping regions
//! - Compound: not a leaf particle
//! - Mixture: not a leaf particle
//! - Mesocrystal: slices outer shape, prepares for computation
//! A virtual method approach would create a dependency from Sample to Resample.
std::vector<std::unique_ptr<IParticle>>
Refine::particlesInSlice(const IParticle* particle, const Span& limits, const Slice* slice,
                         const Material& ambient_mat, const MesoOptions& meso_options)
{
    ASSERT(particle);
    ASSERT(!dynamic_cast<const Compound*>(particle));
    ASSERT(!dynamic_cast<const Mixture*>(particle));

    if (const auto* p = dynamic_cast<const Particle*>(particle)) {
        ASSERT(p->pFormfactor());
        std::unique_ptr<Particle> sliced_particle(
            ::createParticleSlice(p, p->pFormfactor(), limits, slice));
        if (!sliced_particle)
            return {};

        std::unique_ptr<Material> transformed_material(
            p->particleRotation()
                ? p->material()->rotatedMaterial(p->particleRotation()->rotMatrix())
                : std::unique_ptr<Material>(p->material()->clone()));
        sliced_particle->setMaterial(*transformed_material);
        sliced_particle->setAmbientMaterial(ambient_mat);

        std::vector<std::unique_ptr<IParticle>> result;
        result.emplace_back(std::move(sliced_particle));
        return result;
    }

    if (const auto* p = dynamic_cast<const CoreAndShell*>(particle)) {
        const IParticle* core = p->coreParticle();
        const Particle* shell = p->shellParticle();
        ASSERT(core && shell);

        // shell
        std::unique_ptr<Particle> P_shell(shell->clone());
        P_shell->applyParentTransform(p->particleRotation(), p->particlePosition());
        std::vector<std::unique_ptr<IParticle>> shell_slices =
            particlesInSlice(P_shell.get(), limits, slice, ambient_mat, meso_options);
        if (shell_slices.empty())
            return {};
        ASSERT(shell_slices.size() == 1);
        std::unique_ptr<IParticle> sliced_shell(shell_slices.back().release());
        ASSERT(sliced_shell);
        auto* sliced_shell_particle = dynamic_cast<Particle*>(sliced_shell.get());
        ASSERT(sliced_shell_particle);

        // core
        std::unique_ptr<IParticle> P_core(core->clone());
        P_core->applyParentTransform(p->particleRotation(), p->particlePosition());
        std::vector<std::unique_ptr<IParticle>> core_slices = particlesInSlice(
            P_core.get(), limits, slice, *sliced_shell_particle->material(), meso_options);

        // if core out of limits, return sliced shell only
        if (core_slices.empty()) {
            std::vector<std::unique_ptr<IParticle>> result;
            result.emplace_back(std::move(sliced_shell));
            return result;
        }

        double core_volume = 0.0;
        for (const auto& core_slice : core_slices)
            core_volume += core_slice->volumeToAdmix();
        const double shell_volume = sliced_shell_particle->pFormfactor()->volume();
        // The sliced shell still keeps its outer geometry; subtract the sliced core volume from it.
        ASSERT(Numeric::lessOrAlmostEqual(core_volume, shell_volume, 1e3));
        sliced_shell_particle->setVolumeToAdmix(std::max(0.0, shell_volume - core_volume));

        core_slices.emplace_back(std::move(sliced_shell));
        return core_slices;
    }

    if (const auto* p = dynamic_cast<const Mesocrystal*>(particle)) {
        // Create a temporary Particle to use for outer shape slicing.
        // The mesocrystal's outerShape is just a formfactor in mesocrystal-local coords.
        // We need to transform it to slice-local coords using mesocrystal's position/rotation.
        Vacuum dummy_mat;
        Particle temp_particle(dummy_mat, *p->outerShape());
        temp_particle.applyParentTransform(p->particleRotation(), p->particlePosition());
        // temp_particle now has position in slice-local coords

        // createParticleSlice returns a particle with:
        // - position in slice-local coords (adjusted if sliced from bottom)
        // - formfactor in particle-local coords (z in [0, slice_height])
        auto new_shape = ::createParticleSlice(&temp_particle, p->outerShape(), limits, slice);
        if (!new_shape)
            return {};

        // Open the shared top boundary by one ULP so adjacent sliced mesocrystals
        // do not count interface particles twice.
        const bool has_sliced_mesocrystal_above = new_shape->topZ() < temp_particle.topZ();
        if (has_sliced_mesocrystal_above) {
            const auto* ff = new_shape->pFormfactor();
            ASSERT(ff);

            const Rotation* rot = new_shape->particleRotation();
            const double height = ff->zTop(rot) - ff->zBottom(rot);
            const double open_top_height = std::nextafter(height, 0.0);
            new_shape->setFormfactor(doSlice(ff, 0.0, height - open_top_height));
        }

        const auto* sliced_ff = dynamic_cast<const HardFormfactor*>(new_shape->pFormfactor());
        if (!sliced_ff)
            throw std::runtime_error("Mesocrystal slicing requires a hard outer shape");

        // Canonical sliced crystal: sliced outer shape + residual intra-cell basis shift.
        const R3 residual_shift = ::mesocrystalResidualShift(*p, *new_shape);
        const auto& source_crystal = p->particleStructure();
        const auto shifted_crystal =
            std::unique_ptr<Crystal>(source_crystal.transformed(residual_shift, nullptr));
        auto mc = std::make_unique<Mesocrystal>(*shifted_crystal, *sliced_ff);
        mc->applyParentTransform(new_shape->particleRotation(), new_shape->particlePosition());
        mc->setSliceNumber(slice->number());
        mc->setVolumeToAdmix(p->slicedVolumeToAdmix(*new_shape));
        mc->prepareForComputation(ambient_mat, meso_options);

        std::vector<std::unique_ptr<IParticle>> result;
        result.emplace_back(std::move(mc));
        return result;
    }
    ASSERT_NEVER;
}
