//  ************************************************************************************************
//
//  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/Slice/ComputeKz.h"
#include "Sample/Flux/MatrixFlux.h"
#include "Sample/Interface/Roughness.h"
#include "Sample/Multilayer/Layer.h"
#include "Sample/Multilayer/Slice.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 std::vector<std::unique_ptr<const Slice>>& slices,
                                  const std::vector<complex_t>& kzs, bool forward, bool top_exit)
{
    const size_t N = slices.size();
    ASSERT(kzs.size() == N);
    if (N == 0)
        return {};

    // Index reversal for bottom exit (like in scalar case)
    std::vector<size_t> X(N, 0);
    for (size_t i = 0; i < N; ++i)
        X[i] = top_exit ? i : N - 1 - i;

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

    // Build MatrixFlux objects. Index i is the logical index (0=entry, N-1=exit).
    // X[i] maps to physical layer index.
    std::vector<MatrixFlux> TR;
    TR.reserve(N);
    for (size_t i = 0; i < N; ++i) {
        const size_t j = X[i]; // physical layer index
        if (i == 0) {
            TR.emplace_back(kz_sign, k_eigen(kzs[j], 0.0), R3(), 0.0);
        } else {
            const R3 B = (forward ? +1 : -1) * (slices[j]->bField() - B_0);
            const double magnetic_SLD = magneticSLD(B);
            TR.emplace_back(kz_sign, checkForUnderflow(k_eigen(kzs[j], magnetic_SLD)),
                            B.unit_or_null(), magnetic_SLD);
        }
    }
    // Now TR[i] corresponds to logical index i, but stores data for physical layer X[i]

    // Helper: reorder TR from logical to physical order
    auto reorderToPhysical = [&]() {
        std::vector<MatrixFlux> result;
        result.reserve(N);
        for (size_t j = 0; j < N; ++j)
            result.push_back(std::move(TR[top_exit ? j : N - 1 - j]));
        return result;
    };

    // Shortcut: if all kz equal and no B-field, return T=1, R=0 (pure Born).
    // Must check B=0, not just uniform B, since any B affects spin eigenmodes (#1333).
    const complex_t kz0 = kzs[0];
    const bool no_dwba =
        std::all_of(kzs.begin(), kzs.end(), [kz0](complex_t k) { return k == kz0; })
        && std::all_of(X.begin(), X.end(),
                       [&](size_t j) { return slices[j]->bField().mag() == 0; });

    if (no_dwba) {
        for (auto& tr : TR) {
            tr.m_T = SpinMatrix::One();
            tr.m_R = SpinMatrix::Null();
        }
        return reorderToPhysical();
    }

    // Wave along the surface, no penetration to the sample => total "reflection"
    if (kzs[X[0]] == 0.) {
        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 reorderToPhysical();
    }

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

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

    // Parratt recursion from exit (i=N-1) to entry (i=0), using logical indices
    for (size_t i = N - 1; i > 0; --i) {
        const size_t jthis = X[i - 1]; // physical index of layer i-1
        const size_t jlast = X[i];     // physical index of layer i
        const Slice& slice = *slices.at(jlast);
        const auto* roughness = slice.roughness();
        const double rms = slice.rms();
        const auto* r_model = roughness ? roughness->transient() : nullptr;

        // compute the 2x2 blocks of the transfer matrix
        // TR[i-1] and TR[i] are at logical indices, containing data for physical layers jthis,
        // jlast
        const auto [sp, sm] = refractionMatrixBlocks(TR[i - 1], TR[i], rms, r_model);
        const SpinMatrix delta = TR[i - 1].computeDeltaMatrix(slices[jthis]->thickness());

        // 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))) {
            for (size_t j = i; j < N; ++j) {
                TR[j].m_T = SpinMatrix::Null();
                TR[j].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 reorderToPhysical();
}

} // namespace

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

//! Returns fluxes for all sample slices.
//! See PhysRef, chapter "Polarized", section "Fluxes inside the sample".
Fluxes Compute::polarizedFluxes(const std::vector<std::unique_ptr<const Slice>>& slices,
                                const R3& k, bool forward)
{
    const bool top_exit = k.z() <= 0; // true if source or detector pixel are at z>=0
    std::vector<complex_t> kz = Compute::reducedKz(slices, k);
    ASSERT(slices.size() == kz.size());

    Fluxes result;
    for (const MatrixFlux& tr : ::computeTR(slices, kz, forward, top_exit))
        result.emplace_back(std::make_unique<MatrixFlux>(tr));

    markOnlyOutgoing(result, kz, slices, top_exit);
    return result;
}

//! Returns matrix r_0 reflected from the top layer.
//! See PhysRef, chapter "Polarized", section "Generalized Parratt recursion".
SpinMatrix Compute::polarizedReflectivity(const std::vector<std::unique_ptr<const Slice>>& 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) {
        const Slice& slice = *slices.at(i);
        const auto* roughness = slice.roughness();
        const double rms = slice.rms();
        MatrixFlux tr_i = createCoeff(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]->thickness());

        // 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;
}
