//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Sample/Multilayer/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/Multilayer/Sample.h"
#include "Base/Math/IntegratorGK.h"
#include "Base/Util/Assert.h"
#include "Base/Util/StringUtil.h"
#include "Sample/Aggregate/ParticleLayout.h"
#include "Sample/Interface/Roughness.h"
#include "Sample/Material/MaterialUtil.h"
#include "Sample/Multilayer/LayerStack.h"
#include <numbers>

using std::numbers::pi;

namespace {

const AutocorrelationModel* autocorrOf(const Layer* layer)
{
    return layer->roughness()->autocorrelationModel();
}

} // namespace


Sample::Sample()
    : m_outer_stack(new LayerStack)
{
    checkAndProcess();
}

Sample::~Sample() = default;

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

void Sample::checkMaterials(double wavelength) const
{
    m_outer_stack->checkMaterials(wavelength);
}

void Sample::setOuterStack(const LayerStack& outer_stack)
{
    m_outer_stack.reset(outer_stack.clone());
    checkAndProcess();
}

void Sample::addLayer(const Layer& layer)
{
    ASSERT(m_validated);
    if (!numberOfLayers()) // adding the top layer
        if (layer.thickness() != 0.0)
            throw std::runtime_error("Invalid top layer; to indicate that it is semiinfinite,"
                                     " it must have a nominal thickness of 0");
    m_outer_stack->addLayer(layer);
    checkAndProcess();
}

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

double Sample::maxCutoffSpatialFrequencyAt(size_t i_layer) const
{
    ASSERT(m_validated);
    double result = 0;
    for (size_t i = i_layer; i < numberOfLayers(); i++) {
        const auto autocorr = unwrappedLayers.at(i)->roughness()->autocorrelationModel();
        result = std::max(autocorr->maxSpatialFrequency(), result);
    }
    return result;
}

const Layer* Sample::layer(size_t i_layer) const
{
    ASSERT(m_validated);
    return unwrappedLayers.at(i_layer);
}

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

double Sample::roughnessSpectrum(double spatial_f, int i_layer) const
{
    ASSERT(m_validated);
    int size = numberOfLayers();
    const AutocorrelationModel* autocorr = autocorrOf(layer(i_layer));

    if (const auto k_corr = dynamic_cast<const SelfAffineFractalModel*>(autocorr))
        return k_corr->spectralFunction(spatial_f);

    if (dynamic_cast<const LinearGrowthModel*>(autocorr)) {
        if (i_layer == size - 1)
            ASSERT_NEVER;

        int j = i_layer + 1;
        for (; j < size; j++) {
            if (dynamic_cast<const SelfAffineFractalModel*>(autocorrOf(layer(j))))
                break;
        }
        ASSERT(j < size);
        const auto base_k_corr = dynamic_cast<const SelfAffineFractalModel*>(autocorrOf(layer(j)));
        double spectrum_below = base_k_corr->spectralFunction(spatial_f);
        double current_spectrum = spectrum_below;

        // all layers between i and j have linear growth model
        for (int k = j - 1; k >= i_layer; k--) {
            const auto lin_growth = dynamic_cast<const LinearGrowthModel*>(autocorrOf(layer(k)));
            const double thickness = unwrappedLayers.at(k)->thickness();
            current_spectrum = lin_growth->spectralFunction(spectrum_below, thickness, spatial_f);
            spectrum_below = current_spectrum;
        }
        return current_spectrum;
    }
    ASSERT_NEVER;
}

double Sample::roughnessRMS(size_t i_layer) const
{
    ASSERT(i_layer < numberOfLayers());
    const auto autocorr = autocorrOf(layer(i_layer));

    if (auto* k_corr = dynamic_cast<const SelfAffineFractalModel*>(autocorr)) {
        return k_corr->rms();
    } else if (dynamic_cast<const LinearGrowthModel*>(autocorr)) {
        const double maxSpatialFrequency = maxCutoffSpatialFrequencyAt(i_layer);
        return std::sqrt((2 * pi)
                         * RealIntegrator().integrate(
                             [this, i_layer](double spatial_f) -> double {
                                 return spatial_f * roughnessSpectrum(spatial_f, i_layer);
                             },
                             0, maxSpatialFrequency));
    }
    ASSERT_NEVER;
}

std::vector<const INode*> Sample::nodeChildren() const
{
    // skip outer stack in children list to hide it from exporting to python script
    return std::vector<const INode*>() << m_outer_stack->nodeChildren();
}

double Sample::hig(size_t i) const
{
    ASSERT(m_validated);
    ASSERT(i >= 1 && i < numberOfLayers());
    return ZInterfaces.at(i - 1);
}

double Sample::low(size_t i) const
{
    ASSERT(m_validated);
    ASSERT(i < numberOfLayers() - 1);
    return ZInterfaces.at(i);
}

size_t Sample::numberOfLayers() const
{
    ASSERT(m_validated);
    return unwrappedLayers.size();
}

std::string Sample::validate() const
{
    std::vector<std::string> errs;
    if (MaterialUtil::checkMaterialTypes(containedMaterials())
        == MATERIAL_TYPES::InvalidMaterialType)
        errs.push_back("Sample contains incompatible material definitions");

    std::string err = m_outer_stack->validate();
    if (!err.empty())
        errs.push_back("{ sample : " + err + " }");

    if (!errs.empty())
        return "[ " + Base::String::join(errs, ", ") + " ]";

    m_validated = true;
    return "";
}

void Sample::checkAndProcess() const
{
    // call after any change in layered structure

    m_validated = false;
    validateOrThrow();

    // unwrap layers and precompute interface coordinates
    unwrappedLayers = m_outer_stack->unwrapped();
    size_t N = unwrappedLayers.size();
    if (N >= 1)
        ZInterfaces.resize(N - 1);
    if (N >= 2) {
        ZInterfaces[0] = 0;
        for (size_t i = 1; i < N - 1; ++i)
            ZInterfaces.at(i) = ZInterfaces.at(i - 1) - unwrappedLayers.at(i)->thickness();
    }
}

std::string Sample::validateAmbientSubstrate() const
{
    const std::string err = validate();
    if (!err.empty())
        return err;

    // Non-empty stack cannot be the first or last component.
    // There should be ambient/substrate layer.
    auto outer_components = m_outer_stack->components();
    for (int i = outer_components.size() - 1; i >= 0; i--)
        if (const auto* stack = dynamic_cast<const LayerStack*>(outer_components[i]))
            if (stack->unwrapped().empty())
                outer_components.erase(outer_components.begin() + i);

    if (!outer_components.empty()) {
        if (dynamic_cast<const LayerStack*>(outer_components.front()))
            return "Sample: ambient layer is missing";

        if (dynamic_cast<const LayerStack*>(outer_components.back()))
            return "Sample: substrate layer is missing";
    }
    m_validated = true;
    return "";
}