//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Resample/Slice/ComputeKz.cpp
//! @brief     Implements functions in namespace Compute.
//!
//! @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/ComputeKz.h"
#include "Base/Const/PhysicalConstants.h"
#include "Base/Const/Units.h"
#include "Base/Spin/SpinMatrix.h"
#include "Base/Util/Assert.h"
#include "Sample/Material/Material.h"
#include "Sample/Multilayer/Layer.h"
#include "Sample/Multilayer/Slice.h"
#include <numbers>

using std::numbers::pi;

namespace {

using PhysConsts::g_factor_n;
using PhysConsts::h_bar;
using PhysConsts::m_n;
using PhysConsts::mu_N;
using std::numbers::pi;

// Returns normalized SLD in nm^-2
complex_t normalizedSLD(const Material& material)
{
    ASSERT(material.hasNativeSLD());
    complex_t sld = std::conj(material.rawData()) / (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;
}

complex_t ScalarReducedPotential(const Material& mat, const R3& k, double n_ref)
{
    complex_t n = mat.refractiveIndex((2 * pi) / k.mag());
    return n * n - n_ref * n_ref * R3Util::sin2Theta(k);
}

/* currently not used
// The factor 1e-18 is here to have unit: 1/T*nm^-2
constexpr double magnetic_prefactor = (m_n * g_factor_n * mu_N / h_bar / h_bar) * 1e-18;

SpinMatrix polarizedReducedPotential(const Material& mat, const R3& b_field, const R3& k,
                                     double n_ref)
{
    double factor = magnetic_prefactor / k.mag2();
    complex_t unit_factor = ScalarReducedPotential(mat, k, n_ref);
    return MaterialUtil::MagnetizationCorrection(unit_factor, factor, b_field);
    return unit_factor * SpinMatrix::One() + factor * SpinMatrix::PauliMapping(b_field);
}
//*/

} // namespace


std::vector<complex_t> Compute::reducedKz(const std::vector<std::unique_ptr<const Slice>>& 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]->refractingMaterial()->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 layers with the same material as a previously computed layer, reuse that kz value
    // to ensure exact equality (important for the all-same-kz optimization in computeTR).
    for (size_t i = 0; i < N; ++i) {
        // Check if this slice has the same material as an earlier slice
        bool found_same = false;
        for (size_t j = 0; j < i; ++j) {
            if (*slices[i]->refractingMaterial() == *slices[j]->refractingMaterial()) {
                result[i] = result[j];
                found_same = true;
                break;
            }
        }
        if (found_same)
            continue;

        complex_t rad = ::ScalarReducedPotential(*slices[i]->refractingMaterial(), k, n_ref);
        if (i != i_ref)
            rad = checkForUnderflow(rad);
        result[i] = k_base * std::sqrt(rad);
    }
    return result;
}

std::vector<complex_t> Compute::kzFromSLDs(const std::vector<std::unique_ptr<const Slice>>& 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]->refractingMaterial());

    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]->refractingMaterial()));
        result[i] = k_sign * std::sqrt(kz2);
    }
    return result;
}

std::vector<complex_t>
Compute::kzFromRefractivity(const std::vector<std::unique_ptr<const Slice>>& 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]->refractingMaterial()->squaredRefractiveIndex(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]->refractingMaterial()->squaredRefractiveIndex(wl) - n2_ref;
        result[i] = k_sign * std::sqrt(checkForUnderflow(k2 * n2_norm + kz2));
    }
    return result;
}
