//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Sample/Interface/AutocorrelationModels.cpp
//! @brief     Implement AutocorrelationModel classes.
//!
//! @homepage  http://www.bornagainproject.org
//! @license   GNU General Public License v3 or higher (see COPYING)
//! @copyright Forschungszentrum Jülich GmbH 2024
//! @authors   Scientific Computing Group at MLZ (see CITATION, AUTHORS)
//
//  ************************************************************************************************

#include "Sample/Interface/AutocorrelationModels.h"
#include "Base/Const/PhysicalConstants.h"
#include "Base/Py/PyFmt.h"
#include "Base/Util/Assert.h"

using PhysConsts::pi;

AutocorrelationModel::AutocorrelationModel(double maxSpatFrequency)
    : m_max_spatial_frequency(maxSpatFrequency)
{
}

std::vector<std::string> AutocorrelationModel::validationErrs() const
{
    std::vector<std::string> errs;
    requestGt0(errs, m_max_spatial_frequency, "maxSpatialFrequency");
    return errs;
}

std::string AutocorrelationModel::pythonArguments() const
{
    return Py::Fmt::printArguments(
        {{m_max_spatial_frequency, AutocorrelationModel::parDefs()[0].unit}});
}

//-------------------------------------------------------------------------------------------------

//! @param sigma: height scale of the roughness in nanometers (usually close to rms)
//! @param hurstParameter: hurst parameter which describes how jagged the interface,
//! dimensionless (0.0, 1.0], where 0.0 gives more spikes, 1.0 more smoothness
//! @param lateralCorrLength: lateral correlation length of the roughness in nanometers
K_CorrelationModel::K_CorrelationModel(double sigma, double hurst, double lateralCorrLength,
                                       double maxSpatFrequency)
    : AutocorrelationModel(maxSpatFrequency)
    , m_sigma(sigma)
    , m_hurst_parameter(hurst)
    , m_lateral_corr_length(lateralCorrLength)
{
    validateOrThrow();
}

K_CorrelationModel* K_CorrelationModel::clone() const
{
    return new K_CorrelationModel(m_sigma, m_hurst_parameter, m_lateral_corr_length,
                                  m_max_spatial_frequency);
}

std::vector<ParaMeta> K_CorrelationModel::parDefs() const
{
    std::vector<ParaMeta> result = {{"Sigma", "nm"}, {"Hurst", ""}, {"CorrLength", "nm"}};
    const auto base_pars = AutocorrelationModel::parDefs();
    result.insert(result.end(), base_pars.begin(), base_pars.end());
    return result;
}

std::string K_CorrelationModel::validate() const
{
    std::vector<std::string> errs = AutocorrelationModel::validationErrs();
    requestGe0(errs, m_sigma, "sigma");
    if (m_sigma > 0) {
        requestGt0(errs, m_hurst_parameter, "hurst");
        requestIn(errs, m_hurst_parameter, "hurst", 0, 1);
        requestGt0(errs, m_lateral_corr_length, "lateralCorrLength");
    }
    if (!errs.empty())
        return jointError(errs);
    m_validated = true;
    return "";
}

std::string K_CorrelationModel::pythonArguments() const
{
    return Py::Fmt::printArguments({{m_sigma, parDefs()[0].unit},
                                    {m_hurst_parameter, parDefs()[1].unit},
                                    {m_lateral_corr_length, parDefs()[2].unit}})
           + ", " + AutocorrelationModel::pythonArguments();
}

//! Power spectral density of the surface roughness is a result of two-dimensional
//! Fourier transform of the correlation function of the roughness profile.
//!
//! Based on Palasantzas, Phys Rev B, 48, 14472 (1993)
double K_CorrelationModel::spectralFunction(double spatial_f) const
{
    ASSERT(m_validated);
    if (spatial_f > m_max_spatial_frequency)
        return 0;

    const double Qpar2 = pow(2 * pi * spatial_f, 2);
    const double H = m_hurst_parameter;
    const double clength2 = m_lateral_corr_length * m_lateral_corr_length;
    return 4.0 * pi * H * m_sigma * m_sigma * clength2 * std::pow(1 + Qpar2 * clength2, -1 - H);
}

double K_CorrelationModel::rms() const
{
    // integration of spectral function: rms^2 = Integrate[PSD(f) * 2*pi*f, {f, 0, max_frequency}]
    const double H = m_hurst_parameter;
    const double val = 2 * pi * m_lateral_corr_length * m_max_spatial_frequency;
    return m_sigma * std::sqrt(1. - std::pow(1 + val * val, -H));
}

//-------------------------------------------------------------------------------------------------

LinearGrowthModel::LinearGrowthModel(double particle_volume, double damp1, double damp2,
                                     double damp3, double damp4, double maxSpatFrequency)
    : AutocorrelationModel(maxSpatFrequency)
    , m_cluster_volume(particle_volume)
    , m_damp1(damp1)
    , m_damp2(damp2)
    , m_damp3(damp3)
    , m_damp4(damp4)
{
    validateOrThrow();
}

LinearGrowthModel* LinearGrowthModel::clone() const
{
    return new LinearGrowthModel(m_cluster_volume, m_damp1, m_damp2, m_damp3, m_damp4,
                                 m_max_spatial_frequency);
}

std::string LinearGrowthModel::validate() const
{
    std::vector<std::string> errs = AutocorrelationModel::validationErrs();
    requestGe0(errs, m_cluster_volume, "particleVolume");
    requestGe0(errs, m_damp1, "dampingExp1");
    requestGe0(errs, m_damp2, "dampingExp2");
    requestGe0(errs, m_damp3, "dampingExp3");
    requestGe0(errs, m_damp4, "dampingExp4");
    if (!errs.empty())
        return jointError(errs);
    m_validated = true;
    return "";
}

std::vector<ParaMeta> LinearGrowthModel::parDefs() const
{
    std::vector<ParaMeta> result = {{"Particle volume", "nm^3"},
                                    {"Damping 1", ""},
                                    {"Damping 2", "nm"},
                                    {"Damping 3", "nm^2"},
                                    {"Damping 4", "nm^3"}};

    const auto base_pars = AutocorrelationModel::parDefs();
    result.insert(result.end(), base_pars.begin(), base_pars.end());
    return result;
}

std::string LinearGrowthModel::pythonArguments() const
{
    return Py::Fmt::printArguments({{m_cluster_volume, parDefs()[0].unit},
                                    {m_damp1, parDefs()[1].unit},
                                    {m_damp2, parDefs()[2].unit},
                                    {m_damp3, parDefs()[3].unit},
                                    {m_damp4, parDefs()[4].unit}})
           + ", " + AutocorrelationModel::pythonArguments();
    ;
}

double LinearGrowthModel::damping(double spatial_f) const
{
    return m_damp1 * spatial_f + m_damp2 * std::pow(spatial_f, 2) + m_damp3 * std::pow(spatial_f, 3)
           + m_damp4 * std::pow(spatial_f, 4);
}

double LinearGrowthModel::spectralFunction(double spectrum_below, double thickness,
                                           double spatial_f) const
{
    ASSERT(m_validated);
    if (spatial_f > m_max_spatial_frequency)
        return 0;

    if (thickness == 0)
        return spectrum_below;

    double inherited_spectrum = spectrum_below;
    double growth_spectrum = m_cluster_volume * thickness;

    const double damp = damping(spatial_f);
    if (damp > 0) {
        const double exponent = std::exp(-damp * thickness);
        inherited_spectrum = spectrum_below * exponent;
        growth_spectrum = m_cluster_volume * (1. - exponent) / damp;
    }
    return inherited_spectrum + growth_spectrum;
}

double LinearGrowthModel::crosscorrSpectrum(double spectrum_below, double thickness,
                                            double spatial_f) const
{
    ASSERT(m_validated);
    const double damp = damping(spatial_f);
    if (damp == 0 || thickness == 0)
        return spectrum_below;

    return spectrum_below * std::exp(-damp * thickness);
}
