//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      GUI/View/PlotUtil/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/PlotUtil/ColorMap.h"
#include "Device/Data/Datafield.h"
#include "GUI/Model/Axis/AmplitudeAxisItem.h"
#include "GUI/Model/Axis/BasicAxisItem.h"
#include "GUI/Model/Data/IntensityDataItem.h"
#include "GUI/Model/Project/ProjectDocument.h"
#include "GUI/Support/Util/QCP_Util.h"
#include "GUI/View/PlotUtil/PlotConstants.h"
#include "GUI/View/PlotUtil/PlotEventInfo.h"
#include "GUI/View/PlotUtil/RangeUtil.h"
#include "GUI/View/Tool/UpdateTimer.h"
#include <qcustomplot.h>

namespace {

const int replot_update_interval = 10;
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 QCPRange(xmin + dx / 2., xmax - dx / 2.);
}

} // namespace

ColorMap::ColorMap(QWidget* parent)
    : ScientificPlot(parent, PLOT_TYPE::Plot2D)
    , m_customPlot(new QCustomPlot)
    , m_colorMap(nullptr)
    , m_colorScale(nullptr)
    , m_updateTimer(new UpdateTimer(replot_update_interval, this))
    , m_colorBarLayout(new QCPLayoutGrid)
{
    initColorMap();

    auto* vlayout = new QVBoxLayout(this);
    vlayout->setContentsMargins(0, 0, 0, 0);
    vlayout->setSpacing(0);
    vlayout->addWidget(m_customPlot);
    m_customPlot->setAttribute(Qt::WA_NoMousePropagation, false);
    setLayout(vlayout);

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

void ColorMap::setIntensityItem(IntensityDataItem* intensityItem)
{
    ScientificPlot::setIntensityItem(intensityItem);
    setColorMapFromItem();
    connectItem();
}

QRectF ColorMap::viewportRectangleInWidgetCoordinates()
{
    QCPRange xrange = m_customPlot->xAxis->range();
    QCPRange yrange = m_customPlot->yAxis->range();
    double left = xrange.lower;
    double right = xrange.upper;
    double top = yrange.upper;
    double bottom = yrange.lower;

    return QRectF(xAxisCoordToPixel(left), yAxisCoordToPixel(top),
                  xAxisCoordToPixel(right) - xAxisCoordToPixel(left),
                  yAxisCoordToPixel(bottom) - yAxisCoordToPixel(top));
}

PlotEventInfo ColorMap::eventInfo(double xpos, double ypos) const
{
    PlotEventInfo result(plotType());
    const IntensityDataItem* ii = intensityItem();
    if (!ii)
        return result;

    int nx(0), ny(0);
    m_colorMap->data()->coordToCell(xpos, ypos, &nx, &ny);

    result.setX(xpos);
    result.setY(ypos);
    result.setNx(nx);
    result.setNy(ny);

    result.setInAxesRange(axesRangeContains(xpos, ypos));
    result.setValue(m_colorMap->data()->cell(result.nx(), result.ny()));
    result.setLogValueAxis(ii->isLog());

    return result;
}

//! sets logarithmic scale
void ColorMap::setLogz()
{
    const IntensityDataItem* ii = intensityItem();
    if (!ii)
        return;
    bool logz = ii->isLog();
    m_colorBarLayout->setMinimumSize(logz ? colorbar_width_logz : colorbar_width, 10);
    GUI::QCP_Util::setLogz(m_colorScale, logz);
    replot();
}

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

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

void ColorMap::setGradient()
{
    const IntensityDataItem* ii = intensityItem();
    if (!ii)
        return;
    m_colorMap->setGradient(ii->currentGradientQCP());
    replot();
}

void ColorMap::setInterpolation()
{
    const IntensityDataItem* ii = intensityItem();
    if (!ii)
        return;
    m_colorMap->setInterpolate(ii->isInterpolated());
    replot();
}

//! Propagate zmin, zmax back to IntensityDataItem
void ColorMap::onDataRangeChanged(QCPRange newRange)
{
    IntensityDataItem* ii = intensityItem();
    ii->setLowerAndUpperZ(newRange.lower, newRange.upper);
    emit ii->updateOtherPlots(ii);
    gProjectDocument.value()->setModified();
}

//! Propagate xmin, xmax back to IntensityDataItem
void ColorMap::onXaxisRangeChanged(QCPRange newRange)
{
    IntensityDataItem* ii = intensityItem();
    ii->setLowerX(newRange.lower);
    ii->setUpperX(newRange.upper);
    emit ii->updateOtherPlots(ii);
    gProjectDocument.value()->setModified();
}

//! Propagate ymin, ymax back to IntensityDataItem
void ColorMap::onYaxisRangeChanged(QCPRange newRange)
{
    IntensityDataItem* ii = intensityItem();
    ii->setLowerY(newRange.lower);
    ii->setUpperY(newRange.upper);
    emit ii->updateOtherPlots(ii);
    gProjectDocument.value()->setModified();
}

//! Schedule replot for later execution by onTimeReplot() slot.

void ColorMap::replot()
{
    m_updateTimer->scheduleUpdate();
}

//! Replots ColorMap.

void ColorMap::onTimeToReplot()
{
    m_customPlot->replot();
}

//! creates and initializes the color map
void ColorMap::initColorMap()
{
    m_colorMap = new QCPColorMap(m_customPlot->xAxis, m_customPlot->yAxis);
    m_colorScale = new QCPColorScale(m_customPlot);
    m_colorMap->setColorScale(m_colorScale);

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

    m_colorScale->axis()->axisRect()->setMargins(QMargins(0, 0, 0, 0));
    m_colorScale->axis()->axisRect()->setAutoMargins(QCP::msNone);

    m_colorScale->setBarWidth(GUI::Constants::plot_colorbar_size());
    m_colorScale->axis()->setTickLabelFont(
        QFont(QFont().family(), GUI::Constants::plot_tick_label_size()));
    m_customPlot->xAxis->setTickLabelFont(
        QFont(QFont().family(), GUI::Constants::plot_tick_label_size()));
    m_customPlot->yAxis->setTickLabelFont(
        QFont(QFont().family(), GUI::Constants::plot_tick_label_size()));

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

void ColorMap::connectItem()
{
    const IntensityDataItem* ii = intensityItem();
    // data
    connect(ii, &IntensityDataItem::datafieldChanged, this, &ColorMap::onIntensityModified,
            Qt::UniqueConnection);

    // units
    connect(ii, &IntensityDataItem::axesUnitsReplotRequested, this, &ColorMap::onUnitsChanged,
            Qt::UniqueConnection);

    // color scheme
    connect(ii, &IntensityDataItem::gradientChanged, this, &ColorMap::setGradient,
            Qt::UniqueConnection);

    // interpolation
    connect(ii, &IntensityDataItem::interpolationChanged, this, &ColorMap::setInterpolation,
            Qt::UniqueConnection);

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

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

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

    setConnected(true);
}

void ColorMap::setConnected(bool isConnected)
{
    setAxesRangeConnected(isConnected);
    setDataRangeConnected(isConnected);
    setUpdateTimerConnected(isConnected);
}

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

void ColorMap::setAxesRangeConnected(bool isConnected)
{
    if (isConnected) {
        connect(m_customPlot->xAxis,
                static_cast<void (QCPAxis::*)(const QCPRange&)>(&QCPAxis::rangeChanged), this,
                &ColorMap::onXaxisRangeChanged, Qt::UniqueConnection);
        connect(m_customPlot->yAxis,
                static_cast<void (QCPAxis::*)(const QCPRange&)>(&QCPAxis::rangeChanged), this,
                &ColorMap::onYaxisRangeChanged, Qt::UniqueConnection);
    } else {
        disconnect(m_customPlot->xAxis,
                   static_cast<void (QCPAxis::*)(const QCPRange&)>(&QCPAxis::rangeChanged), this,
                   &ColorMap::onXaxisRangeChanged);
        disconnect(m_customPlot->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_colorMap, &QCPColorMap::dataRangeChanged, this, &ColorMap::onDataRangeChanged,
                Qt::UniqueConnection);
    else
        disconnect(m_colorMap, &QCPColorMap::dataRangeChanged, this, &ColorMap::onDataRangeChanged);
}

void ColorMap::setUpdateTimerConnected(bool isConnected)
{
    if (isConnected)
        connect(m_updateTimer, &UpdateTimer::timeToUpdate, this, &ColorMap::onTimeToReplot,
                Qt::UniqueConnection);
    else
        disconnect(m_updateTimer, &UpdateTimer::timeToUpdate, this, &ColorMap::onTimeToReplot);
}

//! to make fixed margins for whole colormap (change in axes labels wont affect axes rectangle)
void ColorMap::setFixedColorMapMargins()
{
    GUI::QCP_Util::setDefaultMargins(m_customPlot);
}

//! Sets initial state of ColorMap to match given intensity item.

void ColorMap::setColorMapFromItem()
{
    setAxesRangeFromItem();
    setAxesZoomFromItem();
    setAxesLabelsFromItem();
    setDataFromItem();
    setColorScaleAppearanceFromItem();
    setDataRangeFromItem();
    setLogz();
}

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

void ColorMap::setAxesRangeFromItem()
{
    const IntensityDataItem* ii = intensityItem();
    if (!ii)
        return;
    m_customPlot->setInteractions(QCP::iRangeDrag | QCP::iRangeZoom);
    m_customPlot->axisRect()->setupFullAxesBox(true);
    m_colorMap->data()->setSize(ii->xSize(), ii->ySize());
    m_colorMap->data()->setRange(qcpRange(ii->xMin(), ii->xMax(), ii->xSize()),
                                 qcpRange(ii->yMin(), ii->yMax(), ii->ySize()));
    replot();
}

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

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

//! Sets X,Y axes labels from item

void ColorMap::setAxesLabelsFromItem()
{
    const IntensityDataItem* ii = intensityItem();
    if (!ii)
        return;

    m_customPlot->xAxis->setLabel(ii->XaxisTitle());
    m_customPlot->yAxis->setLabel(ii->YaxisTitle());
    m_colorScale->setMargins(QMargins(0, 0, 0, 0));

    replot();
}

//! Sets the intensity values to ColorMap.

void ColorMap::setDataFromItem()
{
    const IntensityDataItem* ii = intensityItem();
    if (!ii)
        return;
    const Datafield* data = ii->c_field();
    if (!data) {
        m_colorMap->data()->clear();
        return;
    }

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

    for (int ix = 0; ix < nx; ++ix)
        for (int iy = 0; iy < ny; ++iy)
            m_colorMap->data()->setCell(ix, iy, (*data)[iy + ny * 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_customPlot);
    m_customPlot->axisRect()->setMarginGroup(QCP::msBottom | QCP::msTop, marginGroup);
    m_colorScale->setMarginGroup(QCP::msBottom | QCP::msTop, marginGroup);
}

void ColorMap::setDataRangeFromItem()
{
    if (!intensityItem())
        return;
    setDataRangeConnected(false);
    m_colorMap->setDataRange(GUI::View::RangeUtil::itemDataZoom(intensityItem()));
    setDataRangeConnected(true);
    replot();
}

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

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

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

void ColorMap::marginsChangedNotify()
{
    QMargins axesMargins = m_customPlot->axisRect()->margins();
    //    QMargins colorBarMargins = m_colorScale->margins();
    //    QMargins colorScaleMargins = m_colorScale->axis()->axisRect()->margins();

    double left = axesMargins.left();
    //    double right = axesMargins.right() + colorBarMargins.right() + m_colorScale->barWidth()
    //            + colorScaleMargins.right() + m_colorBarLayout->rect().width();

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

    emit marginsChanged(left, right);
}

IntensityDataItem* ColorMap::intensityItem() const
{
    return intensityDataItem();
}
