//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Resample/Slice/KzComputation.cpp
//! @brief     Implements functions in namespace Compute::Kz.
//!
//! @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/KzComputation.h"
#include "Base/Const/Units.h"
#include "Base/Util/Assert.h"
#include "Resample/Slice/SliceStack.h"
#include <numbers>

using std::numbers::pi;

namespace {

// Returns normalized SLD in nm^-2
complex_t normalizedSLD(const Material& material)
{
    ASSERT(material.typeID() == MATERIAL_TYPES::MaterialBySLD);

    complex_t sld =
        std::conj(material.refractiveIndex_or_SLD()) / (Units::angstrom * Units::angstrom);
    sld *= 4.0 * pi;
    return sld;
}

complex_t checkForUnderflow(complex_t val)
{
    return std::norm(val) < 1e-80 ? complex_t(0.0, 1e-40) : val;
}

} // namespace

//  ************************************************************************************************
// namespace KzComputation
//  ************************************************************************************************

std::vector<complex_t> Compute::Kz::computeReducedKz(const SliceStack& slices, R3 k)
{
    const size_t N = slices.size();

    const size_t i_ref = k.z() > 0. ? N - 1 : 0;
    const double n_ref = slices[i_ref].material().refractiveIndex((2 * pi) / k.mag()).real();
    const double k_base = k.mag() * (k.z() > 0.0 ? -1 : 1);

    std::vector<complex_t> result(N);
    // Calculate refraction angle, expressed as k_z, for each layer.
    for (size_t i = 0; i < N; ++i) {
        complex_t rad = slices[i].scalarReducedPotential(k, n_ref);
        if (i != i_ref)
            rad = checkForUnderflow(rad);
        result[i] = k_base * std::sqrt(rad);
    }
    return result;
}

std::vector<complex_t> Compute::Kz::computeKzFromSLDs(const SliceStack& slices, double kz)
{
    const size_t N = slices.size();
    const double k_sign = kz > 0.0 ? -1 : 1;
    complex_t kz2_base = kz * kz + normalizedSLD(slices[0].material());

    std::vector<complex_t> result(N);
    result[0] = -kz;
    // Calculate refraction angle, expressed as k_z, for each layer.
    for (size_t i = 1; i < N; ++i) {
        complex_t kz2 = checkForUnderflow(kz2_base - normalizedSLD(slices[i].material()));
        result[i] = k_sign * std::sqrt(kz2);
    }
    return result;
}

std::vector<complex_t> Compute::Kz::computeKzFromRefIndices(const SliceStack& slices, R3 k)
{
    const size_t N = slices.size();
    const double kz = k.z();
    const double k_sign = kz > 0.0 ? -1 : 1;
    const double k2 = k.mag2();
    const double kz2 = kz * kz;
    const double wl = (2 * pi) / std::sqrt(k2);
    const complex_t n2_ref = slices[0].material().refractiveIndex2(wl);

    std::vector<complex_t> result(N);
    result[0] = -kz;
    for (size_t i = 1; i < N; ++i) {
        const complex_t n2_norm = slices[i].material().refractiveIndex2(wl) - n2_ref;
        result[i] = k_sign * std::sqrt(checkForUnderflow(k2 * n2_norm + kz2));
    }
    return result;
}
