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

Datafield::Datafield() = default;

Datafield::Datafield(const Frame* frame, const std::vector<double>& values,
                     const std::vector<double>& errSigmas)
    : m_frame(frame)
    , m_values(values.empty() ? std::vector<double>(frame->size(), 0.) : values)
    , m_errSigmas(errSigmas)
{
    ASSERT(m_frame);
    ASSERT(m_values.size() == m_frame->size());
    ASSERT(m_errSigmas.empty() || m_errSigmas.size() == m_values.size());
}

Datafield::Datafield(std::vector<const Scale*>&& axes, const std::vector<double>& values,
                     const std::vector<double>& errSigmas)
    : Datafield(new Frame(std::move(axes)), values, errSigmas)
{
}

Datafield::Datafield(Datafield&&) = default;

Datafield::Datafield(const Datafield& other)
    : Datafield(other.m_frame ? other.m_frame->clone() : nullptr, other.m_values, other.m_errSigmas)
{
}

Datafield::~Datafield() = default;

Datafield& Datafield::operator=(Datafield&& other) noexcept = default;

Datafield* Datafield::clone() const
{
    auto* data = new Datafield(frame().clone(), m_values, m_errSigmas);
    return data;
}

double& Datafield::operator[](size_t i)
{
    return m_values[i];
}

const double& Datafield::operator[](size_t i) const
{
    return m_values[i];
}

void Datafield::setAt(size_t i, double val)
{
    m_values[i] = val;
}

double Datafield::valAt(size_t i) const
{
    return m_values[i];
}

bool Datafield::hasErrorSigmas() const
{
    return !m_errSigmas.empty();
}

const std::vector<double>& Datafield::errorSigmas() const
{
    return m_errSigmas;
}

std::vector<double>& Datafield::errorSigmas()
{
    return m_errSigmas;
}

void Datafield::setAllTo(const double& value)
{
    for (double& v : m_values)
        v = value;
}

void Datafield::setVector(const std::vector<double>& vector)
{
    ASSERT(vector.size() == frame().size());
    m_values = vector;
}

size_t Datafield::rank() const
{
    return frame().rank();
}

size_t Datafield::size() const
{
    ASSERT(frame().size() == m_values.size());
    return frame().size();
}

bool Datafield::empty() const
{
    return size() == 0;
}

const Frame& Datafield::frame() const
{
    ASSERT(m_frame);
    return *m_frame;
}

const Scale& Datafield::axis(size_t k) const
{
    return frame().axis(k);
}

const Scale& Datafield::xAxis() const
{
    return frame().axis(0);
}

const Scale& Datafield::yAxis() const
{
    return frame().axis(1);
}

const std::vector<double>& Datafield::flatVector() const
{
    return m_values;
}

void Datafield::scale(double factor)
{
    size_t N = frame().size();
    for (size_t i = 0; i < N; ++i) {
        m_values[i] *= factor;
        if (!m_errSigmas.empty())
            m_errSigmas[i] *= factor;
    }
}

double Datafield::maxVal() const
{
    return *std::max_element(m_values.begin(), m_values.end());
}

double Datafield::minVal() const
{
    return *std::min_element(m_values.begin(), m_values.end());
}

Datafield* Datafield::crop(double xmin, double ymin, double xmax, double ymax) const
{
    const auto xclipped = std::make_unique<Scale>(xAxis().clipped(xmin, xmax));
    const auto yclipped = std::make_unique<Scale>(yAxis().clipped(ymin, ymax));

    std::vector<double> out(size());
    size_t iout = 0;
    for (size_t i = 0; i < size(); ++i) {
        double x = frame().projectedCoord(i, 0);
        double y = frame().projectedCoord(i, 1);
        if (xclipped->rangeComprises(x) && yclipped->rangeComprises(y))
            out[iout++] = m_values[i];
    }
    return new Datafield(frame().clone(), out);
}

Datafield* Datafield::crop(double xmin, double xmax) const
{
    const auto xclipped = std::make_unique<Scale>(xAxis().clipped(xmin, xmax));

    std::vector<double> out(size());
    size_t iout = 0;
    for (size_t i = 0; i < size(); ++i) {
        const double x = frame().projectedCoord(i, 0);
        if (xclipped->rangeComprises(x))
            out[iout++] = m_values[i];
    }
    return new Datafield(frame().clone(), out);
}

#ifdef BORNAGAIN_PYTHON

#include "PyCore/Embed/PyInterpreter.h" // Numpy::arrayND, Numpy::getDataPtr

PyObject* Datafield::npArray() const
{
    // TODO: Thoroughly check this function regarding index manipulations

    PyInterpreter::Numpy::initialize();

    ASSERT(rank() <= 2);

    std::vector<size_t> dimensions;
    for (size_t i = 0; i < rank(); i++)
        dimensions.push_back(axis(i).size());

    // for rot90 of 2-dim arrays to conform with numpy
    if (dimensions.size() == 2)
        std::swap(dimensions[0], dimensions[1]);

    // creating ndarray objects describing size of dimensions
    PyObjectPtr pyarray{PyInterpreter::Numpy::arrayND(dimensions)};
    if (!pyarray.valid())
        return nullptr;

    // get the pointer to the data buffer of the array (assumed to be C-contiguous)
    double* data{PyInterpreter::Numpy::getDataPtr(pyarray.get())};

    if (!data)
        return nullptr;

    double* array_buffer = data;

    // filling numpy array with output_data
    if (rank() == 2) {
        for (size_t i = 0; i < size(); ++i) {
            std::vector<int> axes_indices = frame().allIndices(i);
            size_t offset =
                axes_indices[0] + axis(0).size() * (axis(1).size() - 1 - axes_indices[1]);
            array_buffer[offset] = (*this)[i];
        }
    } else if (rank() == 1) {
        for (size_t i = 0; i < size(); ++i)
            *array_buffer++ = (*this)[i];
    }

    // returns a _new_ reference; ie. caller is responsible for the ref-count
    return pyarray.release();
}

#endif // BORNAGAIN_PYTHON

Datafield* Datafield::xProjection() const
{
    return create_xProjection(0, static_cast<int>(xAxis().size()) - 1);
}

Datafield* Datafield::xProjection(double yvalue) const
{
    int ybin_selected = static_cast<int>(yAxis().closestIndex(yvalue));
    return create_xProjection(ybin_selected, ybin_selected);
}

Datafield* Datafield::xProjection(double ylow, double yup) const
{
    int ybinlow = static_cast<int>(yAxis().closestIndex(ylow));
    int ybinup = static_cast<int>(yAxis().closestIndex(yup));
    return create_xProjection(ybinlow, ybinup);
}

Datafield* Datafield::yProjection() const
{
    return create_yProjection(0, static_cast<int>(xAxis().size()) - 1);
}

Datafield* Datafield::yProjection(double xvalue) const
{
    int xbin_selected = static_cast<int>(xAxis().closestIndex(xvalue));
    return create_yProjection(xbin_selected, xbin_selected);
}

Datafield* Datafield::yProjection(double xlow, double xup) const
{
    int xbinlow = static_cast<int>(xAxis().closestIndex(xlow));
    int xbinup = static_cast<int>(xAxis().closestIndex(xup));
    return create_yProjection(xbinlow, xbinup);
}

Datafield* Datafield::create_xProjection(int ybinlow, int ybinup) const
{
    std::vector<double> out(xAxis().size());
    for (size_t i = 0; i < size(); ++i) {
        int ybin = static_cast<int>(frame().projectedIndex(i, 1));
        if (ybin >= ybinlow && ybin <= ybinup) {
            double x = frame().projectedCoord(i, 0);
            ASSERT(xAxis().rangeComprises(x));
            size_t iout = xAxis().closestIndex(x);
            out[iout] += valAt(i);
        }
    }
    return new Datafield({xAxis().clone()}, out);
}

Datafield* Datafield::create_yProjection(int xbinlow, int xbinup) const
{
    std::vector<double> out(yAxis().size());
    for (size_t i = 0; i < size(); ++i) {
        int xbin = static_cast<int>(frame().projectedIndex(i, 0));
        if (xbin >= xbinlow && xbin <= xbinup) {

            // TODO: duplicates code from create_xProjection, move to Frame ?

            double y = frame().projectedCoord(i, 1);
            ASSERT(yAxis().rangeComprises(y));
            size_t iout = yAxis().closestIndex(y);
            out[iout] += valAt(i);
        }
    }
    return new Datafield({yAxis().clone()}, out);
}
