//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      GUI/View/Plotter/ColorMap.cpp
//! @brief     Implements class ColorMap.
//!
//! @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 "GUI/View/Plotter/ColorMap.h"
#include "Base/Py/PyFmt.h"
#include "Base/Util/Assert.h"
#include "Device/Data/Datafield.h"
#include "GUI/Model/Axis/AmplitudeAxisItem.h"
#include "GUI/Model/Axis/BasicAxisItem.h"
#include "GUI/Model/Data/Data2DItem.h"
#include "GUI/Model/Project/ProjectDocument.h"
#include "GUI/View/Base/Fontsize.h"
#include "GUI/View/Plotter/RangeUtil.h"
#include "GUI/View/Widget/AppConfig.h"
#include <qcustomplot.h>

namespace {

const int colorbar_width_logz = 50;
const int colorbar_width = 80;

// Converts xmin (low edge of first bin) and xmax (upper edge of last bin) to the
// range expected by QCPColorMapData::setRange.
QCPRange qcpRange(double xmin, double xmax, int nbins)
{
    double dx = (xmax - xmin) / nbins;
    return {xmin + dx / 2., xmax - dx / 2.};
}

} // namespace


ColorMap::ColorMap()
    : ScientificPlot(PLOT_TYPE::Plot2D)
    , m_qcp_map(new QCPColorMap(m_qcp->xAxis, m_qcp->yAxis))
    , m_color_scale(new QCPColorScale(m_qcp))
    , m_color_bar_layout(new QCPLayoutGrid)
{
    m_qcp->xAxis->setTickLabelFont(QFont(QFont().family(), GUI::Style::fontSizeSmall()));
    m_qcp->yAxis->setTickLabelFont(QFont(QFont().family(), GUI::Style::fontSizeSmall()));
    m_qcp->setAttribute(Qt::WA_NoMousePropagation, false);

    connect(m_qcp, &QCustomPlot::afterReplot, this, &ColorMap::marginsChangedNotify);

    m_color_scale->axis()->axisRect()->setMargins(QMargins(0, 0, 0, 0));
    m_color_scale->axis()->axisRect()->setAutoMargins(QCP::msNone);
    m_color_scale->setBarWidth(GUI::Style::SizeOfLetterM().width());
    m_color_scale->axis()->setTickLabelFont(QFont(QFont().family(), GUI::Style::fontSizeSmall()));
    m_qcp_map->setColorScale(m_color_scale);

    m_color_bar_layout->addElement(0, 0, m_color_scale);
    m_color_bar_layout->setMinimumSize(colorbar_width_logz, 10);
    const auto base_size = GUI::Style::SizeOfLetterM(this).width() * 0.5;
    m_color_bar_layout->setMargins(QMargins(base_size, 0, base_size, 0));

    auto* vlayout = new QVBoxLayout(this);
    vlayout->setContentsMargins(0, 0, 0, 0); // essential for matching idealized and pixelized maps
    vlayout->setSpacing(0);
    vlayout->addWidget(m_qcp);

    setMouseTrackingEnabled(true);
    //    setFixedColorMapMargins();
}

void ColorMap::itemToMap(Data2DItem* item)
{
    ASSERT(item);
    m_data_item = item;

    // set nullptr at destruction
    disconnect(m_data_item, &DataItem::destroyed, this, nullptr);
    connect(m_data_item, &DataItem::destroyed, [this] { m_data_item = nullptr; });

    //... Set initial state to match given intensity item.
    setAxesRangeFromItem();
    setAxesZoomFromItem();
    setAxesLabelsFromItem();
    setDataFromItem();
    setColorScaleAppearanceFromItem();
    setDataRangeFromItem();
    setLogz();

    // data
    connect(item, &Data2DItem::datafieldChanged, this, &ColorMap::onIntensityModified,
            Qt::UniqueConnection);

    // units
    connect(item, &Data2DItem::axesUnitsReplotRequested, this, &ColorMap::onUnitsChanged,
            Qt::UniqueConnection);

    // color scheme
    connect(gApp.get(), &AppConfig::gradientChanged, this, &ColorMap::setGradient,
            Qt::UniqueConnection);

    // interpolation
    connect(item, &Data2DItem::interpolationChanged, this, &ColorMap::setInterpolation,
            Qt::UniqueConnection);

    // x axis
    connect(item->axItemX(), &BasicAxisItem::axisRangeChanged, this, &ColorMap::setAxesZoomFromItem,
            Qt::UniqueConnection);
    connect(item->axItemX(), &BasicAxisItem::axisTitleChanged, this,
            &ColorMap::setAxesLabelsFromItem, Qt::UniqueConnection);

    // y axis
    connect(item->axItemY(), &BasicAxisItem::axisRangeChanged, this, &ColorMap::setAxesZoomFromItem,
            Qt::UniqueConnection);
    connect(item->axItemY(), &BasicAxisItem::axisTitleChanged, this,
            &ColorMap::setAxesLabelsFromItem, Qt::UniqueConnection);

    // z axis
    connect(item->zAxisItem(), &BasicAxisItem::axisRangeChanged, this,
            &ColorMap::setDataRangeFromItem, Qt::UniqueConnection);
    connect(item->zAxisItem(), &AmplitudeAxisItem::logScaleChanged, this, &ColorMap::setLogz,
            Qt::UniqueConnection);
    connect(item->zAxisItem(), &BasicAxisItem::axisVisibilityChanged, this,
            &ColorMap::setColorScaleVisible, Qt::UniqueConnection);

    setAxesRangeConnected(true);
    setDataRangeConnected(true);
}

QString ColorMap::infoString(const QPointF& pos) const
{
    double x = m_qcp->xAxis->pixelToCoord(pos.x());
    double y = m_qcp->yAxis->pixelToCoord(pos.y());

    int ix, iy;
    m_qcp_map->data()->coordToCell(x, y, &ix, &iy);

    double intensity = m_qcp_map->data()->cell(ix, iy);
    QString intensityString =
        data2DItem()->isLog() ? QString::fromStdString(Py::Fmt::printScientificDouble(intensity))
                              : QString::number(intensity, 'f', 2);
    return QString(" [x: %1, y: %2]    [binx: %3, biny:%4]    [value: %5]")
        .arg(QString::number(x, 'f', 4))
        .arg(QString::number(y, 'f', 4), 2)
        .arg(ix, 2)
        .arg(iy, 2)
        .arg(intensityString);
}

//! sets logarithmic scale
void ColorMap::setLogz()
{
    const Data2DItem* ii = data2DItem();
    if (!ii)
        return;
    bool logz = ii->isLog();
    m_color_bar_layout->setMinimumSize(logz ? colorbar_width_logz : colorbar_width, 10);
    GUI::QCP_RangeUtil::setLogz(m_color_scale, logz);
    replot();
}

void ColorMap::onIntensityModified()
{
    setDataFromItem();
    setAxesRangeFromItem();
}

void ColorMap::onUnitsChanged()
{
    setAxesRangeFromItem();
    setAxesZoomFromItem();
    setAxesLabelsFromItem();
}

void ColorMap::setGradient()
{
    const Data2DItem* ii = data2DItem();
    if (!ii)
        return;
    m_qcp_map->setGradient(gApp->currentColorGradient());
    replot();
}

void ColorMap::setInterpolation()
{
    const Data2DItem* ii = data2DItem();
    if (!ii)
        return;
    m_qcp_map->setInterpolate(ii->isInterpolated());
    replot();
}

//! Propagate zmin, zmax back to Data2DItem
void ColorMap::onDataRangeChanged(QCPRange range)
{
    Data2DItem* ii = data2DItem();
    ii->setZrange(range.lower, range.upper);
    ii->zAxisItem()->adjustLogRangeOrders();
    emit ii->updateOtherPlots(ii);
    gDoc->setModified();
}

//! Propagate xmin, xmax back to Data2DItem
void ColorMap::onXaxisRangeChanged(QCPRange range)
{
    Data2DItem* ii = data2DItem();
    ii->setXrange(range.lower, range.upper);
    emit ii->updateOtherPlots(ii);
    gDoc->setModified();
}

//! Propagate ymin, ymax back to Data2DItem
void ColorMap::onYaxisRangeChanged(QCPRange range)
{
    Data2DItem* ii = data2DItem();
    ii->setYrange(range.lower, range.upper);
    emit ii->updateOtherPlots(ii);
    gDoc->setModified();
}

//! Connects/disconnects signals related to ColorMap's X,Y axes rectangle change.

void ColorMap::setAxesRangeConnected(bool isConnected)
{
    if (isConnected) {
        connect(m_qcp->xAxis,
                static_cast<void (QCPAxis::*)(const QCPRange&)>(&QCPAxis::rangeChanged), this,
                &ColorMap::onXaxisRangeChanged, Qt::UniqueConnection);
        connect(m_qcp->yAxis,
                static_cast<void (QCPAxis::*)(const QCPRange&)>(&QCPAxis::rangeChanged), this,
                &ColorMap::onYaxisRangeChanged, Qt::UniqueConnection);
    } else {
        disconnect(m_qcp->xAxis,
                   static_cast<void (QCPAxis::*)(const QCPRange&)>(&QCPAxis::rangeChanged), this,
                   &ColorMap::onXaxisRangeChanged);
        disconnect(m_qcp->yAxis,
                   static_cast<void (QCPAxis::*)(const QCPRange&)>(&QCPAxis::rangeChanged), this,
                   &ColorMap::onYaxisRangeChanged);
    }
}

//! Connects/disconnects signals related to ColorMap's Z-axis (min,max) change.

void ColorMap::setDataRangeConnected(bool isConnected)
{
    if (isConnected)
        connect(m_qcp_map, &QCPColorMap::dataRangeChanged, this, &ColorMap::onDataRangeChanged,
                Qt::UniqueConnection);
    else
        disconnect(m_qcp_map, &QCPColorMap::dataRangeChanged, this, &ColorMap::onDataRangeChanged);
}

//! Sets (xmin,xmax,nbins) and (ymin,ymax,nbins) of ColorMap from intensity item.

void ColorMap::setAxesRangeFromItem()
{
    Data2DItem* ii = data2DItem();
    if (!ii)
        return;
    m_qcp->setInteractions(QCP::iRangeDrag | QCP::iRangeZoom);
    m_qcp->axisRect()->setupFullAxesBox(true);
    m_qcp_map->data()->setSize(ii->xSize(), ii->ySize());
    m_qcp_map->data()->setRange(qcpRange(ii->xMin(), ii->xMax(), ii->xSize()),
                                qcpRange(ii->yMin(), ii->yMax(), ii->ySize()));
    ii->setAxesRangeToData();
    replot();
}

//! Sets zoom range of X,Y axes as in intensity item.

void ColorMap::setAxesZoomFromItem()
{
    const Data2DItem* ii = data2DItem();
    if (!ii)
        return;
    setAxesRangeConnected(false);
    m_qcp->xAxis->setRange(ii->lowerX(), ii->upperX());
    m_qcp->yAxis->setRange(ii->lowerY(), ii->upperY());
    setAxesRangeConnected(true);
    replot();
}

//! Sets X,Y axes labels from item

void ColorMap::setAxesLabelsFromItem()
{
    const Data2DItem* ii = data2DItem();
    if (!ii)
        return;

    m_qcp->xAxis->setLabel(ii->xAxisLabel());
    m_qcp->yAxis->setLabel(ii->yAxisLabel());
    m_color_scale->setMargins(QMargins(0, 0, 0, 0));

    replot();
}

//! Sets the intensity values to ColorMap.

void ColorMap::setDataFromItem()
{
    const Data2DItem* ii = data2DItem();
    if (!ii)
        return;
    const Datafield* df = ii->c_field();
    if (!df) {
        m_qcp_map->data()->clear();
        return;
    }

    int nx(ii->xSize()); // outside of the loop because of slow retrieval
    int ny(ii->ySize());
    m_qcp_map->data()->setSize(nx, ny);

    for (int iy = 0; iy < ny; ++iy)
        for (int ix = 0; ix < nx; ++ix)
            m_qcp_map->data()->setCell(ix, iy, (*df)[iy * nx + ix]);
}

//! Sets the appearance of color scale (visibility, gradient type) from intensity item.

void ColorMap::setColorScaleAppearanceFromItem()
{
    setColorScaleVisible();
    setGradient();
    setInterpolation();
    // make sure the axis rect and color scale synchronize their bottom and top margins (so they
    // line up):
    auto* marginGroup = new QCPMarginGroup(m_qcp);
    m_qcp->axisRect()->setMarginGroup(QCP::msBottom | QCP::msTop, marginGroup);
    m_color_scale->setMarginGroup(QCP::msBottom | QCP::msTop, marginGroup);
}

void ColorMap::setDataRangeFromItem()
{
    const Data2DItem* ii = data2DItem();
    if (!ii)
        return;
    setDataRangeConnected(false);
    m_qcp_map->setDataRange(QCPRange(ii->lowerZ(), ii->upperZ()));
    setDataRangeConnected(true);
    replot();
}

void ColorMap::setColorScaleVisible()
{
    const Data2DItem* ii = data2DItem();
    if (!ii)
        return;
    bool visibility_flag = ii->zAxisItem()->isVisible();

    m_color_bar_layout->setVisible(visibility_flag);
    if (visibility_flag) {
        // add it to the right of the main axis rect
        if (!m_qcp->plotLayout()->hasElement(0, 1))
            m_qcp->plotLayout()->addElement(0, 1, m_color_bar_layout);
    } else {
        for (int i = 0; i < m_qcp->plotLayout()->elementCount(); ++i)
            if (m_qcp->plotLayout()->elementAt(i) == m_color_bar_layout)
                m_qcp->plotLayout()->takeAt(i);
        m_qcp->plotLayout()->simplify();
    }
    replot();
}

//! Calculates left, right margins around color map to report to projection plot.

void ColorMap::marginsChangedNotify()
{
    QMargins axesMargins = m_qcp->axisRect()->margins();

    double left = axesMargins.left();
    double right = axesMargins.right() + m_color_bar_layout->rect().width();

    emit marginsChanged(left, right);
}
