//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Resample/Specular/ComputeFluxMagnetic.cpp
//! @brief     Implements functions to compute polarized fluxes
//!
//! @homepage  http://www.bornagainproject.org
//! @license   GNU General Public License v3 or higher (see COPYING)
//! @copyright Forschungszentrum Jülich GmbH 2020
//! @authors   Scientific Computing Group at MLZ (see CITATION, AUTHORS)
//
//  ************************************************************************************************

#include "Resample/Specular/ComputeFluxMagnetic.h"
#include "Base/Const/PhysicalConstants.h"
#include "Base/Spin/SpinMatrix.h"
#include "Base/Util/Assert.h"
#include "Resample/Flux/MatrixFlux.h"
#include "Resample/Slice/KzComputation.h"
#include "Resample/Slice/SliceStack.h"
#include "Resample/Specular/TransitionMagneticNevot.h"
#include "Resample/Specular/TransitionMagneticTanh.h"
#include "Sample/Interface/LayerRoughness.h"
#include <algorithm>
#include <numbers>
using std::numbers::pi;

namespace {

// The factor 1e-18 is here to have unit: 1/T*nm^-2
constexpr double magnetic_prefactor = PhysConsts::m_n * PhysConsts::g_factor_n * PhysConsts::mu_N
                                      / PhysConsts::h_bar / PhysConsts::h_bar / 4. / pi * 1e-18;

double magneticSLD(R3 B_field)
{
    return magnetic_prefactor * B_field.mag();
}

Spinor k_eigen(complex_t kz, double magnetic_SLD)
{
    const complex_t a = kz * kz;
    return {std::sqrt(a + 4. * pi * magnetic_SLD), std::sqrt(a - 4. * pi * magnetic_SLD)};
}

Spinor checkForUnderflow(const Spinor& eigenvs)
{
    auto k_eigen = [](complex_t value) { return std::abs(value) < 1e-40 ? 1e-40 : value; };
    return {k_eigen(eigenvs.u), k_eigen(eigenvs.v)};
}

std::pair<SpinMatrix, SpinMatrix> refractionMatrixBlocks(const MatrixFlux& tr_a,
                                                         const MatrixFlux& tr_b, double sigma,
                                                         const RoughnessModel& r_model)
{
    ASSERT(sigma >= 0);
    if (sigma < 10 * std::numeric_limits<double>::epsilon()) {
        const SpinMatrix kk = tr_a.computeInverseKappa() * tr_b.computeKappa();
        const SpinMatrix sp = (SpinMatrix::One() + kk) / 2;
        const SpinMatrix sm = (SpinMatrix::One() - kk) / 2;
        return {sp, sm};
    }

    if (r_model == RoughnessModel::NEVOT_CROCE)
        return Compute::refractionMatrixBlocksNevot(tr_a, tr_b, sigma);

    return Compute::refractionMatrixBlocksTanh(tr_a, tr_b, sigma);
}

//! Split off from Compute::polarizedFluxes just to facilitate the conversion from
//! vector<MatrixFluxes> to vector<Fluxes*>.
std::vector<MatrixFlux> computeTR(const SliceStack& slices, const std::vector<complex_t>& kzs,
                                  const RoughnessModel& r_model, bool forward)
{
    const size_t N = slices.size();
    ASSERT(kzs.size() == N);
    if (N == 0)
        return {};

    std::vector<MatrixFlux> TR;
    TR.reserve(N);

    const double kz_sign = kzs.front().real() >= 0.0 ? 1.0 : -1.0; // save sign to restore it later

    const R3& B_0 = slices.front().bField();
    TR.emplace_back(kz_sign, k_eigen(kzs.front(), 0.0), R3(), 0.0);
    for (size_t i = 1; i < slices.size(); ++i) {
        const R3 B = (forward ? +1 : -1) * (slices[i].bField() - B_0);
        const double magnetic_SLD = magneticSLD(B);
        TR.emplace_back(kz_sign, checkForUnderflow(k_eigen(kzs[i], magnetic_SLD)), B.unit_or_null(),
                        magnetic_SLD);
    }

    if (N == 1) {
        TR[0].m_T = SpinMatrix::One();
        TR[0].m_R = SpinMatrix::Null();
        return TR;
    }

    if (kzs[0] == 0.) { // TODO: cover by test
        TR[0].m_T = SpinMatrix::One();
        TR[0].m_R = SpinMatrix::Diag(-1, -1);
        for (size_t i = 1; i < N; ++i) {
            TR[i].m_T = SpinMatrix::Null();
            TR[i].m_R = SpinMatrix::Null();
        }
        return TR;
    }

    std::vector<SpinMatrix> FMatrices(N - 1);
    std::vector<complex_t> Norms(N - 1);

    // bottom boundary condition
    TR[N - 1].m_R = SpinMatrix::Null(); // holds x = t^-1 * r.

    for (size_t i = N - 1; i > 0; --i) {
        double sigma = 0.;
        if (const auto* roughness = slices.bottomRoughness(i))
            sigma = roughness->sigma();

        // compute the 2x2 blocks of the transfer matrix
        const auto [sp, sm] = refractionMatrixBlocks(TR[i - 1], TR[i], sigma, r_model);
        const SpinMatrix delta = TR[i - 1].computeDeltaMatrix(slices[i - 1].thicknessOr0());

        // compute the rotation matrix
        SpinMatrix E = sp + sm * TR[i].m_R;
        SpinMatrix F(E.d, -E.b, -E.c, E.a);
        const complex_t norm = F.determinant();
        F = F * delta;

        // store the rotation matrix and normalization constant in order to rotate
        // the coefficients for all lower slices at the end of the computation
        FMatrices[i - 1] = F;
        Norms[i - 1] = norm;

        // compute the reflection matrix and
        // rotate the polarization such that we have pure incoming states (T = I)
        F /= norm;

        // T is always equal to the identity at this point, no need to store
        TR[i - 1].m_R = delta * (sm + sp * TR[i].m_R) * F;
    }

    // now correct all amplitudes in forward direction by dividing with the remaining
    // normalization constants. In addition rotate the polarization by the amount
    // that was rotated above the current interface
    // if the normalization overflows, all amplitudes below that point are set to zero
    TR[0].m_T = SpinMatrix::One();
    complex_t normProduct = 1;
    SpinMatrix F = SpinMatrix::One();
    for (size_t i = 1; i < N; ++i) {
        normProduct = normProduct * Norms[i - 1];
        F = FMatrices[i - 1] * F;

        if (std::isinf(std::norm(normProduct))) {
            std::for_each(TR.begin() + i, TR.end(), [](auto& tr) {
                tr.m_T = SpinMatrix::Null();
                tr.m_R = SpinMatrix::Null();
            });
            break;
        }

        TR[i].m_T = F / normProduct; // T * F omitted, since T is always I
        TR[i].m_R *= F / normProduct;
    }

    return TR;
}

} // namespace

//  ************************************************************************************************
//  implemention of public interface
//  ************************************************************************************************

//! Returns fluxes for all sample slices.
//! See PhysRef, chapter "Polarized", section "Fluxes inside the sample".
Fluxes Compute::polarizedFluxes(const SliceStack& slices, const R3& k, bool forward)
{
    if (slices.size() > 1 && k.z() > 0)
        throw std::runtime_error(
            "source or detector below horizon not yet implemented for polarized scattering");
    std::vector<complex_t> kz = Compute::Kz::computeReducedKz(slices, k);
    ASSERT(slices.size() == kz.size());

    Fluxes result;
    for (const MatrixFlux& tr : ::computeTR(slices, kz, slices.roughnessModel(), forward))
        result.emplace_back(new MatrixFlux(tr));

    return result;
}

//! Returns matrix r_0 reflected from the top layer.
//! See PhysRef, chapter "Polarized", section "Generalized Parratt recursion".
SpinMatrix Compute::polarizedReflectivity(const SliceStack& slices,
                                          const std::vector<complex_t>& kzs, bool forward)
{
    const RoughnessModel& r_model = slices.roughnessModel();

    ASSERT(slices.size() == kzs.size());
    const size_t N = slices.size();
    if (N == 1)
        return {};
    if (kzs[0] == 0.)
        return -SpinMatrix::One();

    const R3& B_0 = slices.front().bField();
    const double kz_sign = kzs.front().real() >= 0.0 ? 1.0 : -1.0; // save sign to restore it later

    auto createCoeff = [&slices, &kzs, kz_sign, B_0, forward](size_t i) {
        const R3 B = (forward ? +1 : -1) * (slices[i].bField() - B_0);
        const double magnetic_SLD = ::magneticSLD(B);

        return MatrixFlux(kz_sign, ::checkForUnderflow(::k_eigen(kzs[i], magnetic_SLD)),
                          B.unit_or_null(), magnetic_SLD);
    };

    MatrixFlux tr_i1 = createCoeff(N - 1);

    // bottom boundary condition
    tr_i1.m_R = SpinMatrix::Null(); // holds x = t^-1 * r.

    for (size_t i = N - 1; i > 0; --i) {
        MatrixFlux tr_i = createCoeff(i - 1);

        double sigma = 0.;
        if (const auto* const roughness = slices.bottomRoughness(i - 1))
            sigma = roughness->sigma();

        // compute the 2x2 blocks of the transfer matrix
        const auto [sp, sm] = ::refractionMatrixBlocks(tr_i, tr_i1, sigma, r_model);
        const SpinMatrix delta = tr_i.computeDeltaMatrix(slices[i - 1].thicknessOr0());

        // compute the rotation matrix
        SpinMatrix E = sp + sm * tr_i1.m_R;
        SpinMatrix F(E.d, -E.b, -E.c, E.a);
        const complex_t norm = F.determinant();
        F = F * delta / norm;

        tr_i.m_R = delta * (sm + sp * tr_i1.m_R) * F;
        tr_i1 = tr_i;
    }
    return tr_i1.m_R;
}
