//  ************************************************************************************************
//
//  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/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"

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()
    : m_qcp_map(new QCPColorMap(xAxis, yAxis))
    , m_color_scale(new QCPColorScale(this))
    , m_color_bar_layout(new QCPLayoutGrid)
{
    const QFont font(QFont().family(), GUI::Font::fontSizeSmall());
    xAxis->setTickLabelFont(font);
    yAxis->setTickLabelFont(font);
    setAttribute(Qt::WA_NoMousePropagation, false);

    connect(this, &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::Font::SizeOfLetterM().width());
    m_color_scale->axis()->setTickLabelFont(font);
    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::Font::SizeOfLetterM(this).width() * 0.5;
    m_color_bar_layout->setMargins(QMargins(base_size, 0, base_size, 0));

    setMouseTrackingEnabled(true);
}

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

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

    // set nullptr at destruction
    connect(m_data_item, &QObject::destroyed, this, &ColorMap::onDataDestroyed,
            Qt::UniqueConnection);

    // data
    connect(item, &Data2DItem::datafieldChanged, this, &ColorMap::onIntensityModified,
            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(), &BasicAxisItem::logScaleChanged, this, &ColorMap::setLogZ,
            Qt::UniqueConnection);
    connect(item->zAxisItem(), &AmplitudeAxisItem::axisVisibilityChanged, this,
            &ColorMap::setColorScaleVisible, Qt::UniqueConnection);

    setAxesRangeConnected(true);
    setDataRangeConnected(true);
}

QString ColorMap::infoString(double x, double y) const
{
    int ix, iy;
    cmData()->coordToCell(x, y, &ix, &iy);

    double intensity = cmData()->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);
}

QCPColorMapData* ColorMap::cmData() const
{
    return m_qcp_map->data();
}

//! 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();
    setAxesLabelsFromItem(); // when datafield is loaded, axes with labels are loaded as well
}

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(xAxis, static_cast<void (QCPAxis::*)(const QCPRange&)>(&QCPAxis::rangeChanged),
                this, &ColorMap::onXaxisRangeChanged, Qt::UniqueConnection);
        connect(yAxis, static_cast<void (QCPAxis::*)(const QCPRange&)>(&QCPAxis::rangeChanged),
                this, &ColorMap::onYaxisRangeChanged, Qt::UniqueConnection);
    } else {
        disconnect(xAxis, static_cast<void (QCPAxis::*)(const QCPRange&)>(&QCPAxis::rangeChanged),
                   this, &ColorMap::onXaxisRangeChanged);
        disconnect(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;
    setInteractions(QCP::iRangeDrag | QCP::iRangeZoom);
    axisRect()->setupFullAxesBox(true);
    cmData()->setSize(ii->xSize(), ii->ySize());
    cmData()->setRange(::qcpRange(ii->xMin(), ii->xMax(), ii->xSize()),
                       ::qcpRange(ii->yMin(), ii->yMax(), ii->ySize()));
    if (ii->c_field() && !ii->isArgRangeLocked())
        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);
    xAxis->setRange(ii->lowerX(), ii->upperX());
    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;

    xAxis->setLabel(ii->xAxisLabel());
    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) {
        cmData()->clear();
        return;
    }

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

    for (int iy = 0; iy < ny; ++iy)
        for (int ix = 0; ix < nx; ++ix)
            cmData()->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(this);
    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 (!plotLayout()->hasElement(0, 1))
            plotLayout()->addElement(0, 1, m_color_bar_layout);
    } else {
        for (int i = 0; i < plotLayout()->elementCount(); ++i)
            if (plotLayout()->elementAt(i) == m_color_bar_layout)
                plotLayout()->takeAt(i);
        plotLayout()->simplify();
    }
    replot();
}

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

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

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

    emit marginsChanged(left, right);
}

void ColorMap::onDataDestroyed()
{
    m_data_item = nullptr;
}