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

using std::numbers::pi;

namespace {

// Relative mean-square map change used by the IAAFT loop.
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

// Wraps a flat C++ vector as an Arrayf64Wrapper with the requested NumPy-style dimensions.
Arrayf64Wrapper arrayExport(const std::vector<int>& dimensions, const std::vector<double>& flatData,
                            const bool copydata)
{
    const std::size_t rank = dimensions.size();
    ASSERT(rank <= 3);

    return {rank, dimensions.data(), &flatData, copydata};
}

#endif // BORNAGAIN_PYTHON

// Preserve the Fourier phase of a coefficient. A zero coefficient carries no phase information, so
// use a deterministic neutral phase; this case is rare and only affects degenerate bins.
complex_t unitPhase(complex_t value)
{
    const double magnitude = std::abs(value);
    return magnitude > 0 ? value / magnitude : complex_t(1.0, 0.0);
}

bool isSelfConjugateBin(size_t ix, size_t iy, size_t x_points, size_t y_points)
{
    const bool self_x = ix == 0 || (x_points % 2 == 0 && ix == x_points / 2);
    const bool self_y = iy == 0 || (y_points % 2 == 0 && iy == y_points / 2);
    return self_x && self_y;
}

// Returns a unit phase orthogonal to reference in the complex plane. The sign is taken from
// phase_hint, so conjugate rFFT bins remain conjugate and the independent part is not redrawn. For
// self-conjugate DC/Nyquist bins the coefficient is real, so a non-zero orthogonal residual cannot
// be represented; those bins fall back to the phase hint.
complex_t orthogonalUnitPhase(complex_t phase_hint, complex_t reference, size_t ix, size_t iy,
                              size_t x_points, size_t y_points)
{
    if (::isSelfConjugateBin(ix, iy, x_points, y_points))
        return ::unitPhase(phase_hint);

    const double reference_magnitude = std::abs(reference);
    if (reference_magnitude <= 0)
        return ::unitPhase(phase_hint);

    const complex_t reference_phase = reference / reference_magnitude;
    const complex_t orthogonal_phase(-reference_phase.imag(), reference_phase.real());
    const double side = std::imag(phase_hint * std::conj(reference_phase));
    return side < 0 ? -orthogonal_phase : orthogonal_phase;
}

// rFFT stores only non-negative x frequencies; y frequencies still wrap around the Nyquist row.
double spatialFrequency(size_t ix, size_t iy, size_t y_points, double Lx, double Ly)
{
    const size_t mirrored_y = iy <= y_points / 2 ? iy : y_points - iy;
    return std::hypot(static_cast<double>(ix) / Lx, static_cast<double>(mirrored_y) / Ly);
}

// Spectral projection for an interface with a rough interface below. The projected spectrum is the
// correlated inherited component C/S_below * H_below plus an independent residual whose power is
// S_current - C^2/S_below. For each non-self-conjugate Fourier bin the residual is placed
// orthogonally to H_below, so the concrete map, not only the ensemble average, satisfies both the
// target auto-spectrum and the real adjacent cross-spectrum, except for the self-conjugate bins
// where the residual direction is necessarily real.
complex2d_t projectUpperSpectrum(const complex2d_t& independent_spectrum_hint,
                                 const complex2d_t& below_spectrum, size_t x_points,
                                 size_t y_points, double Lx, double Ly, double fft_factor,
                                 size_t i_layer, size_t i_below,
                                 const std::vector<std::unique_ptr<Layer>>& layers)
{
    const size_t spectrum_x_points = x_points / 2 + 1;
    complex2d_t projected(y_points, std::vector<complex_t>(spectrum_x_points));

    for (size_t j = 0; j < y_points; ++j) {
        for (size_t i = 0; i < spectrum_x_points; ++i) {
            // Read the target auto-spectra for this Fourier bin.
            const double spatial_f = ::spatialFrequency(i, j, y_points, Lx, Ly);
            const double spectrum_current =
                spatial_f == 0 ? 0 : RoughnessUtil::spectrum(layers, spatial_f, i_layer);
            const double spectrum_below =
                spatial_f == 0 ? 0 : RoughnessUtil::spectrum(layers, spatial_f, i_below);
            if (spectrum_current <= 0 || spectrum_below <= 0) {
                projected[j][i] = complex_t(0.0, 0.0);
                continue;
            }

            // Split the upper auto-spectrum into the part inherited from the lower interface and
            // the remaining independent power prescribed by the adjacent cross-spectrum.
            const double thickness = layers[i_layer]->hig() - layers[i_below]->hig();
            const double cross_spectrum_raw = RoughnessUtil::crossSpectrumWithInterfaceBelow(
                *layers[i_layer]->roughness(), thickness, spectrum_current, spectrum_below,
                spatial_f);
            const double cross_spectrum =
                std::clamp(cross_spectrum_raw, 0.0, std::sqrt(spectrum_current * spectrum_below));
            const double inherited_factor = cross_spectrum / spectrum_below;
            const double independent_spectrum =
                std::max(0.0, spectrum_current - cross_spectrum * cross_spectrum / spectrum_below);

            // Take the random sign of the independent phase from the existing spectral generator,
            // but force it to be orthogonal to the lower-interface phase.
            const complex_t independent_phase = ::orthogonalUnitPhase(
                independent_spectrum_hint[j][i], below_spectrum[j][i], i, j, x_points, y_points);
            projected[j][i] = below_spectrum[j][i] * inherited_factor
                              + std::sqrt(independent_spectrum) * fft_factor * independent_phase;
        }
    }
    return projected;
}

} // namespace


// Stores map dimensions, physical patch size, selected interface, and deterministic RNG state.
RoughnessMap::RoughnessMap(size_t x_points, size_t y_points, double Lx, double Ly,
                           const Sample& sample, size_t i_layer, unsigned 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)
{
    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");
}

// Generates all rough interface maps in one bottom-up pass. The lowest rough interface is generated
// by the original single-interface algorithm; each next interface is built from the one below plus
// an orthogonal independent residual with the prescribed remaining spectrum.
std::vector<double2d_t> RoughnessMap::generateInterfaceMaps(size_t x_points, size_t y_points,
                                                            double Lx, double Ly,
                                                            const Sample& sample, unsigned seed)
{
    const auto& layers = sample.unwrapped();
    std::vector<double2d_t> result(layers.size());

    // Work only on finite rough interfaces. Layer 0 is the ambient semi-infinite layer.
    std::vector<size_t> rough_layers;
    rough_layers.reserve(layers.size());
    for (size_t i = 1; i < layers.size(); ++i)
        if (layers[i]->roughness() && layers[i]->rms() >= MIN_RMS)
            rough_layers.push_back(i);
    if (rough_layers.empty())
        return result;

    const double fft_factor =
        static_cast<double>(x_points) * static_cast<double>(y_points) / std::sqrt(Lx * Ly);

    std::mt19937 rng(seed);
    FourierTransform ft;

    complex2d_t below_spectrum;
    size_t i_below = 0;
    for (size_t pos = rough_layers.size(); pos-- > 0;) {
        const size_t i_layer = rough_layers[pos];

        // The lowest rough interface has no generated lower reference, so keep the original
        // single-interface IAAFT algorithm unchanged and use the public seed directly.
        if (below_spectrum.empty()) {
            result[i_layer] =
                RoughnessMap(x_points, y_points, Lx, Ly, sample, i_layer, seed).generateMap();
            below_spectrum = ft.rfft(result[i_layer]);
            i_below = i_layer;
            continue;
        }

        // For upper interfaces, reuse the existing spectral random-map generator only to obtain a
        // random phase hint for the independent residual.
        const double2d_t independent_map =
            RoughnessMap(x_points, y_points, Lx, Ly, sample, i_layer, rng()).mapFromSpectrum();
        const complex2d_t independent_hint = ft.rfft(independent_map);
        complex2d_t projected =
            ::projectUpperSpectrum(independent_hint, below_spectrum, x_points, y_points, Lx, Ly,
                                   fft_factor, i_layer, i_below, layers);

        // The generated upper interface becomes the reference for the next interface above.
        result[i_layer] = ft.irfft(projected, static_cast<int>(x_points));
        below_spectrum = std::move(projected);
        i_below = i_layer;
    }

    return result;
}

#ifdef BORNAGAIN_PYTHON

// Python-facing stack generator: selects requested interface maps and returns one rank-3 array with
// shape (number of requested interfaces, y_points, x_points).
Arrayf64Wrapper RoughnessMap::generateInterfaceMaps(size_t x_points, size_t y_points, double Lx,
                                                    double Ly, const Sample& sample,
                                                    const std::vector<int>& interface_indices,
                                                    unsigned seed)
{
    const std::vector<double2d_t> maps =
        generateInterfaceMaps(x_points, y_points, Lx, Ly, sample, seed);

    std::vector<double> flat;
    flat.reserve(interface_indices.size() * x_points * y_points);
    for (int i_layer : interface_indices) {
        if (i_layer < 0 || static_cast<size_t>(i_layer) >= maps.size())
            throw std::runtime_error("RoughnessMap::generateInterfaceMaps: interface index "
                                     + std::to_string(i_layer) + " is out of range");

        const double2d_t& map = maps[static_cast<size_t>(i_layer)];
        if (map.empty())
            throw std::runtime_error("RoughnessMap::generateInterfaceMaps: interface index "
                                     + std::to_string(i_layer) + " has no generated roughness map");

        std::vector<double> map_flat = FieldUtil::flatten(map);
        flat.insert(flat.end(), map_flat.begin(), map_flat.end());
    }

    const std::vector<int> dimensions = {int(interface_indices.size()), int(y_points),
                                         int(x_points)};
    return ::arrayExport(dimensions, flat, /* copydata= */ true);
}

#endif // BORNAGAIN_PYTHON

// Generates one interface map with the original single-interface IAAFT algorithm.
double2d_t RoughnessMap::generateMap()
{
    if (m_sample.unwrappedLayer(m_i_layer)->rms() < MIN_RMS)
        return double2d_t(m_y_points, std::vector<double>(m_x_points));

    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.
    return s_map;
}

#ifdef BORNAGAIN_PYTHON

// Python-facing single-interface generator: returns the generated map as a rank-2 array.
Arrayf64Wrapper RoughnessMap::generate()
{
    double2d_t map = generateMap();
    std::vector<int> dimensions = {int(map.size()), int(map[0].size())};
    return ::arrayExport(dimensions, FieldUtil::flatten(map), /* copydata= */ true);
}

#endif // BORNAGAIN_PYTHON

// Builds a random height field sampled from the selected interface transient distribution.
double2d_t RoughnessMap::mapFromHeights() const
{
    const size_t z_steps = 3999;
    const TransientModel* transient = m_sample.unwrappedLayer(m_i_layer)->roughness()->transient();
    const double rms = m_sample.unwrappedLayer(m_i_layer)->rms();
    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;
}

// Builds a random height field with the selected interface auto-spectrum and random phases.
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(
                RoughnessUtil::spectrum(m_sample.unwrapped(), 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);
}

// Spectral IAAFT projection: keeps phases from h_map and replaces amplitudes by those from s_map.
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());
}

// Height IAAFT projection: keeps the rank order of s_map and replaces its sorted values by those of
// h_map.
double2d_t RoughnessMap::applyHeightsToSpectrum(const double2d_t& h_map, const double2d_t& s_map)
{
    // 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());
}
