//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Sample/Sample/Sample.cpp
//! @brief     Implements class Sample.
//!
//! @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 "Sample/Sample/Sample.h"
#include "Base/Util/Assert.h"
#include "Sample/Aggregate/Struct2D.h"
#include "Sample/Aggregate/Struct3D.h"
#include "Sample/Material/Material.h"
#include "Sample/Multilayer/Layer.h"
#include "Sample/Multilayer/LayerStack.h"
#include "Sample/Multilayer/RoughnessUtil.h"
#include "Sample/Particle/Compound.h"
#include "Sample/Particle/CoreAndShell.h"
#include "Sample/Particle/Crystal.h"
#include "Sample/Particle/Mesocrystal.h"
#include "Sample/Particle/Mixture.h"
#include "Sample/Particle/Particle.h"
#include <algorithm>
#include <sstream>

namespace {

constexpr double inf = std::numeric_limits<double>::infinity();

void addMaterialsFromParticle(const IParticle* p, std::vector<const Material*>& out)
{
    if (!p)
        return;
    if (const auto* par = dynamic_cast<const Particle*>(p)) {
        if (const auto* mat = par->material())
            out.push_back(mat);
    } else if (const auto* cas = dynamic_cast<const CoreAndShell*>(p)) {
        addMaterialsFromParticle(cas->coreParticle(), out);
        addMaterialsFromParticle(cas->shellParticle(), out);
    } else if (const auto* comp = dynamic_cast<const Compound*>(p)) {
        for (const auto* cp : comp->components())
            addMaterialsFromParticle(cp, out);
    } else if (const auto* mix = dynamic_cast<const Mixture*>(p)) {
        for (const auto* mp : mix->particles())
            addMaterialsFromParticle(mp, out);
    } else if (const auto* meso = dynamic_cast<const Mesocrystal*>(p)) {
        addMaterialsFromParticle(meso->particleStructure().basis(), out);
    }
}

void addMaterialsFromILayer(const ILayer& ilayer, std::vector<const Material*>& out)
{
    if (const auto* lay = dynamic_cast<const Layer*>(&ilayer)) {
        if (const auto* mat = lay->material())
            out.push_back(mat);
        for (const auto& s : lay->structsAtBottom())
            addMaterialsFromParticle(s->particle(), out);
        for (const auto& s : lay->structsAtTop())
            addMaterialsFromParticle(s->particle(), out);
        for (const auto& s : lay->structs3D())
            addMaterialsFromParticle(s->particle(), out);
    } else if (const auto* stack = dynamic_cast<const LayerStack*>(&ilayer)) {
        for (const auto& sub : stack->sublayers())
            addMaterialsFromILayer(*sub, out);
    }
}

std::vector<const Material*> sampleMaterials(const Sample& sample)
{
    std::vector<const Material*> out;
    for (const auto& sub : sample.outerStack().sublayers())
        addMaterialsFromILayer(*sub, out);
    return out;
}

} // namespace


//! Takes ownership of given stack
Sample::Sample(LayerStack* outer_stack)
    : m_outer_stack(outer_stack)
{
    ASSERT(outer_stack);
}

Sample::Sample()
    : Sample(new LayerStack)
{
}

Sample::~Sample() = default;

void Sample::addLayer(const Layer& layer)
{
    ASSERT(!m_validated);
    m_outer_stack->addLayer(layer);
}

void Sample::addStack(const LayerStack& substack)
{
    ASSERT(!m_validated);
    m_outer_stack->addStack(substack);
}

void Sample::setExternalField(const R3& ext_field)
{
    ASSERT(!m_validated);
    m_ext_field = ext_field;
}

Sample* Sample::clone() const
{
    auto* result = new Sample(outerStack().clone());
    result->setExternalField(externalField());
    return result;
}

void Sample::validate() const
{
    if (m_validated)
        return;

    // Checks.
    ASSERT(!outerStack().sublayers().empty()); // ensured by LayerStack::validate
    if (outerStack().sublayers().size() < 2)
        throw std::runtime_error("Sample must have at least two layers (top and bottom)");
    const auto* layT = dynamic_cast<const Layer*>(outerStack().sublayers().front().get());
    if (!layT)
        throw std::runtime_error("Top layer is not a simple layer");
    else if (std::isfinite(layT->thickness()))
        throw std::runtime_error("Top layer has finite thickness");
    const auto* layB = dynamic_cast<const Layer*>(outerStack().sublayers().back().get());
    if (!layB)
        throw std::runtime_error("Bottom layer is not a simple layer");
    else if (std::isfinite(layB->thickness()))
        throw std::runtime_error("Bottom layer has finite thickness");
    // Check that inner layers have finite thickness (no infinite layers in the middle)
    for (size_t i = 1; i < outerStack().sublayers().size() - 1; ++i) {
        const auto* lay = dynamic_cast<const Layer*>(outerStack().sublayers()[i].get());
        if (lay && !std::isfinite(lay->thickness()))
            throw std::runtime_error("Inner layer " + std::to_string(i)
                                     + " has infinite thickness");
    }
    m_unwrapped = outerStack().flatClones();

    // Set z spans of unwrapped layers.
    const size_t N = m_unwrapped.size();
    ASSERT(N >= 2); // Ensured by validation above
    m_unwrapped[0]->setZSpan({0, inf});
    for (size_t i = 1; i < N; ++i) {
        const double top = m_unwrapped[i - 1]->low();
        m_unwrapped[i]->setZSpan({top - m_unwrapped[i]->thickness(), top});
    }

    for (size_t i = 1; i < N; ++i) {
        const double rms = RoughnessUtil::rms(m_unwrapped, i);
        m_unwrapped[i]->setRms(rms);
    }

    m_validated = true;
}

void Sample::checkRefractiveIndex(double wavelength) const
{
    for (const auto* mat : sampleMaterials(*this)) {
        const complex_t n = mat->refractiveIndex(wavelength);
        if (n.real() < 0.9 || n.real() > 1.1) {
            std::stringstream msg;
            msg << "Refractive index " << n << " at wavelength " << wavelength
                << " is too far from 1. Invalid material data?";
            throw std::runtime_error(msg.str());
        }
    }
}

bool Sample::isMagnetic() const
{
    const auto materials = sampleMaterials(*this);
    return std::any_of(materials.cbegin(), materials.cend(),
                       [](const Material* mat) { return mat->magnetized(); });
}

size_t Sample::nUnwrapped() const
{
    return unwrapped().size();
}

const Layer* Sample::unwrappedLayer(size_t i) const
{
    return unwrapped().at(i).get();
}

const std::vector<std::unique_ptr<Layer>>& Sample::unwrapped() const
{
    validate();
    return m_unwrapped;
}
