//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Sample/Multilayer/RoughnessUtil.cpp
//! @brief     Implements utility functions for roughness calculations.
//!
//! @homepage  http://www.bornagainproject.org
//! @license   GNU General Public License v3 or higher (see COPYING)
//! @copyright Forschungszentrum Jülich GmbH 2024
//! @authors   Scientific Computing Group at MLZ (see CITATION, AUTHORS)
//
//  ************************************************************************************************

#include "Sample/Multilayer/RoughnessUtil.h"
#include "Base/Math/Integrator.h"
#include "Base/Util/Assert.h"
#include "Sample/Interface/AutocorrelationModels.h"
#include "Sample/Interface/Roughness.h"
#include "Sample/Multilayer/Layer.h"
#include <algorithm>
#include <numbers>

using std::numbers::pi;

double RoughnessUtil::spectrum(const std::vector<std::unique_ptr<Layer>>& flat, double spatial_f,
                               size_t i_layer)
{
    auto autocorrOf = [](const Layer* layer) -> const AutocorrelationModel* {
        ASSERT(layer->roughness());
        return layer->roughness()->autocorrelationModel();
    };

    const size_t N = flat.size();
    ASSERT(N >= 1);
    const AutocorrelationModel* autocorr = autocorrOf(flat.at(i_layer).get());

    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 == N - 1)
            ASSERT_NEVER;

        size_t j = i_layer + 1;
        for (; j < N; j++)
            if (dynamic_cast<const SelfAffineFractalModel*>(autocorrOf(flat.at(j).get())))
                break;
        ASSERT(j < N);
        const auto* base_k_corr =
            dynamic_cast<const SelfAffineFractalModel*>(autocorrOf(flat.at(j).get()));
        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 = int(j) - 1; k >= int(i_layer); k--) {
            const auto* lin_growth =
                dynamic_cast<const LinearGrowthModel*>(autocorrOf(flat.at(k).get()));
            const double thickness = flat.at(k)->thickness();
            current_spectrum = lin_growth->spectralFunction(spectrum_below, thickness, spatial_f);
            spectrum_below = current_spectrum;
        }
        return current_spectrum;
    }
    ASSERT_NEVER;
}

double RoughnessUtil::rms(const std::vector<std::unique_ptr<Layer>>& flat, std::size_t i_layer)
{
    ASSERT(i_layer < flat.size());
    if (!flat.at(i_layer)->roughness())
        return 0;

    const auto* autocorr = flat.at(i_layer)->roughness()->autocorrelationModel();

    if (const auto* k_corr = dynamic_cast<const SelfAffineFractalModel*>(autocorr)) {
        return k_corr->rms();
    } else if (dynamic_cast<const LinearGrowthModel*>(autocorr)) {
        double maxSpatialFrequency = 0;
        for (size_t i = i_layer; i < flat.size(); i++) {
            if (const Roughness* r = flat.at(i)->roughness()) {
                const auto* ac = r->autocorrelationModel();
                maxSpatialFrequency = std::max(maxSpatialFrequency, ac->maxSpatialFrequency());
            }
        }
        if (maxSpatialFrequency == 0)
            return 0;
        return std::sqrt((2 * pi)
                         * RealIntegrator().integrate(
                             [&flat, i_layer](double spatial_f) -> double {
                                 return spatial_f
                                        * RoughnessUtil::spectrum(flat, spatial_f,
                                                                  static_cast<int>(i_layer));
                             },
                             0, maxSpatialFrequency));
    }
    ASSERT_NEVER;
}

double RoughnessUtil::crossSpectrumWithInterfaceBelow(const Roughness& roughness, double thickness,
                                                      double spectrum_up, double spectrum_low,
                                                      double spatial_f)
{
    if (spectrum_up <= 0 || spectrum_low <= 0)
        return 0;

    if (const auto* crosscorr = dynamic_cast<const SpatialFrequencyCrosscorrelation*>(
            roughness.crosscorrelationModel()))
        return crosscorr->crosscorrSpectrum(spectrum_up, spectrum_low, thickness, spatial_f);

    if (const auto* linear_growth =
            dynamic_cast<const LinearGrowthModel*>(roughness.autocorrelationModel()))
        return linear_growth->crosscorrSpectrum(spectrum_low, thickness, spatial_f);

    return 0;
}
