#include "stretchtoolbar.h"
#include <QVBoxLayout>
#include <QDebug>
#include <QToolButton>
#include "imageringlist.h"

const float BLACK_POINT_SIGMA = -2.8f;
const float MAD_TO_SIGMA = 1.4826f;
const float TARGET_BACKGROUND = 0.25f;

float MTF(float x, float m)
{
    if(x < 0)return 0;
    if(x > 1)return 1;
    return ((m - 1) * x) / ((2 * m - 1) * x - m);
}

StretchToolbar::StretchToolbar(QWidget *parent) : QToolBar(tr("Stretch toolbar"), parent)
{
    setObjectName("stretchtoolbar");
    QWidget *lum = new QWidget(this);
    QVBoxLayout *vbox1 = new QVBoxLayout(lum);
    m_stfSlider = new STFSlider(Qt::white, this);
    vbox1->addWidget(m_stfSlider);

    /*m_stfSliderR = new STFSlider(Qt::red, this);
    m_stfSliderG = new STFSlider(Qt::green, this);
    m_stfSliderB = new STFSlider(Qt::blue, this);
    QWidget *rgb = new QWidget(this);
    QVBoxLayout *vbox2 = new QVBoxLayout(rgb);
    vbox2->setSpacing(0);
    vbox2->addWidget(m_stfSliderR);
    vbox2->addWidget(m_stfSliderG);
    vbox2->addWidget(m_stfSliderB);*/

    m_stack = new QStackedWidget(this);
    m_stack->addWidget(lum);
    //m_stack->addWidget(rgb);
    //m_stack->setCurrentIndex(0);
    addWidget(m_stack);

    connect(m_stfSlider, &STFSlider::paramChanged, [this](float blackPoint, float midPoint, float whitePoint){
        m_mtfParam.blackPoint[0] = m_mtfParam.blackPoint[1] = m_mtfParam.blackPoint[2] = blackPoint;
        m_mtfParam.midPoint[0] = m_mtfParam.midPoint[1] = m_mtfParam.midPoint[2] = midPoint;
        m_mtfParam.whitePoint[0] = m_mtfParam.whitePoint[1] = m_mtfParam.whitePoint[2] = whitePoint;
        emit paramChanged(m_mtfParam);
    });
    /*connect(m_stfSliderR, &STFSlider::paramChanged, [this](float blackPoint, float midPoint, float whitePoint){
        m_mtfParam.blackPoint[0] = blackPoint;
        m_mtfParam.midPoint[0] = midPoint;
        m_mtfParam.whitePoint[0] = whitePoint;
        emit paramChanged(m_mtfParam);
    });
    connect(m_stfSliderG, &STFSlider::paramChanged, [this](float blackPoint, float midPoint, float whitePoint){
        m_mtfParam.blackPoint[1] = blackPoint;
        m_mtfParam.midPoint[1] = midPoint;
        m_mtfParam.whitePoint[1] = whitePoint;
        emit paramChanged(m_mtfParam);
    });
    connect(m_stfSliderB, &STFSlider::paramChanged, [this](float blackPoint, float midPoint, float whitePoint){
        m_mtfParam.blackPoint[2] = blackPoint;
        m_mtfParam.midPoint[2] = midPoint;
        m_mtfParam.whitePoint[2] = whitePoint;
        emit paramChanged(m_mtfParam);
    });*/

    QAction *autoStretchButton = addAction(QIcon(":/nuke.png"), tr("Auto Stretch F12"));
    autoStretchButton->setShortcut(Qt::Key_F12);
    connect(autoStretchButton, &QAction::triggered, this, &StretchToolbar::autoStretch);

    QAction *resetButton = addAction(style()->standardIcon(QStyle::SP_DialogResetButton), tr("Reset Screen Transfer Function F11"));
    resetButton->setShortcut(Qt::Key_F11);
    connect(resetButton, &QAction::triggered, this, &StretchToolbar::resetMTF);

    QAction *invertButton = addAction(QIcon(":/invert.png"), tr("Invert colors"));
    invertButton->setCheckable(true);
    connect(invertButton, SIGNAL(toggled(bool)), this, SIGNAL(invert(bool)));

    QAction *falseColorButton = addAction(QIcon(":/falsecolor.png"), tr("False color"));
    falseColorButton->setCheckable(true);
    connect(falseColorButton, SIGNAL(toggled(bool)), this, SIGNAL(falseColor(bool)));

    QAction *superPixelButton = addAction(QIcon(":/bayer.png"), tr("Debayer CFA"));
    superPixelButton->setCheckable(true);
    connect(superPixelButton, SIGNAL(toggled(bool)), this, SIGNAL(superPixel(bool)));

    m_autoStretchOnLoad = addAction(QIcon(":/nuke_a.png"), tr("Apply auto stretch on load"));
    m_autoStretchOnLoad->setCheckable(true);
}

void StretchToolbar::stretchImage(Image *img)
{
    if(img)
    {
        if(img->rawImage())
        {
            const RawImage::Stats &stats = img->rawImage()->imageStats();
            int ch = img->rawImage()->channels() == 1 ? 1 : 3;
            float bp2 = 0;
            float mid2 = 0;
            float max2 = 0;
            for(int i=0; i < ch; i++)
            {
                double median, mad, max;
                median = stats.m_median[i];
                mad = stats.m_mad[i];
                max = stats.m_max[i];
                median /= img->rawImage()->norm();
                mad /= img->rawImage()->norm();
                max /= img->rawImage()->norm();
                if(max>1.0f)max = 1.0f;
                float bp = median + mad * BLACK_POINT_SIGMA * MAD_TO_SIGMA;
                float mid = MTF(median - bp, TARGET_BACKGROUND);
                m_mtfParam.blackPoint[i] = bp;
                m_mtfParam.midPoint[i] = mid;
                m_mtfParam.whitePoint[i] = max;
                bp2 += bp;
                mid2 += mid;
                max2 += max;
            }
            if(m_stack->currentIndex() == 0)
            {
                m_mtfParam.blackPoint[0] = m_mtfParam.blackPoint[1] = m_mtfParam.blackPoint[2] = bp2 / ch;
                m_mtfParam.midPoint[0] = m_mtfParam.midPoint[1] = m_mtfParam.midPoint[2] = mid2 / ch;
                m_mtfParam.whitePoint[0] = m_mtfParam.whitePoint[1] = m_mtfParam.whitePoint[2] = max2 / ch;
                m_stfSlider->setMTFParams(m_mtfParam.blackPoint[0], m_mtfParam.midPoint[0], m_mtfParam.whitePoint[0]);
            }
            else
            {
                m_stfSliderR->setMTFParams(m_mtfParam.blackPoint[0], m_mtfParam.midPoint[0], m_mtfParam.whitePoint[0]);
                m_stfSliderG->setMTFParams(m_mtfParam.blackPoint[1], m_mtfParam.midPoint[1], m_mtfParam.whitePoint[1]);
                m_stfSliderB->setMTFParams(m_mtfParam.blackPoint[2], m_mtfParam.midPoint[2], m_mtfParam.whitePoint[2]);
            }
            emit paramChanged(m_mtfParam);
        }
    }
}

void StretchToolbar::resetMTF()
{
    MTFParam params;
    m_mtfParam = params;
    m_stfSlider->setMTFParams(0, 0.5, 1);
    m_stfSliderR->setMTFParams(0, 0.5, 1);
    m_stfSliderG->setMTFParams(0, 0.5, 1);
    m_stfSliderB->setMTFParams(0, 0.5, 1);
    emit paramChanged(params);
}

void StretchToolbar::imageLoaded(Image *img)
{
    if(m_autoStretchOnLoad->isChecked())
        stretchImage(img);
}