//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Resample/Specular/ComputeFluxScalar.cpp
//! @brief     Implements functions to compute scalar fluxes.
//!
//! @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/Specular/ComputeFluxScalar.h"
#include "Base/Math/Functions.h"
#include "Base/Util/Assert.h"
#include "Resample/Flux/ScalarFlux.h"
#include "Resample/Slice/KzComputation.h"
#include "Resample/Slice/SliceStack.h"
#include "Sample/Interface/LayerRoughness.h"
#include "Sample/Multilayer/Layer.h"
#include <numbers>

using std::numbers::pi;

namespace {

//! See PhysRef, chapter "Scattering by rough interfaces", section "Interface with tanh profile".
std::pair<complex_t, complex_t> transition(complex_t kzi, complex_t kzi1, double sigma,
                                           const RoughnessModel* r_model)
{
    const complex_t kz_ratio = kzi1 / kzi;

    if (sigma == 0)
        return {1. + kz_ratio, 1. - kz_ratio};

    ASSERT(sigma > 0);
    ASSERT(r_model);

    if (dynamic_cast<const ErfRoughness*>(r_model))
        // Roughness is modelled by a Gaussian profile, i.e. Nevot-Croce factors for the
        // reflection coefficients.
        // Implementation follows A. Gibaud and G. Vignaud, in X-ray and Neutron Reflectivity,
        // edited by J. Daillant and A. Gibaud, volume 770 of Lecture Notes in Physics (2009)
        return {(1. + kz_ratio) * std::exp(-pow((kzi1 - kzi) * sigma, 2) / 2.),
                (1. - kz_ratio) * std::exp(-pow((kzi1 + kzi) * sigma, 2) / 2.)};

    // TANH model assumed
    // Roughness is modelled by tanh profile
    // [e.g. Bahr, Press, et al, Phys. Rev. B, vol. 47 (8), p. 4385 (1993)].
    // but the scale factor is adjusted to give rms == sigma.
    const double sigeff = std::sqrt(3.0) * sigma;
    const complex_t roughness = std::sqrt(Math::tanhc(sigeff * kzi1) / Math::tanhc(sigeff * kzi));

    return {1. / roughness + kz_ratio * roughness, 1. / roughness - kz_ratio * roughness};
}

std::vector<Spinor> computeTR(const SliceStack& slices, const std::vector<complex_t>& kz,
                              bool top_exit)
{
    const size_t N = slices.size();
    std::vector<Spinor> TR(N, {1., 0.});

    if (N == 1) // If only one layer present, there's nothing left to calculate
        return TR;

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

    if (kz[X[0]] == 0.0) { // If kz in layer 0 is zero, R0 = -T0 and all others equal to 0
        TR[X[0]] = {1.0, -1.0};
        for (size_t i = 1; i < N; ++i)
            TR[X[i]] = Spinor(0, 0);
        return TR;
    }

    // Parratt algorithm, pass 1: compute t/t factors and r/t ratios from bottom to top.
    std::vector<complex_t> tfactor(N - 1); // transmission damping t_{i+1} / t_{i}
    std::vector<complex_t> ratio(N);       // Parratt's x = r/t
    for (size_t i = N - 1; i > 0; i--) {
        const size_t jthis = X[i - 1];
        const size_t jlast = X[i];
        const auto* roughness = slices.bottomRoughness(jthis);
        const double sigma = roughness ? roughness->sigma() : 0.;
        const auto r_model = roughness ? roughness->roughnessModel() : nullptr;

        const auto [slp, slm] = transition(kz[jthis], kz[jlast], sigma, r_model);

        const complex_t delta = exp_I(kz[jthis] * slices[jthis].thicknessOr0());
        const complex_t f = delta / (slp + slm * ratio[jlast]);
        tfactor[i - 1] = 2. * f;
        ratio[jthis] = delta * (slm + slp * ratio[jlast]) * f;
    }

    // Parratt algorithm, pass 2: compute r and t from top to bottom.
    TR[X[0]] = Spinor(1., ratio[X[0]]);
    for (size_t i = 1; i < N; ++i) {
        TR[X[i]].u = TR[X[i - 1]].u * tfactor[i - 1]; // Spinor.u is t
        TR[X[i]].v = ratio[X[i]] * TR[X[i]].u;        // Spinor.v is r
    }

    return TR;
}

} // namespace

//  ************************************************************************************************
//  implementation of public interface
//  ************************************************************************************************

Fluxes Compute::scalarFluxes(const SliceStack& slices, const R3& k)
{
    const bool top_exit = k.z() <= 0; // true if source or detector pixel are at z>=0
    std::vector<complex_t> kz = Compute::Kz::computeReducedKz(slices, k);
    ASSERT(slices.size() == kz.size());

    const std::vector<Spinor> TR = ::computeTR(slices, kz, top_exit);

    Fluxes result;
    for (size_t i = 0; i < kz.size(); ++i)
        result.push_back(new ScalarFlux(kz[i], TR[i]));
    return result;
}

complex_t Compute::scalarReflectivity(const SliceStack& slices, const std::vector<complex_t>& kz)
{
    ASSERT(slices.size() == kz.size());
    const size_t N = slices.size();
    if (N == 1)
        return 0.; // only one layer present, there's nothing left to calculate
    if (kz[0] == 0.)
        return -1.;

    complex_t R_i1 = 0.;

    for (size_t i = N - 1; i > 0; i--) {
        const auto* roughness = slices.bottomRoughness(i - 1);
        const double sigma = roughness ? roughness->sigma() : 0.;
        const auto r_model = roughness ? roughness->roughnessModel() : nullptr;

        const auto [sp, sm] = ::transition(kz[i - 1], kz[i], sigma, r_model);

        const complex_t delta = exp_I(kz[i - 1] * slices[i - 1].thicknessOr0());

        R_i1 = pow(delta, 2) * (sm + sp * R_i1) / (sp + sm * R_i1);
    }

    return R_i1;
}
