//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Sample/Interface/RoughnessMap.cpp
//! @brief     Implements RoughnessMap class.
//!
//! @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/RoughnessMap.h"
#include "Base/Util/Assert.h"
#include "Sample/Interface/Roughness.h"
#include "Sample/Multilayer/Layer.h"
#include <algorithm>
#include <numbers>

using std::numbers::pi;

namespace {

bool converged(const double2d_t& h_old, const double2d_t& h_new, double threshold)
{
    ASSERT(h_old.size() == h_new.size());
    ASSERT(h_old[0].size() == h_new[0].size());

    double sum_diff = 0;
    double sum_base = 0;
    int n = h_old.size() * h_old[0].size();

    for (size_t j = 0; j < h_old.size(); j++)
        for (size_t i = 0; i < h_old[0].size(); i++) {
            sum_diff += pow(h_old[j][i] - h_new[j][i], 2) / n;
            sum_base += pow(h_old[j][i], 2) / n;
        }

    return sum_diff < threshold * sum_base;
}

#ifdef BORNAGAIN_PYTHON

Arrayf64Wrapper arrayExport(const std::vector<std::size_t>& dimensions,
                            const std::vector<double>& flatData, const bool owndata)
{
    const std::size_t n_dims = dimensions.size();
    ASSERT(n_dims <= 2);

    return {flatData.size(), n_dims, dimensions.data(), flatData.data(), owndata};
}

#endif // BORNAGAIN_PYTHON

} // namespace

RoughnessMap::RoughnessMap(size_t x_points, size_t y_points, double Lx, double Ly,
                           const Sample& sample, int i_layer, int seed)
    : m_x_points(x_points)
    , m_y_points(y_points)
    , m_lx(Lx)
    , m_ly(Ly)
    , m_sample(sample)
    , m_i_layer(i_layer)
    , m_gen(seed < 0 ? m_rd() : seed)
{
    if (x_points == 0)
        throw std::runtime_error("Number of points along X must be >=1");
    if (y_points == 0)
        throw std::runtime_error("Number of points along Y must be >=1");
    if (Lx <= 0)
        throw std::runtime_error("Sample X length must be > 0");
    if (Ly <= 0)
        throw std::runtime_error("Sample Y length must be > 0");
}

double2d_t RoughnessMap::generateMap()
{
    createMap();
    return m_rough_map;
}

double2d_t RoughnessMap::mapFromHeights() const
{
    const size_t z_steps = 3999;
    const TransientModel* transient = m_sample.layer(m_i_layer)->roughness()->transient();
    const double rms = m_sample.roughnessRMS(m_i_layer);
    const double sigma_factor = transient->sigmaRange();
    const double z_limit = rms * sigma_factor;
    const double step = 2 * z_limit / (z_steps - 1);

    // create mesh of values
    std::vector<double> z_points(z_steps);
    for (size_t i = 0; i < z_steps; i++)
        z_points[i] = -z_limit + step * i;

    // fill mesh with weights
    std::vector<double> z_weights(z_steps);
    for (size_t i = 0; i < z_steps; i++)
        z_weights[i] = transient->distribution(z_points[i], rms);

    // fill map with random values
    std::discrete_distribution<int> d(z_weights.begin(), z_weights.end());
    double2d_t result(m_y_points, std::vector<double>(m_x_points));
    for (int j = 0; j < m_y_points; j++)
        for (int i = 0; i < m_x_points; i++)
            result[j][i] = z_points[d(m_gen)];

    return result;
}

double2d_t RoughnessMap::mapFromSpectrum() const
{
    const double fft_factor = m_x_points * m_y_points / std::sqrt(m_lx * m_ly);
    const double dfx = 1. / m_lx;
    const double dfy = 1. / m_ly;

    const int N = m_x_points / 2 + 1;
    const int M = m_y_points / 2 + 1;

    std::vector<double> fx(N);
    for (int i = 0; i < N; i++)
        fx[i] = i * dfx;

    std::vector<double> fy(M);
    for (int j = 0; j < M; j++)
        fy[j] = j * dfy;

    double2d_t psd_mag(m_y_points, std::vector<double>(N));
    for (int i = 0; i < N; i++) {
        for (int j = 0; j < M; j++)
            psd_mag[j][i] =
                std::sqrt(m_sample.roughnessSpectrum(std::hypot(fx[i], fy[j]), m_i_layer));
        for (int j = M; j < m_y_points; j++)
            psd_mag[j][i] = psd_mag[m_y_points - j][i];
    }
    psd_mag[0][0] = 0; // average height (amplitude at zero frequency) is null

    std::uniform_real_distribution<double> d(-pi, pi);
    double2d_t phase(m_y_points, std::vector<double>(N));

    // main axes
    // x axis
    for (int pos = 1; pos < (N - 1); pos++)
        phase[0][pos] = d(m_gen);

    // y axis
    for (int pos = 1; pos < (M - 1); pos++) {
        phase[pos][0] = d(m_gen);
        phase[m_y_points - pos][0] = -phase[pos][0];
    }

    // inner area
    for (int pos_j = 1; pos_j < m_y_points; pos_j++) {
        int sym_pos_j = m_y_points - pos_j;
        for (int pos_i = 1; pos_i < N; pos_i++) {
            int sym_pos_i = m_x_points - pos_i;
            if (pos_i < sym_pos_i || pos_j < sym_pos_j)
                phase[pos_j][pos_i] = d(m_gen);
            else
                phase[pos_j][pos_i] = +phase[sym_pos_j][pos_i];
        }
    }
    complex2d_t spectrum(m_y_points, std::vector<complex_t>(N));
    for (int j = 0; j < m_y_points; j++)
        for (int i = 0; i < N; i++)
            spectrum[j][i] = psd_mag[j][i] * std::exp(I * phase[j][i]) * fft_factor;

    return m_ft.irfft(spectrum, m_x_points);
}

double2d_t RoughnessMap::applySpectrumToHeights(const double2d_t& h_map,
                                                const double2d_t& s_map) const
{
    // get spectrum
    auto h_fft = m_ft.rfft(h_map);
    auto s_fft = m_ft.rfft(s_map);

    // rescale components
    for (size_t j = 0; j < h_fft.size(); j++)
        for (size_t i = 0; i < h_fft[0].size(); i++)
            if (std::abs(h_fft[j][i]) != 0)
                h_fft[j][i] *= std::abs(s_fft[j][i] / h_fft[j][i]);

    return m_ft.irfft(h_fft, h_map[0].size());
}

double2d_t RoughnessMap::applyHeightsToSpectrum(const double2d_t& h_map,
                                                const double2d_t& s_map) const
{
    // flatten spectral map
    auto s_map_flat = FieldUtil::flatten(s_map);

    // sort and remember the original positions
    std::vector<std::pair<double, size_t>> s_map_indexed(s_map_flat.size());
    for (size_t i = 0; i < s_map_flat.size(); i++)
        s_map_indexed[i] = std::make_pair(s_map_flat[i], i);

    std::sort(s_map_indexed.begin(), s_map_indexed.end());

    // replace heights of spectral map
    auto h_map_flat = FieldUtil::flatten(h_map);
    std::sort(h_map_flat.begin(), h_map_flat.end());

    for (size_t i = 0; i < s_map_flat.size(); i++)
        s_map_flat[s_map_indexed[i].second] = h_map_flat[i];

    return FieldUtil::reshapeTo2D(s_map_flat, s_map.size());
}

void RoughnessMap::createMap()
{
    if (m_sample.roughnessRMS(m_i_layer) < 1e-10) {
        m_rough_map = double2d_t(m_y_points, std::vector<double>(m_x_points));
        return;
    }

    double2d_t h_map = mapFromHeights();
    double2d_t s_map = mapFromSpectrum();

    for (int i = 0;; ++i) {
        // number of iterations is limited even if no convergence
        ASSERT(i < 100);
        double2d_t h_map_old = h_map;
        s_map = applySpectrumToHeights(h_map, s_map);
        h_map = applyHeightsToSpectrum(h_map, s_map);

        // adjust tolerance for proper speed/accuracy
        if (::converged(h_map_old, h_map, 1e-4))
            break;
    }

    // 's_map' has "perfect" original spectrum and tolerable height statistics.
    // 'h_map' has original height statistics but noisy spectrum.
    m_rough_map = s_map;
}

#ifdef BORNAGAIN_PYTHON

Arrayf64Wrapper RoughnessMap::generate()
{
    createMap();
    std::vector<std::size_t> dimensions = {m_rough_map.size(), m_rough_map[0].size()};
    return ::arrayExport(dimensions, FieldUtil::flatten(m_rough_map), /* owndata= */ true);
}

#endif // BORNAGAIN_PYTHON