//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Device/Resolution/ConvolutionDetectorResolution.cpp
//! @brief     Implements class ConvolutionDetectorResolution.
//!
//! @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 "Device/Resolution/ConvolutionDetectorResolution.h"
#include "Base/Axis/Frame.h"
#include "Base/Axis/Scale.h"
#include "Base/Util/Assert.h"
#include "Device/Data/Datafield.h"
#include "Device/Resolution/Convolve.h"

ConvolutionDetectorResolution::ConvolutionDetectorResolution(cumulative_DF_1d res_function_1d)
    : m_rank(1)
    , m_res_function_1d(res_function_1d)
{
}

ConvolutionDetectorResolution::ConvolutionDetectorResolution(
    const IResolutionFunction2D& res_function_2d)
    : m_rank(2)
    , m_res_function_1d(nullptr)
{
    setResolutionFunction(res_function_2d);
}

ConvolutionDetectorResolution::~ConvolutionDetectorResolution() = default;

ConvolutionDetectorResolution::ConvolutionDetectorResolution(
    const ConvolutionDetectorResolution& other)
    : m_rank(other.m_rank)
    , m_res_function_1d(other.m_res_function_1d)
{
    if (other.m_res_function_2d)
        setResolutionFunction(*other.m_res_function_2d);
}

ConvolutionDetectorResolution* ConvolutionDetectorResolution::clone() const
{
    return new ConvolutionDetectorResolution(*this);
}

std::vector<const INode*> ConvolutionDetectorResolution::nodeChildren() const
{
    return std::vector<const INode*>() << m_res_function_2d;
}

void ConvolutionDetectorResolution::execDetectorResolution(Datafield* df) const
{
    ASSERT(df->rank() == m_rank);
    if (m_rank == 1)
        apply1dConvolution(df);
    else if (m_rank == 2)
        apply2dConvolution(df);
    else
        ASSERT_NEVER;
}

void ConvolutionDetectorResolution::setResolutionFunction(const IResolutionFunction2D& resFunc)
{
    m_res_function_2d.reset(resFunc.clone());
}

void ConvolutionDetectorResolution::apply1dConvolution(Datafield* df) const
{
    ASSERT(m_res_function_1d);
    ASSERT(df->rank() == 1);

    const Scale& axis = df->axis(0);
    const size_t n = df->size();
    if (n < 2)
        return; // No convolution for sets of zero or one element

    // Construct kernel vector from resolution function
    ASSERT(axis.size() == n);
    double step_size = std::abs(axis.binCenter(0) - axis.binCenter(n - 1)) / (n - 1);
    double mid_value = axis.binCenter(n / 2); // because Convolve expects zero at midpoint
    std::vector<double> kernel;
    for (size_t index = 0; index < n; ++index)
        kernel.push_back(getIntegratedPDF1d(axis.binCenter(index) - mid_value, step_size));
    // Calculate convolution
    std::vector<double> result;
    Convolve().fftconvolve1D(df->flatVector(), kernel, result);
    // Truncate negative values that can arise because of finite precision of Fourier Transform
    for (double& e : result)
        e = std::max(0.0, e);

    df->setVector(result);
}

void ConvolutionDetectorResolution::apply2dConvolution(Datafield* df) const
{
    ASSERT(m_res_function_2d);
    ASSERT(df->rank() == 2);
    const Scale& X = df->axis(0);
    const Scale& Y = df->axis(1);
    const size_t nx = X.size();
    const size_t ny = Y.size();
    if (nx < 2 && ny < 2)
        return;

    // Construct kernel vector from resolution function
    std::vector<std::vector<double>> kernel;
    kernel.resize(ny);
    double mid_value1 = X.binCenter(nx / 2); // because Convolve expects zero at midpoint
    double mid_value2 = Y.binCenter(ny / 2); // because Convolve expects zero at midpoint
    double dx = std::abs(X.binCenter(0) - X.binCenter(nx - 1)) / (nx - 1);
    double dy = std::abs(Y.binCenter(0) - Y.binCenter(ny - 1)) / (ny - 1);
    for (size_t iy = 0; iy < ny; ++iy) {
        double y = Y.binCenter(iy) - mid_value2;
        std::vector<double> row_vector;
        row_vector.resize(nx, 0.0);
        for (size_t ix = 0; ix < nx; ++ix) {
            double x = X.binCenter(ix) - mid_value1;
            row_vector[ix] = getIntegratedPDF2d(x, dx, y, dy);
        }
        kernel[iy] = row_vector;
    }

    // Calculate convolution
    std::vector<std::vector<double>> result;
    Convolve().fftconvolve2D(df->values2D(), kernel, result);

    df->setVector2D(result);
}

double ConvolutionDetectorResolution::getIntegratedPDF1d(double x, double step) const
{
    double halfstep = step / 2.0;
    double xmin = x - halfstep;
    double xmax = x + halfstep;
    ASSERT(m_res_function_1d != nullptr);
    return m_res_function_1d(xmax) - m_res_function_1d(xmin);
}

double ConvolutionDetectorResolution::getIntegratedPDF2d(double x, double step_x, double y,
                                                         double step_y) const
{
    double halfstepx = step_x / 2.0;
    double halfstepy = step_y / 2.0;
    double xmin = x - halfstepx;
    double xmax = x + halfstepx;
    double ymin = y - halfstepy;
    double ymax = y + halfstepy;
    double result =
        m_res_function_2d->evaluateCDF(xmax, ymax) - m_res_function_2d->evaluateCDF(xmax, ymin)
        - m_res_function_2d->evaluateCDF(xmin, ymax) + m_res_function_2d->evaluateCDF(xmin, ymin);
    return result;
}
