//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Sim/Evaluate/EvaluateRoughnessScattering.cpp
//! @brief     Implements function Evaluate::roughnessScattering.
//!
//! @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 "Sim/Evaluate/EvaluateRoughnessScattering.h"
#include "Base/Util/Assert.h"
#include "Resample/Sample/ReSample.h"
#include "Sample/Element/DiffuseElement.h"
#include "Sample/Flux/ScalarFlux.h"
#include "Sample/Interface/Roughness.h"
#include "Sample/Material/Material.h"
#include "Sample/Multilayer/Layer.h"
#include "Sample/Multilayer/RoughnessUtil.h"
#include "Sample/Multilayer/Slice.h"
#include "Sample/Sample/Sample.h"
#include <cerf.h>
#include <numbers>

using std::numbers::pi;

// As we say in our 2020 paper (Sect 5.6), diffuse scattering from rough interfaces
// is modelled after Schlomka et al, Phys Rev B, 51, 2311 (1995). They give credit
// for the basic modelling idea to Sinha et al (1988), and for the matrix elements
// to Holy et al (1993), Holy and Baumbach (1994), Sinha (1994) and de Boer (unpublished).
//
// Specifically, we implemented the differential cross section according to lines 2-3
// in Eq 3 of Schlomka et al. Unfortunately, this equation has an incorrect prefactor k^2,
// and so had our implementation up to release 20.0. The correct prefactor is k^4, as in
// Eq 17 of Holy et al (1993). This was corrected in release 20.1 (issue #553).

namespace {

complex_t hAbove(complex_t z)
{
    return 0.5 * cerfcx(-mul_I(z) / std::sqrt(2.0));
}
complex_t hBelow(complex_t z)
{
    return 0.5 * cerfcx(mul_I(z) / std::sqrt(2.0));
}

complex_t refractiveTerm(const ReSample& re_sample, size_t i_layer, double wavelength)
{
    return re_sample.avgeSlice(i_layer).refractingMaterial()->squaredRefractiveIndex(wavelength)
           - re_sample.avgeSlice(i_layer + 1)
                 .refractingMaterial()
                 ->squaredRefractiveIndex(wavelength);
}

complex_t sum8Terms(const ReSample& re_sample, size_t i_layer, const DiffuseElement& ele)
{
    // Abbreviations:
    //   i/f : initial/final beam
    //   A/B : above/below the interface
    const auto* i_A = dynamic_cast<const ScalarFlux*>(ele.fluxIn(i_layer));
    const auto* f_A = dynamic_cast<const ScalarFlux*>(ele.fluxOut(i_layer));
    const auto* i_B = dynamic_cast<const ScalarFlux*>(ele.fluxIn(i_layer + 1));
    const auto* f_B = dynamic_cast<const ScalarFlux*>(ele.fluxOut(i_layer + 1));
    if (!(i_A && f_A && i_B && f_B))
        throw std::runtime_error(
            "Rough interfaces not yet supported for polarized simulation (issue #564)");

    const complex_t kiz_A = i_A->getScalarKz();
    const complex_t kfz_A = f_A->getScalarKz();
    const complex_t qz1_A = -kiz_A - kfz_A;
    const complex_t qz2_A = -kiz_A + kfz_A;
    const complex_t qz3_A = -qz2_A;
    const complex_t qz4_A = -qz1_A;

    complex_t T_i_A = i_A->getScalarT();
    complex_t R_i_A = i_A->getScalarR();
    complex_t T_f_A = f_A->getScalarT();
    complex_t R_f_A = f_A->getScalarR();

    if (const double thickness = re_sample.avgeSlice(i_layer).thickness();
        std::isfinite(thickness)) {
        T_i_A *= exp_I(kiz_A * thickness);
        R_i_A *= exp_I(-kiz_A * thickness);
        T_f_A *= exp_I(kfz_A * thickness);
        R_f_A *= exp_I(-kfz_A * thickness);
    }

    const complex_t kiz_B = i_B->getScalarKz();
    const complex_t kfz_B = f_B->getScalarKz();
    const complex_t qz1_B = -kiz_B - kfz_B;
    const complex_t qz2_B = -kiz_B + kfz_B;
    const complex_t qz3_B = -qz2_B;
    const complex_t qz4_B = -qz1_B;

    const double rms = re_sample.slices().at(i_layer + 1)->rms();
    const complex_t term1 = T_i_A * T_f_A * ::hAbove(qz1_A * rms);
    const complex_t term2 = T_i_A * R_f_A * ::hAbove(qz2_A * rms);
    const complex_t term3 = R_i_A * T_f_A * ::hAbove(qz3_A * rms);
    const complex_t term4 = R_i_A * R_f_A * ::hAbove(qz4_A * rms);
    const complex_t term5 = i_B->getScalarT() * f_B->getScalarT() * ::hBelow(qz1_B * rms);
    const complex_t term6 = i_B->getScalarT() * f_B->getScalarR() * ::hBelow(qz2_B * rms);
    const complex_t term7 = i_B->getScalarR() * f_B->getScalarT() * ::hBelow(qz3_B * rms);
    const complex_t term8 = i_B->getScalarR() * f_B->getScalarR() * ::hBelow(qz4_B * rms);

    return term1 + term2 + term3 + term4 + term5 + term6 + term7 + term8;
}

} // namespace


double Evaluate::roughnessScattering(const ReSample& re_sample, const DiffuseElement& ele)
{
    if (ele.alphaMean() < 0.0)
        return 0;

    const size_t n_slices = re_sample.numberOfSlices();
    const double spatial_f = ele.meanQ().magxy() / (2 * pi);
    const double wavelength = ele.wavelength();
    double autocorr_sum = 0;
    double crosscorr_sum = 0;

    std::vector<const Slice*> roughStack; // only slices with rough top interfaces
    std::vector<complex_t> rterm;
    std::vector<complex_t> sterm;

    roughStack.reserve(n_slices - 1);
    rterm.reserve(n_slices - 1);
    sterm.reserve(n_slices - 1);

    for (size_t i = 0; i < n_slices - 1; i++)
        if (re_sample.avgeSlice(i + 1).roughness()) {
            roughStack.emplace_back(&re_sample.avgeSlice(i + 1));
            rterm.push_back(::refractiveTerm(re_sample, i, wavelength));
            sterm.push_back(::sum8Terms(re_sample, i, ele));
        }

    const size_t n_interfaces = roughStack.size();
    std::vector<double> spectrum(n_interfaces, 0);
    std::vector<double> crosscorr_with_interface_below(n_interfaces, 0); // the last is zero

    // precompute autocorrelation terms
    for (int i = n_interfaces - 1; i >= 0; i--) {
        const auto* autocorr_model = roughStack[i]->roughness()->autocorrelationModel();
        if (const auto* k_corr = dynamic_cast<const SelfAffineFractalModel*>(autocorr_model)) {
            spectrum[i] = k_corr->spectralFunction(spatial_f);
        } else if (const auto* lin_growth =
                       dynamic_cast<const LinearGrowthModel*>(autocorr_model)) {
            ASSERT(i < int(n_interfaces - 1));
            const double thickness = roughStack[i]->hig() - roughStack[i + 1]->hig();
            spectrum[i] = lin_growth->spectralFunction(spectrum[i + 1], thickness, spatial_f);
        }
    }

    // auto correlation in each layer (first term in final expression in Eq (3) of Schlomka et al)
    for (size_t i = 0; i < n_interfaces; i++)
        autocorr_sum += std::norm(rterm[i] * sterm[i]) * spectrum[i];

    // precompute crosscorrelation terms
    for (int j = n_interfaces - 2; j >= 0; j--) {
        const double thickness = roughStack[j]->hig() - roughStack[j + 1]->hig();
        crosscorr_with_interface_below[j] = RoughnessUtil::crossSpectrumWithInterfaceBelow(
            *roughStack[j]->roughness(), thickness, spectrum[j], spectrum[j + 1], spatial_f);
    }

    // cross correlation between layers (second term in loc. cit.)
    for (size_t j = 0; j < n_interfaces - 1; j++) {
        if (crosscorr_with_interface_below[j] == 0)
            continue;
        double crosscorr_spectrum = crosscorr_with_interface_below[j];
        for (size_t k = j + 1; k < n_interfaces; k++) {
            crosscorr_sum +=
                (rterm[j] * sterm[j] * std::conj(rterm[k] * sterm[k])).real() * crosscorr_spectrum;
            if (crosscorr_with_interface_below[k] == 0 || spectrum[k] == 0)
                break;
            crosscorr_spectrum *= crosscorr_with_interface_below[k] / spectrum[k];
        }
    }
    const double k0 = (2 * pi) / wavelength;
    return pow(k0, 4) / 16 / pi / pi * (autocorr_sum + 2 * crosscorr_sum);
}
