//  ************************************************************************************************
//
//  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/Math/Functions.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 "Sample/Interface/Roughness.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;


//! Returns refraction matrix blocks s_{ab}^+-.
//! See PhysRef, chapter "Polarized", section "Interface with tanh profile".
std::pair<SpinMatrix, SpinMatrix> refractionMatrixBlocksTanh(const MatrixFlux& TR_a,
                                                             const MatrixFlux& TR_b, double sigma)
{
    ASSERT(sigma > 0);

    const double sigeff = std::sqrt(3.0) * sigma;
    complex_t rau = std::sqrt(Math::tanhc(sigeff * TR_a.k_eigen_up()));
    complex_t rad = std::sqrt(Math::tanhc(sigeff * TR_a.k_eigen_dn()));
    complex_t rbu = std::sqrt(Math::tanhc(sigeff * TR_b.k_eigen_up()));
    complex_t rbd = std::sqrt(Math::tanhc(sigeff * TR_b.k_eigen_dn()));

    SpinMatrix Rkk =
        TR_b.eigenToMatrix({rbu * TR_b.k_eigen_up(), rbd * TR_b.k_eigen_dn()})
        * TR_a.eigenToMatrix({1. / rau / TR_a.k_eigen_up(), 1. / rad / TR_a.k_eigen_dn()});

    SpinMatrix RInv = TR_a.eigenToMatrix({rau, rad}) * TR_b.eigenToMatrix({1. / rbu, 1. / rbd});

    const SpinMatrix sp = (RInv + Rkk) / 2.;
    const SpinMatrix sm = (RInv - Rkk) / 2.;

    return {sp, sm};
}

//! Returns refraction matrix blocks s_{ab}^+-.
//! See PhysRef, chapter "Polarized", section "Nevot-Croce approximation".
std::pair<SpinMatrix, SpinMatrix> refractionMatrixBlocksNevot(const MatrixFlux& TR_a,
                                                              const MatrixFlux& TR_b, double sigma)
{
    ASSERT(sigma > 0);

    auto roughness_matrix = [&TR_a, &TR_b, sigma](double sign) -> SpinMatrix {
        complex_t alpha_a = TR_a.k_eigen_up() + TR_a.k_eigen_dn();
        complex_t alpha_b = TR_b.k_eigen_up() + TR_b.k_eigen_dn();
        complex_t beta_a = TR_a.k_eigen_up() - TR_a.k_eigen_dn();
        complex_t beta_b = TR_b.k_eigen_up() - TR_b.k_eigen_dn();

        const complex_t alpha = alpha_b + sign * alpha_a;
        C3 b = beta_b * TR_b.field() + sign * beta_a * TR_a.field();

        auto square = [](auto& v) { return v.x() * v.x() + v.y() * v.y() + v.z() * v.z(); };
        complex_t beta = std::sqrt(square(b));
        if (std::abs(beta) < std::numeric_limits<double>::epsilon() * 10.) {
            const complex_t alpha_pp = -(alpha * alpha) * sigma * sigma / 8.;
            return {std::exp(alpha_pp), 0, 0, std::exp(alpha_pp)};
        }

        b /= beta;

        const complex_t alpha_pp = -(alpha * alpha + beta * beta) * sigma * sigma / 8.;
        const complex_t beta_pp = -alpha * beta * sigma * sigma / 4.;
        SpinMatrix Q(b.z() + 1., b.x() - I * b.y(), b.x() + I * b.y(), -1. - b.z());
        const SpinMatrix M(std::exp(beta_pp), 0, 0, std::exp(-beta_pp));

        return std::exp(alpha_pp) * Q * M * Q.adjoint() / (2. * (1. + b.z()));
    };

    const auto kk = SpinMatrix(TR_a.computeInverseKappa() * TR_b.computeKappa());
    const auto sp = 0.5 * (SpinMatrix::One() + kk) * roughness_matrix(-1.);
    const auto sm = 0.5 * (SpinMatrix::One() - kk) * roughness_matrix(+1.);

    return {sp, sm};
}

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 TransientModel* 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};
    }
    ASSERT(r_model);

    if (dynamic_cast<const ErfTransient*>(r_model))
        return refractionMatrixBlocksNevot(tr_a, tr_b, sigma);

    return 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,
                                  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) {
        const auto* roughness = slices.bottomRoughness(i);
        const double rms = slices.bottomRMS(i);
        const auto* r_model = roughness ? roughness->transient() : nullptr;

        // compute the 2x2 blocks of the transfer matrix
        const auto [sp, sm] = refractionMatrixBlocks(TR[i - 1], TR[i], rms, 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, forward))
        result.push_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)
{
    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);
        const auto* roughness = slices.bottomRoughness(i - 1);
        const double rms = slices.bottomRMS(i - 1);
        const auto* r_model = roughness ? roughness->transient() : nullptr;

        // compute the 2x2 blocks of the transfer matrix
        const auto [sp, sm] = ::refractionMatrixBlocks(tr_i, tr_i1, rms, 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;
}