//  ************************************************************************************************
//
//  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
namespace {

Arrayf64Wrapper arrayExport(const Frame& frame, const std::vector<double>& flatData,
                            const bool owndata)
{
    ASSERT(frame.rank() <= 2);

    const std::size_t n_dims = frame.rank();
    std::vector<std::size_t> dimensions(n_dims);

    for (std::size_t i = 0; i < n_dims; i++) {
        dimensions[i] = frame.axis(i).size();
    }

    // for rot90 of 2-dim arrays to conform with numpy
    // TODO: why this rotation is required?
    if (n_dims == 2)
        std::swap(dimensions[0], dimensions[1]);

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

} // namespace

#endif // BORNAGAIN_PYTHON


Datafield::Datafield(const std::string& title, const Frame& frame,
                     const std::vector<double>& values, const std::vector<double>& errSigmas)
    : m_title(title)
    , m_frame(frame.clone())
    , m_values(values)
    , m_err_sigmas(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(const std::string& title, const Frame& frame)
    : Datafield(title, frame, std::vector<double>(frame.size(), 0.))
{
}

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

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

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

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

Datafield::Datafield(const std::string& xlabel, const std::string& ylabel, const double2d_t& 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>(axes);

    setVector2D(vec);
}

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

Datafield::Datafield(const Datafield& other)
    : Datafield(other.title(), *other.m_frame, other.m_values, other.m_err_sigmas)
{
}

Datafield::~Datafield() = default;

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

void Datafield::setTitle(const std::string& title)
{
    m_title = title;
}

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

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

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 double2d_t& in)
{
    m_values = FieldUtil::flatten(in);
}

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

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

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);
}

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();
    const Frame outframe = 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();
    const Frame outframe = Frame(xclipped);

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

#ifdef BORNAGAIN_PYTHON

Arrayf64Wrapper Datafield::xCenters() const
{
    // NOTE: the array data is produced on the fly;
    // therefore, the ArrayDescriptor must store the data.
    return ::arrayExport(Frame(xAxis().clone()), xAxis().binCenters(),
                         /* owndata= */ true);
}

Arrayf64Wrapper Datafield::dataArray() const
{
    return ::arrayExport(frame(), flatVector(), /* owndata= */ false);
}

Arrayf64Wrapper Datafield::errors() const
{
    return ::arrayExport(frame(), errorSigmas(), /* owndata= */ false);
}

#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(std::vector<const Scale*>{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(std::vector<const Scale*>{yAxis().clone()}, out);
}

Datafield Datafield::plottableField() const
{
    return {title(), frame().plottableFrame(), 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(), 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(), outval, errval};
}

double2d_t Datafield::values2D() const
{
    return FieldUtil::reshapeTo2D(m_values, axis(1).size());
}