//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Img3D/Build/BuilderUtil.cpp
//! @brief     Implements Img3D::BuilderUtils namespace.
//!
//! @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 "Img3D/Build/BuilderUtil.h"
#include "Base/Const/Units.h"
#include "Base/Util/Assert.h"
#include "Img3D/Build/Particle3DContainer.h"
#include "Img3D/Model/ParticleFromFF.h"
#include "Img3D/Model/Particles.h"
#include "Sample/Lattice/Lattice3D.h"
#include "Sample/Particle/Compound.h"
#include "Sample/Particle/CoreAndShell.h"
#include "Sample/Particle/Crystal.h"
#include "Sample/Particle/IFormfactor.h"
#include "Sample/Particle/Mesocrystal.h"
#include "Sample/Particle/Particle.h"
#include "Sample/Scattering/Rotations.h"

using Img3D::F3;
using Img3D::Particle3DContainer;

namespace {

R3 to_kvector(const F3& origin)
{
    return {static_cast<double>(origin.x()), static_cast<double>(origin.y()),
            static_cast<double>(origin.z())};
}

//! Returns the Euler angles from an IRotation object.
//! The result is encoded as F3, but does not have vector semantics.
Img3D::F3 EulerAngles(const IRotation*& rotation)
{
    double alpha = 0;
    double beta = 0;
    double gamma = 0;

    if (const auto* rotX = dynamic_cast<const RotationX*>(rotation)) {
        beta = rotX->angle(); // about x-axis
    } else if (const auto* rotY = dynamic_cast<const RotationY*>(rotation)) {
        alpha = Units::deg2rad(90);
        beta = rotY->angle(); // about y-axis
        gamma = Units::deg2rad(-90);
    } else if (const auto* rotZ = dynamic_cast<const RotationZ*>(rotation)) {
        alpha = rotZ->angle(); // about z-axis
    } else if (const auto* rotEuler = dynamic_cast<const RotationEuler*>(rotation)) {
        alpha = rotEuler->alpha();
        beta = rotEuler->beta();
        gamma = rotEuler->gamma();
    }
    return Img3D::F3fromR3({alpha, beta, gamma});
}

//! Apply transformations (translation, rotation) to a 3D Particle
//! or to a particle belonging to a Compound
void applyParticleTransformations(const Particle& particle, Img3D::PlotParticle& particle3D,
                                  const R3& origin)
{
    // rotation
    F3 particle_rotate;
    const IRotation* rotation = particle.rotation();

    if (rotation)
        particle_rotate = EulerAngles(rotation);

    // translation
    F3 position = Img3D::F3fromR3(particle.particlePosition() + origin);

    // If the particle belongs to a particle composition, along with the particle's
    // intrinsic transformations, position() and rotation() methods also account for the
    // translation and rotation (if present) of the particle composition as the
    // particleComposition's decompose() method already does this

    particle3D.addTransform(particle_rotate, position);
}

//! Apply transformations (translation, rotation) to a particle (core/shell) in a CoreAndShell
void applyCoreAndShellTransformations(const Particle& particle, Img3D::PlotParticle& particle3D,
                                      const CoreAndShell& particleCoreShell, const R3& origin)
{
    std::unique_ptr<Particle> P_clone(particle.clone()); // clone of the current particle

    // rotation
    F3 particle_rotate;
    const IRotation* rotationCoreShell = particleCoreShell.rotation();

    if (rotationCoreShell)
        P_clone->rotate(*rotationCoreShell);

    const IRotation* rotation = P_clone->rotation();

    if (rotation)
        particle_rotate = EulerAngles(rotation);

    // translation
    R3 positionCoreShell = particleCoreShell.particlePosition();

    P_clone->translate(positionCoreShell);

    F3 position = Img3D::F3fromR3(P_clone->particlePosition() + origin);

    particle3D.transform(particle_rotate, position);
}

} // namespace


Img3D::BuilderUtils::BuilderUtils(std::function<QColor(const QString&)> fnColorFromMaterialName)
    : m_fn_color_from_material_name(fnColorFromMaterialName)
{
    ASSERT(fnColorFromMaterialName);
}

void Img3D::BuilderUtils::applyParticleColor(const Particle& particle,
                                             Img3D::PlotParticle& particle3D, double alpha)
{
    // assign correct color to the particle from the knowledge of its material
    const Material* particle_material = particle.material();
    auto color =
        m_fn_color_from_material_name(QString::fromStdString(particle_material->materialName()));
    color.setAlphaF(alpha);
    particle3D.setColor(color);
}

Particle3DContainer Img3D::BuilderUtils::singleParticle3DContainer(const Particle& particle,
                                                                   double total_abundance,
                                                                   const F3& origin)
{
    std::unique_ptr<Particle> P_clone(particle.clone()); // clone of the particle

    const IFormfactor* ff = P_clone->pFormfactor();

    auto particle3D = Img3D::particle3DfromFF(ff);
    applyParticleTransformations(*P_clone, *particle3D, to_kvector(origin));
    applyParticleColor(*P_clone, *particle3D);

    Particle3DContainer result;
    result.addParticle3D(particle3D.release());
    result.setCumulativeAbundance(P_clone->abundance() / total_abundance);

    return result;
}

Particle3DContainer
Img3D::BuilderUtils::particleCoreShell3DContainer(const CoreAndShell& particleCoreShell,
                                                  double total_abundance, const F3& origin)
{
    // clone of the particleCoreShell
    std::unique_ptr<CoreAndShell> PCS_clone(particleCoreShell.clone());

    const auto* coreff = PCS_clone->coreParticle()->pFormfactor();
    const auto* shellff = PCS_clone->shellParticle()->pFormfactor();

    auto outCore = Img3D::particle3DfromFF(coreff);
    auto outShell = Img3D::particle3DfromFF(shellff);

    // core
    applyCoreAndShellTransformations(*PCS_clone->coreParticle(), *outCore, *PCS_clone,
                                     to_kvector(origin));
    applyParticleColor(*PCS_clone->coreParticle(), *outCore);

    // shell (set an alpha value of 0.5 for transparency)
    applyCoreAndShellTransformations(*PCS_clone->shellParticle(), *outShell, *PCS_clone,
                                     to_kvector(origin));
    applyParticleColor(*PCS_clone->shellParticle(), *outShell, 0.5);

    Particle3DContainer result;

    result.addParticle3D(outCore.release());  // index 0
    result.addParticle3D(outShell.release()); // index 1
    result.setCumulativeAbundance(PCS_clone->abundance() / total_abundance);

    return result;
}

Particle3DContainer
Img3D::BuilderUtils::particleComposition3DContainer(const Compound& particleComposition,
                                                    double total_abundance, const F3& origin)
{
    // clone of the particleComposition
    std::unique_ptr<Compound> PC_clone(particleComposition.clone());

    Particle3DContainer result;

    for (const auto* pc_particle : PC_clone->decompose()) {
        ASSERT(pc_particle);
        Particle3DContainer particle3DContainer;
        if (const auto* p = dynamic_cast<const CoreAndShell*>(pc_particle)) {
            particle3DContainer = particleCoreShell3DContainer(*p, 1.0, origin);
        } else if (dynamic_cast<const Mesocrystal*>(pc_particle)) {
            // TODO: Implement method to populate Mesocrystal from CORE and NOT from MesocrystalItem
            // as it is done currently in BuilderUtils::mesocrystal3DContainer
            throw std::runtime_error("Mesocrystal inside compound particle is not supported");
        } else if (const auto* p = dynamic_cast<const Particle*>(pc_particle))
            particle3DContainer = singleParticle3DContainer(*p, 1.0, origin);
        else
            ASSERT_NEVER;

        // keep result flat
        for (size_t i = 0; i < particle3DContainer.containerSize(); ++i) {
            result.addParticle3D(particle3DContainer.createParticle(i).release());
        }
    }

    // set the correct abundance for the entire Compound
    // (no abundances are associated with the individual components of Compound)
    result.setCumulativeAbundance(PC_clone->abundance() / total_abundance);

    return result;
}

Particle3DContainer Img3D::BuilderUtils::mesocrystal3DContainer(Mesocrystal* const mesocrystal,
                                                                double total_abundance,
                                                                const F3& origin)
{
    std::unique_ptr<Mesocrystal> M_clone(mesocrystal->clone()); // clone of the mesocrystal

    // These methods DO NOT add rotation/translation of the mesocrystal to its children
    // and hence they need to be added manually
    const auto* particleBasis = mesocrystal->particleStructure().basis();
    const auto* outerShapeff = mesocrystal->outerShape();

    const auto* mesocrystal_rotation = M_clone->rotation();
    const auto mesocrystal_translation = Img3D::F3fromR3(M_clone->particlePosition());

    Particle3DContainer outBasis;

    if (const auto* p = dynamic_cast<const Compound*>(particleBasis)) {
        outBasis = particleComposition3DContainer(*p, 1.0, origin);

    } else if (const auto* p = dynamic_cast<const CoreAndShell*>(particleBasis)) {
        outBasis = particleCoreShell3DContainer(*p, 1.0, origin);

    } else if (dynamic_cast<const Mesocrystal*>(particleBasis)) {
        // TODO: Implement method to populate Mesocrystal from CORE and NOT from MesocrystalItem
        // as it is done currently in Img3D::BuilderUtils::mesocrystal3DContainer
        throw std::runtime_error("Mesocrystal inside mesocrystal is not supported");

    } else if (const auto* p = dynamic_cast<const Particle*>(particleBasis)) {
        outBasis = singleParticle3DContainer(*p, 1.0, origin);

    } else
        ASSERT_NEVER;

    Particle3DContainer result;

    // Add inner particles for visualization
    std::vector<R3> basisPositions = mesocrystal->calcBasisPositions();
    for (const R3& position : basisPositions)
        for (size_t it = 0; it < outBasis.containerSize(); ++it) {
            auto particle3D = outBasis.createParticle(it);
            particle3D->addTranslation(Img3D::F3fromR3(position));
            particle3D->addExtrinsicRotation(EulerAngles(mesocrystal_rotation));
            particle3D->addTranslation(mesocrystal_translation);
            result.addParticle3D(particle3D.release());
        }

    // Add outer shape for visualization
    auto outerShape3D = Img3D::particle3DfromFF(outerShapeff);
    outerShape3D->addTransform(EulerAngles(mesocrystal_rotation), mesocrystal_translation + origin);

    // assign grey (default) color to the outer shape
    QColor color = {};
    color.setAlphaF(0.3);
    outerShape3D->setColor(color);
    result.addParticle3D(outerShape3D.release());

    // set the correct abundance for the Mesocrystal
    result.setCumulativeAbundance(M_clone->abundance() / total_abundance);

    return result;
}