//  ************************************************************************************************
//
//  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/Slice/ComputeKz.h"
#include "Sample/Flux/ScalarFlux.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 {

//! 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 rms,
                                           const TransientModel* r_model)
{
    const complex_t kz_ratio = kzi1 / kzi;

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

    ASSERT(rms > 0);
    ASSERT(r_model);

    if (dynamic_cast<const ErfTransient*>(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) * rms, 2) / 2.),
                (1. - kz_ratio) * std::exp(-pow((kzi1 + kzi) * rms, 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) * rms;
    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 std::vector<std::unique_ptr<const Slice>>& slices,
                              const std::vector<complex_t>& kz, bool top_exit)
{
    const size_t N = slices.size();
    std::vector<Spinor> TR(N, {1., 0.});

    // If all layers have the same kz (same SLD), there are no reflecting interfaces.
    // Return simple T=1, R=0 for all layers (pure Born approximation case).
    const complex_t kz0 = kz[0];
    if (std::all_of(kz.begin(), kz.end(), [kz0](complex_t k) { return k == kz0; }))
        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 Slice& slice = *slices.at(jlast);
        const auto* roughness = slice.roughness();
        const double rms = slice.rms();
        const auto* r_model = roughness ? roughness->transient() : nullptr;

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

        if (const double thickness = slices[jthis]->thickness(); std::isfinite(thickness)) {
            const complex_t delta = exp_I(kz[jthis] * thickness);
            const complex_t f = delta / (slp + slm * ratio[jlast]);
            tfactor[i - 1] = 2. * f;
            ratio[jthis] = delta * (slm + slp * ratio[jlast]) * f;
        } else {
            const complex_t f = 1. / (slp + slm * ratio[jlast]);
            tfactor[i - 1] = 2. * f;
            ratio[jthis] = (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 std::vector<std::unique_ptr<const Slice>>& 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::reducedKz(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.emplace_back(std::make_unique<ScalarFlux>(kz[i], TR[i]));

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

complex_t Compute::scalarReflectivity(const std::vector<std::unique_ptr<const Slice>>& 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 Slice& slice = *slices.at(i);
        const auto* roughness = slice.roughness();
        const double rms = slice.rms();
        const auto* r_model = roughness ? roughness->transient() : nullptr;

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

        if (const double thickness = slices[i - 1]->thickness(); std::isfinite(thickness)) {
            const complex_t delta = exp_I(kz[i - 1] * thickness);
            R_i1 = pow(delta, 2) * (sm + sp * R_i1) / (sp + sm * R_i1);
        } else
            R_i1 = (sm + sp * R_i1) / (sp + sm * R_i1);
    }

    return R_i1;
}
