//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Resample/Processed/Slicer.cpp
//! @brief     Implements function Compute::Slicing::sliceFormFactor
//!
//! @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/Processed/Slicer.h"
#include "Base/Math/Functions.h"
#include "Base/Types/Span.h"
#include "Base/Util/Assert.h"
#include "Base/Vector/RotMatrix.h"
#include "Resample/Particle/ReCompound.h"
#include "Resample/Particle/ReMesocrystal.h"
#include "Resample/Particle/ReParticle.h"
#include "Resample/Slice/ZLimits.h"
#include "Sample/HardParticle/HardParticles.h"
#include "Sample/Material/MaterialUtils.h"
#include "Sample/Particle/Compound.h"
#include "Sample/Particle/CoreAndShell.h"
#include "Sample/Particle/Crystal.h"
#include "Sample/Particle/Mesocrystal.h"
#include "Sample/Particle/Particle.h"
#include "Sample/Scattering/Rotations.h"

namespace {

IFormFactor* doSlice(const IFormFactor* 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 = 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,
                                M_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->radius() + dz_bottom,
                                      f->radius() - dz_top);
    }
    if (const auto* f = dynamic_cast<const Sphere*>(ff)) {
        return new TruncatedSphere(f->radius(), f->height() - dz_bottom, dz_top);
    }
    if (const auto* f = dynamic_cast<const Spheroid*>(ff)) {
        double flattening = f->height() / (2.0 * f->radius());
        return new TruncatedSpheroid(f->radius(), f->height() - dz_bottom, flattening, dz_top);
    }
    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 TruncatedSphere*>(ff)) {
        return new TruncatedSphere(f->radius(), f->untruncated_height() - dz_bottom,
                                   dz_top + f->removedTop());
    }
    if (const auto* f = dynamic_cast<const TruncatedSpheroid*>(ff)) {
        return new TruncatedSpheroid(f->radius(), f->untruncated_height() - dz_bottom,
                                     f->heightFlattening(), dz_top + f->removedTop());
    } else
        throw std::runtime_error("Slicing of " + ff->className() + " not supported");
}

ReParticle* createParticleSlice(const IFormFactor* ff, ZLimits limits, const R3& position,
                                const IRotation* rot)
{
    const RotMatrix* rotMatrix =
        rot && !rot->isIdentity() ? new RotMatrix(rot->rotMatrix()) : nullptr;
    const Span span = ff->spanZ(rot) + position.z();

    if (span.hig() <= limits.low() || span.low() >= limits.hig())
        return nullptr;

    if (limits.low() <= span.low() && span.hig() <= limits.hig())
        return new ReParticle(ff->clone(), new R3(position), rotMatrix);

    if (!ff->canSliceAnalytically(rot))
        throw std::runtime_error("Slicing of " + ff->className()
                                 + " not supported for the given rotation");

    const double height = span.hig() - span.low();
    R3 new_position(position);
    double z_bottom = position.z();
    double z_top = position.z() + height;
    double dz_top = std::isinf(limits.hig()) ? -1 : z_top - limits.hig();
    double dz_bottom = std::isinf(limits.low()) ? -1 : limits.low() - z_bottom;
    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)
        new_position.setZ(limits.low());

    IFormFactor* slicedff = doSlice(ff, dz_bottom, dz_top);
    return new ReParticle(slicedff, new R3(new_position), rotMatrix);
}

//! Recursively processes the basis of a mesocrystal.
IReParticle* processBasis(const IParticle* basis, const Material& ambientMat)
{
    if (const auto* b = dynamic_cast<const Compound*>(basis)) {
        const auto& particles = b->decompose();
        ASSERT(!particles.empty());
        auto* result = new ReCompound;
        for (const auto* particle : particles) {
            std::unique_ptr<IReParticle> re(processBasis(particle, ambientMat));
            result->addFormFactor(*re);
        }
        return result;
    } else if (dynamic_cast<const CoreAndShell*>(basis)) {
        throw std::runtime_error("Mesocrystal with CoreAndShell basis not yet supported");
    }
    // Remaining case: the basis is a single particle.
    const auto* p = dynamic_cast<const Particle*>(basis);
    ASSERT(p);
    const IRotation* rot = p->rotation();
    auto particleSlice = std::make_unique<ReParticle>(
        p->pFormfactor()->clone(), new const R3(p->particlePosition()),
        rot && !rot->isIdentity() ? std::move(new const RotMatrix(rot->rotMatrix())) : nullptr);
    if (!particleSlice)
        return {};
    double volume = particleSlice->volume();
    Material transformed_material(p->rotation()
                                      ? p->material()->rotatedMaterial(p->rotation()->rotMatrix())
                                      : *p->material());
    particleSlice->setMaterial(transformed_material);
    particleSlice->setAdmixedFraction(volume);
    particleSlice->setAdmixedMaterial(transformed_material);
    particleSlice->setAmbientMaterial(ambientMat);
    return particleSlice.release();
}

} // namespace


OwningVector<IReParticle> Compute::Slicing::particlesInSlice(const IParticle* particle,
                                                             const ZLimits& limits,
                                                             const Material& ambientMat)
{
    if (const auto* p = dynamic_cast<const Particle*>(particle)) {
        ASSERT(p->pFormfactor());
        std::unique_ptr<ReParticle> particleSlice(
            createParticleSlice(p->pFormfactor(), limits, p->particlePosition(), p->rotation()));
        if (!particleSlice)
            return {};
        double volume = particleSlice->volume();
        Material transformed_material(
            p->rotation() ? p->material()->rotatedMaterial(p->rotation()->rotMatrix())
                          : *p->material());
        particleSlice->setMaterial(transformed_material);
        particleSlice->setAdmixedFraction(volume);
        particleSlice->setAdmixedMaterial(transformed_material);
        particleSlice->setAmbientMaterial(ambientMat);
        OwningVector<IReParticle> result;
        result.emplace_back(particleSlice.release());
        return result;

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

        // shell
        std::unique_ptr<Particle> P_shell(shell->clone());
        if (p->rotation())
            P_shell->rotate(*p->rotation());
        P_shell->translate(p->particlePosition());
        OwningVector<IReParticle> shell_slices =
            particlesInSlice(P_shell.get(), limits, ambientMat);
        if (shell_slices.empty())
            return {};
        ASSERT(shell_slices.size() == 1);
        IReParticle* sliced_shell = shell_slices.release_back();
        if (!sliced_shell)
            return {};
        const Material& shell_material = sliced_shell->admixed().material;

        // core
        std::unique_ptr<Particle> P_core(core->clone());
        if (p->rotation())
            P_core->rotate(*p->rotation());
        P_core->translate(p->particlePosition());
        OwningVector<IReParticle> core_slices =
            particlesInSlice(P_core.get(), limits, shell_material);
        ASSERT(core_slices.size() == 1);
        IReParticle* sliced_core = core_slices.release_back();

        // if core out of limits, return sliced shell
        if (!sliced_core) {
            OwningVector<IReParticle> result;
            result.emplace_back(sliced_shell);
            return result;
        }

        sliced_shell->setAdmixedFraction(sliced_shell->admixedFraction()
                                         - sliced_core->admixedFraction());

        OwningVector<IReParticle> result;
        result.emplace_back(sliced_shell);
        result.emplace_back(sliced_core);
        return result;

    } else if (dynamic_cast<const Compound*>(particle)) {
        throw std::runtime_error("Compound does not yet support slicing");

    } else if (const auto* p = dynamic_cast<const Mesocrystal*>(particle)) {
        const Crystal* crystal = &p->particleStructure();
        const IFormFactor* meso_formfactor = p->outerShape();
        ASSERT(crystal && meso_formfactor);

        double unit_cell_volume = crystal->lattice()->unitCellVolume();
        if (unit_cell_volume <= 0)
            return {};

        std::unique_ptr<ReParticle> new_shape(
            createParticleSlice(meso_formfactor, limits, p->particlePosition(), p->rotation()));

        const std::unique_ptr<Crystal> new_crystal(
            crystal->transformed(p->particlePosition(), p->rotation()));

        const std::unique_ptr<IReParticle> new_basis(
            processBasis(new_crystal->basis(), ambientMat));

        auto mc = std::make_unique<ReMesocrystal>(std::optional<size_t>{}, *new_crystal->lattice(),
                                                  *new_basis, *new_shape,
                                                  new_crystal->position_variance());

        std::vector<double> weights;
        std::vector<Material> materials;
        double volume = 0;
        for (const auto* p : crystal->basis()->decompose()) {
            volume += p->volume();
            weights.push_back(p->volume());
            materials.push_back(p->avgeMaterial());
        }
        mc->setAdmixedFraction(volume / unit_cell_volume);
        mc->setAdmixedMaterial(
            MaterialUtils::averagedMaterial("MesocrystalAvgeMat", weights, materials));

        OwningVector<IReParticle> result;
        result.emplace_back(mc.release());
        return result;

    } else
        ASSERT(0);
}