//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Sim/Computation/RoughMultiLayerContribution.cpp
//! @brief     Implements class RoughMultiLayerContribution.
//!
//! @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/Computation/RoughMultiLayerContribution.h"
#include "3rdparty/Core/cerfcpp.h"
#include "Base/Const/PhysicalConstants.h"
#include "Base/Util/Assert.h"
#include "Resample/Element/DiffuseElement.h"
#include "Resample/Flux/ScalarFlux.h"
#include "Resample/Processed/ReSample.h"
#include "Sample/Interface/LayerRoughness.h"
#include "Sample/Multilayer/Layer.h"
#include "Sample/Multilayer/Sample.h"

using PhysConsts::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 h_above(complex_t z)
{
    return 0.5 * cerfcx(-mul_I(z) / std::sqrt(2.0));
}
complex_t h_below(complex_t z)
{
    return 0.5 * cerfcx(mul_I(z) / std::sqrt(2.0));
}

complex_t get_refractive_term(const ReSample& re_sample, size_t i_layer, double wavelength)
{
    return re_sample.avgeSlice(i_layer).material().refractiveIndex2(wavelength)
           - re_sample.avgeSlice(i_layer + 1).material().refractiveIndex2(wavelength);
}

complex_t get_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;

    const double thickness = re_sample.avgeSlice(i_layer).thicknessOr0();
    const complex_t T_i_A = i_A->getScalarT() * exp_I(kiz_A * thickness);
    const complex_t R_i_A = i_A->getScalarR() * exp_I(-kiz_A * thickness);
    const complex_t T_f_A = f_A->getScalarT() * exp_I(kfz_A * thickness);
    const complex_t R_f_A = f_A->getScalarR() * 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.averageSlices().bottomRMS(i_layer);
    const complex_t term1 = T_i_A * T_f_A * ::h_above(qz1_A * rms);
    const complex_t term2 = T_i_A * R_f_A * ::h_above(qz2_A * rms);
    const complex_t term3 = R_i_A * T_f_A * ::h_above(qz3_A * rms);
    const complex_t term4 = R_i_A * R_f_A * ::h_above(qz4_A * rms);
    const complex_t term5 = i_B->getScalarT() * f_B->getScalarT() * ::h_below(qz1_B * rms);
    const complex_t term6 = i_B->getScalarT() * f_B->getScalarR() * ::h_below(qz2_B * rms);
    const complex_t term7 = i_B->getScalarR() * f_B->getScalarT() * ::h_below(qz3_B * rms);
    const complex_t term8 = i_B->getScalarR() * f_B->getScalarR() * ::h_below(qz4_B * rms);

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

} // namespace

double Compute::roughMultiLayerContribution(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<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).topRoughness()) {
            roughStack.push_back(re_sample.avgeSlice(i + 1));
            rterm.push_back(::get_refractive_term(re_sample, i, wavelength));
            sterm.push_back(::get_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--) {
        auto* autocorr_model = roughStack[i].topRoughness()->autocorrelationModel();
        if (auto* k_corr = dynamic_cast<const K_CorrelationModel*>(autocorr_model)) {
            spectrum[i] = k_corr->spectralFunction(spatial_f);
        } else if (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();
        if (auto* spat_freq_cc = dynamic_cast<const SpatialFrequencyCrosscorrelation*>(
                roughStack[j].topRoughness()->crosscorrelationModel())) {
            crosscorr_with_interface_below[j] =
                spat_freq_cc->crosscorrSpectrum(spectrum[j], spectrum[j + 1], thickness, spatial_f);
        } else if (auto* lin_growth_autocorr = dynamic_cast<const LinearGrowthModel*>(
                       roughStack[j].topRoughness()->autocorrelationModel())) {
            crosscorr_with_interface_below[j] =
                lin_growth_autocorr->crosscorrSpectrum(spectrum[j + 1], thickness, 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);
}
