//  ************************************************************************************************
//
//  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/HardParticles.h"
#include "Sample/Lattice/Lattice3D.h"
#include "Sample/Material/Material.h"
#include "Sample/Material/Vacuum.h"
#include "Sample/Multilayer/Layer.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 <numbers>

using std::numbers::pi;

namespace {

//! 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());
    }
    if (const auto* f = dynamic_cast<const Bipyramid4*>(ff)) {
        if (dz_bottom > f->base_height()) {
            double dbase_edge = 2 * (dz_bottom - f->base_height()) * Math::cot(f->alpha());
            return new Pyramid4(f->length() - dbase_edge, f->height() - dz_bottom - dz_top,
                                f->alpha());
        }
        if (dz_top > f->heightRatio() * f->base_height()) {
            double dbase_edge = 2 * (f->base_height() - dz_bottom) * Math::cot(f->alpha());
            return new Pyramid4(f->length() - dbase_edge, f->height() - dz_bottom - dz_top,
                                pi - f->alpha());
        }
        {
            return new Bipyramid4(f->length(), f->base_height() - dz_bottom,
                                  f->heightRatio() * f->base_height() - dz_top, f->alpha());
        }
    }
    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);
    }
    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))
            throw std::runtime_error("Slicing of " + ff->className()
                                     + " not supported for the given rotation");

        // 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));
    }

    // 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->rotate(p->particleRotation());
        P_shell->translate(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->rotate(p->particleRotation());
        P_core->translate(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.rotate(p->particleRotation());
        temp_particle.translate(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->rotate(new_shape->particleRotation());
        mc->translate(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;
}
