//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Resample/Slice/ProfileHelper.cpp
//! @brief     Implements class ProfileHelper.
//!
//! @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/Slice/ProfileHelper.h"
#include "Base/Math/Functions.h"
#include "Base/Util/Assert.h"
#include "Sample/Interface/LayerRoughness.h"
#include <numbers>

using std::numbers::pi;

namespace {

double Transition(double x, const LayerRoughness* roughness)
{
    if (!roughness)
        return Math::Heaviside(x);

    return roughness->roughnessModel()->transient(x, roughness->sigma());
}

const std::string SLD = "SLD";
const std::string X = "X";
const std::string Y = "Y";
const std::string Z = "Z";

complex_t quantity(const Material& mat, std::string q)
{
    if (q == SLD)
        return mat.refractiveIndex_or_SLD();
    if (q == X)
        return mat.magnetization().x();
    if (q == Y)
        return mat.magnetization().y();
    if (q == Z)
        return mat.magnetization().z();
    ASSERT_NEVER;
}

std::vector<double> vec2real(const std::vector<complex_t>& v)
{
    std::vector<double> result(v.size());
    for (size_t i = 0; i < v.size(); i++)
        result[i] = real(v[i]);
    return result;
}

} // namespace


ProfileHelper::ProfileHelper(SliceStack stack)
    : m_stack(std::move(stack))
{
}

std::vector<complex_t> ProfileHelper::profile(const std::vector<double>& z_values,
                                              std::string component) const
{
    const complex_t top_value =
        !m_stack.empty() ? quantity(m_stack.at(0).material(), component) : 0.0;
    std::vector<complex_t> result(z_values.size(), top_value);
    for (size_t i = 1; i < m_stack.size(); ++i) {
        const Slice& slice = m_stack.at(i);
        const Slice& sliceAbove = m_stack.at(i - 1);
        const complex_t diff =
            quantity(slice.material(), component) - quantity(sliceAbove.material(), component);
        for (size_t j = 0; j < z_values.size(); ++j) {
            const double arg = (slice.hig() - z_values[j]);
            const double t = Transition(arg, slice.topRoughness());
            result[j] += diff * t;
        }
    }
    return result;
}

// Note: for refractive index materials, the material interpolation actually happens at the level
// of n^2. To first order in delta and beta, this implies the same smooth interpolation of delta
// and beta, as is done here.
std::vector<complex_t> ProfileHelper::calculateSLDProfile(const std::vector<double>& z_values) const
{
    return profile(z_values, SLD);
}

std::vector<double>
ProfileHelper::calculateMagnetizationProfile(const std::vector<double>& z_values,
                                             std::string xyz) const
{
    if (xyz != X && xyz != Y && xyz != Z)
        throw std::runtime_error("Incorrect magnetization component \"" + xyz + "\".\nOnly \"" + X
                                 + "\", \"" + Y + "\" or \"" + Z + "\" are allowed.");

    return ::vec2real(profile(z_values, xyz));
}

std::pair<double, double> ProfileHelper::defaultLimits() const
{
    if (m_stack.size() < 2)
        return {0.0, 0.0};
    double interface_span = m_stack.front().low() - m_stack.back().hig();
    double default_margin = interface_span > 0.0 ? interface_span / 20.0 : 10.0;
    const LayerRoughness* topRoughness = m_stack.at(1).topRoughness();
    const LayerRoughness* bottomRoughness = m_stack.back().topRoughness();

    double top_margin =
        topRoughness && topRoughness->sigma() > 0 ? 5.0 * topRoughness->sigma() : default_margin;
    double bottom_margin = bottomRoughness && bottomRoughness->sigma() > 0
                               ? 5.0 * bottomRoughness->sigma()
                               : default_margin;
    double z_min = m_stack.back().hig() - bottom_margin;
    double z_max = m_stack.front().low() + top_margin;
    return {z_min, z_max};
}
