//  ************************************************************************************************
//
//  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/MakeScale.h"
#include "Base/Axis/Scale.h"
#include "Base/Util/Assert.h"
#include <algorithm>
#include <random>

#ifdef BORNAGAIN_PYTHON

#include "PyCore/Embed/PyInterpreter.h" // Numpy::arrayND, Numpy::getDataPtr
#include "PyCore/Embed/PyObjectPtr.h"
#include <Python.h>

namespace {

PyObject* npExport(const Frame& frame, const std::vector<double>& flatData)
{
    if (flatData.empty())
        return Py_BuildValue("");

    // TODO: Thoroughly check this function regarding index manipulations

    PyInterpreter::Numpy::initialize();

    ASSERT(frame.rank() <= 2);

    std::vector<size_t> dimensions;
    for (size_t i = 0; i < frame.rank(); i++)
        dimensions.push_back(frame.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)};
    ASSERT(pyarray.valid());

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

    double* array_buffer = data;

    // filling numpy array with output_data
    if (frame.rank() == 2) {
        for (size_t i = 0; i < frame.size(); ++i)
            *array_buffer++ = flatData[i];
    } else if (frame.rank() == 1) {
        for (size_t i = 0; i < frame.size(); ++i)
            *array_buffer++ = flatData[i];
    } else
        ASSERT_NEVER;

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

} // namespace

#endif // BORNAGAIN_PYTHON


Datafield::Datafield(std::string title, const Frame* frame, std::vector<double> values,
                     std::vector<double> errSigmas)
    : m_title(std::move(title))
    , m_frame(frame)
    , m_values(std::move(values))
    , m_err_sigmas(std::move(errSigmas))
{
    ASSERT(m_frame);
    ASSERT(m_values.size() == m_frame->size());
    ASSERT(m_err_sigmas.empty() || m_err_sigmas.size() == m_values.size());
}

Datafield::Datafield(std::string title, const Frame* frame)
    : Datafield(title, frame, std::vector<double>(frame->size(), 0.))
{
}

Datafield::Datafield(const Frame* frame, std::vector<double> values, std::vector<double> errSigmas)
    : Datafield("", frame, values, errSigmas)
{
}

Datafield::Datafield(const Frame* frame)
    : Datafield("", frame)
{
}

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

Datafield::Datafield(std::vector<const Scale*> axes)
    : Datafield(new Frame(axes))
{
}

Datafield::Datafield(const std::string& xlabel, const std::string& ylabel,
                     const std::vector<std::vector<double>>& vec)
{
    const size_t nrows = vec.size();
    const size_t ncols = vec[0].size();
    ASSERT(nrows > 0);
    ASSERT(ncols > 0);

    std::vector<const Scale*> axes;
    if (nrows < 2)
        axes = {newEquiDivision(xlabel, ncols, 0.0, (double)ncols)};
    else if (ncols < 2)
        axes = {newEquiDivision(ylabel, nrows, 0.0, (double)nrows)};
    else
        axes = {newEquiDivision(xlabel, ncols, 0.0, (double)ncols),
                newEquiDivision(ylabel, nrows, 0.0, (double)nrows)};
    m_frame = std::make_unique<Frame>(std::move(axes));

    setVector2D(vec);
}

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

Datafield::Datafield(const Datafield& other)
    : Datafield(other.title(), other.m_frame->clone(), other.m_values, other.m_err_sigmas)
{
}

Datafield::~Datafield() = default;

Datafield* Datafield::clone() const
{
    return new Datafield(title(), frame().clone(), m_values, m_err_sigmas);
}

void Datafield::setTitle(const std::string& title)
{
    m_title = title;
}
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_err_sigmas.empty();
}

const std::vector<double>& Datafield::errorSigmas() const
{
    return m_err_sigmas;
}
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;
}

void Datafield::setVector2D(const std::vector<std::vector<double>>& in)
{
    const size_t ncols = axis(0).size();
    const size_t nrows = axis(1).size();

    ASSERT(in.size() == nrows);
    m_values.clear();
    m_values.reserve(nrows * ncols);
    for (size_t j = 0; j < nrows; ++j) {
        ASSERT(in[j].size() == ncols);
        for (size_t i = 0; i < ncols; ++i)
            m_values.push_back(in[j][i]);
    }
}

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_err_sigmas.empty())
            m_err_sigmas[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
{
    ASSERT(rank() == 2);

    const size_t N = size();
    std::vector<double> out;
    for (size_t i = 0; i < N; ++i) {
        const Bin1D& x_bin = frame().projectedBin(i, 0);
        const Bin1D& y_bin = frame().projectedBin(i, 1);
        if (xmin <= x_bin.max() && x_bin.min() <= xmax && ymin <= y_bin.max()
            && y_bin.min() <= ymax)
            out.push_back(m_values[i]);
    }

    const Scale* xclipped = xAxis().clipped(xmin, xmax).clone();
    const Scale* yclipped = yAxis().clipped(ymin, ymax).clone();
    Frame* outframe = new Frame(xclipped, yclipped);

    ASSERT(outframe->size() == out.size());
    return new Datafield(outframe, out);
}

Datafield* Datafield::crop(double xmin, double xmax) const
{
    ASSERT(rank() == 1);

    const size_t N = size();
    std::vector<double> out;
    std::vector<double> errout;
    for (size_t i = 0; i < N; ++i) {
        const Bin1D& x_bin = frame().projectedBin(i, 0);
        if (xmin <= x_bin.max() && x_bin.min() <= xmax) {
            out.push_back(m_values[i]);
            if (hasErrorSigmas())
                errout.push_back(m_err_sigmas[i]);
        }
    }
    const Scale* xclipped = xAxis().clipped(xmin, xmax).clone();
    Frame* outframe = new Frame(xclipped);

    ASSERT(outframe->xAxis().size() == out.size());
    return new Datafield(outframe, out, errout);
}

#ifdef BORNAGAIN_PYTHON

PyObject* Datafield::npXcenters() const
{
    return ::npExport(Frame(xAxis().clone()), xAxis().binCenters());
}

PyObject* Datafield::npArray() const
{
    return ::npExport(frame(), flatVector());
}

PyObject* Datafield::npErrors() const
{
    return ::npExport(frame(), errorSigmas());
}

#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);
            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);
            size_t iout = yAxis().closestIndex(y);
            out[iout] += valAt(i);
        }
    }
    return new Datafield({yAxis().clone()}, out);
}

Datafield Datafield::plottableField() const
{
    return {title(), frame().plottableFrame(), m_values, m_err_sigmas};
}

Datafield Datafield::angularField(double lambda, double alpha_i) const
{
    return {title(), frame().angularFrame(lambda, alpha_i), m_values, m_err_sigmas};
}

Datafield Datafield::qSpaceField(double lambda, double alpha_i) const
{
    return {title(), frame().qSpaceFrame(lambda, alpha_i), m_values, m_err_sigmas};
}

Datafield Datafield::flat() const
{
    return {title(), frame().flat(), m_values, m_err_sigmas};
}

Datafield Datafield::noisy(double prefactor, double minimum) const
{
    auto const seed = 123;
    static auto urbg = std::mt19937(seed);

    std::vector<double> outval(size());
    std::vector<double> errval(size());
    for (size_t i = 0; i < size(); ++i) {
        double mean = valAt(i);
        double stdv = prefactor * sqrt(mean);
        auto norm = std::normal_distribution<double>(mean, stdv);
        outval[i] = std::max(norm(urbg), minimum);
        errval[i] = stdv;
    }
    return {title(), frame().clone(), outval, errval};
}

Datafield Datafield::normalizedToMax() const
{
    const double maxval = *std::max_element(m_values.begin(), m_values.end());
    std::vector<double> outval(size());
    std::vector<double> errval(size());
    for (size_t i = 0; i < size(); ++i) {
        outval[i] = m_values[i] / maxval;
        if (hasErrorSigmas())
            errval[i] = m_err_sigmas[i] / maxval;
    }
    return {title(), frame().clone(), outval, errval};
}

std::vector<std::vector<double>> Datafield::values2D() const
{
    ASSERT(rank() == 2);
    std::vector<std::vector<double>> result;

    const size_t nrows = axis(1).size();
    const size_t ncols = axis(0).size();

    result.resize(nrows);

    for (size_t row = 0; row < nrows; ++row) {
        result[row].resize(ncols, 0.0);
        for (size_t col = 0; col < ncols; ++col)
            result[row][col] = m_values[row * ncols + col];
    }

    return result;
}
