#include "imagescrollareagl.h"
#include <QOpenGLFunctions>
#include <QDebug>
#include <QKeyEvent>
#include <QOpenGLDebugLogger>
#include <QOpenGLPixelTransferOptions>
#include <QOpenGLFramebufferObject>
#include <QGridLayout>
#include <QMimeData>
#include <QMessageBox>
#include <QCoreApplication>
#include <QPainter>

struct RawImageType
{
    QOpenGLTexture::PixelFormat pixelFormat;
    QOpenGLTexture::TextureFormat textureFormat;
    QOpenGLTexture::PixelType dataType;
    bool bw;
};

const RawImageType rawImageTypes[] = {
    {QOpenGLTexture::Red, QOpenGLTexture::R8_UNorm, QOpenGLTexture::UInt8, true},
    {QOpenGLTexture::Red, QOpenGLTexture::R16_UNorm, QOpenGLTexture::UInt16, true},
    {QOpenGLTexture::Red, QOpenGLTexture::R32F, QOpenGLTexture::Float32, true},
    {QOpenGLTexture::RGB, QOpenGLTexture::RGB8_UNorm, QOpenGLTexture::UInt8, false},
    {QOpenGLTexture::BGRA,QOpenGLTexture::RGB8_UNorm, QOpenGLTexture::UInt8, false},
    {QOpenGLTexture::RGB, QOpenGLTexture::RGB16_UNorm, QOpenGLTexture::UInt16, false},
    {QOpenGLTexture::RGB, QOpenGLTexture::RGB32F, QOpenGLTexture::Float32, false}
};

void setScrollRange(QScrollBar *scrollBar, int newRange)
{
    int page = scrollBar->pageStep();
    int pos = scrollBar->value() + page/2;
    int range = scrollBar->maximum() + page;
    float relPos = (float)pos/(float)range;

    if(page >= newRange)
        scrollBar->hide();
    else
        scrollBar->show();

    scrollBar->setRange(0, newRange - page);
    scrollBar->setValue(relPos*newRange - page/2);
}

ImageWidget::ImageWidget(QWidget *parent) : QOpenGLWidget(parent)
{
    setFocusPolicy(Qt::ClickFocus);
    m_range = UINT16_MAX;
    m_low = 0;
    m_mid = 0.5;
    m_high = 1;
    m_dx = m_dy = 0;
    m_scale = 1.0f;
    m_blockRepaint = false;
    m_range = UINT16_MAX;
    m_imgWidth = m_imgHeight = -1;
    m_superpixel = m_invert = false;
    m_showThumbnails = false;
    m_thumbnailCount = 0;
    m_updateTimer = new QTimer(this);
    m_updateTimer->setInterval(500);
    m_updateTimer->setSingleShot(true);
    connect(m_updateTimer, SIGNAL(timeout()), this, SLOT(update()));
    setAcceptDrops(true);
    QTimer::singleShot(1000, [this](){
        if(!isValid())
        {
            QMessageBox::critical(this, tr("OpenGL error"), tr("Could not initialize OpenGL 3.3 context. Ensure that proper GPU driver is installed."));
             QCoreApplication::exit(-1);
        }
    });
}

ImageWidget::~ImageWidget()
{
    makeCurrent();
}

void ImageWidget::setImage(const RawImage *image, int index)
{
    if(image == nullptr)return;

    m_imgWidth = image->width();
    m_imgHeight = image->height();
    m_currentImg = index;

    const RawImageType &rawImageType = rawImageTypes[image->type()];

    m_image->destroy();
    m_image->setFormat(rawImageType.textureFormat);
    m_image->setSize(image->width(), image->height());
    m_image->setMipLevels([&](){ int c = 0; int s = std::min(m_imgWidth, m_imgHeight); while(s>>=1)c++; return c; }());
    m_image->allocateStorage();
    m_image->setMinMagFilters(QOpenGLTexture::LinearMipMapLinear, QOpenGLTexture::Linear);
    m_image->setWrapMode(QOpenGLTexture::ClampToEdge);
    m_image->setBorderColor(0, 0, 0, 0);
    m_image->setData(0, rawImageType.pixelFormat, rawImageType.dataType, image->data(), m_transferOptions.get());
    m_image->setLevelOfDetailRange(m_superpixel ? 1 : 0, m_image->mipMaxLevel());
    m_image->generateMipMaps();
    m_bwImg = rawImageType.bw;
    update();
}

void ImageWidget::setScale(float scale)
{
    m_scale = scale;
    update();
}

void ImageWidget::blockRepaint(bool block)
{
    m_blockRepaint = block;
    if(!block)update();
}

void ImageWidget::allocateThumbnails(const QStringList &&names)
{
    int count = names.size();
    m_Names = std::move(names);
    m_thumbnailTexture->destroy();
    m_thumbnailTexture->create();
    m_thumbnailTexture->setFormat(QOpenGLTexture::RGB16_UNorm);
    m_thumbnailTexture->setSize(THUMB_SIZE, THUMB_SIZE);
    m_thumbnailTexture->setLayers(names.size());
    m_thumbnailTexture->allocateStorage();
    m_bufferSizes->bind();
    float *tmp = new float[count*3];
    memset(tmp, 0, count * sizeof(float)*3);
    m_bufferSizes->allocate(tmp, count * sizeof(float)*3);
    delete [] tmp;
    m_thumbnailCount = count;
}

void ImageWidget::setMTFParams(float low, float mid, float high)
{
    m_low = low;
    m_mid = mid;
    m_high = high;
    update();
}

void ImageWidget::setOffset(int dx, int dy)
{
    m_dx = dx;
    m_dy = dy;
    update();
}

void ImageWidget::superPixel(bool enable)
{
    m_superpixel = enable;
    m_image->setLevelOfDetailRange(enable ? 1 : 0, m_image->mipMaxLevel());
    update();
}

void ImageWidget::invert(bool enable)
{
    m_invert = enable;
    update();
}

QImage ImageWidget::renderToImage()
{
    if(m_imgWidth < 0)return QImage();
    makeCurrent();
    QOpenGLFramebufferObject fbo(m_imgWidth, m_imgHeight);
    fbo.bind();

    f->glViewport(0, 0, m_imgWidth, m_imgHeight);

    m_program->bind();
    m_program->setUniformValue("viewport", (float)m_imgWidth, (float)m_imgHeight);
    m_program->setUniformValue("offset", 0.0f, 0.0f);
    m_program->setUniformValue("zoom", 1.0f);

    m_image->bind(0);
    f->glDrawArrays(GL_TRIANGLE_STRIP, 0, 4);

    fbo.bindDefault();
    return fbo.toImage(true);
}

void ImageWidget::thumbnailLoaded(const Image *image)
{
    const RawImage *raw = image->thumbnail();
    m_thumbnailTexture->setData(0, image->number(), QOpenGLTexture::RGB, QOpenGLTexture::UInt16, raw->data(), m_transferOptions.get());
    float a = raw->thumbAspect();
    int sizes[3] = { std::max(1, a > 1.0f ? THUMB_SIZE : (int)(THUMB_SIZE * a)), std::max(1, a < 1.0f ? THUMB_SIZE : (int)(THUMB_SIZE / a)), image->number() };
    m_bufferSizes->bind();
    m_bufferSizes->write(image->number() * sizeof(sizes), sizes, sizeof(sizes));
    if(!m_updateTimer->isActive())m_updateTimer->start();
}

void ImageWidget::showThumbnail(bool enable)
{
    m_showThumbnails = enable;
    update();
}

void ImageWidget::paintGL()
{
    if(m_blockRepaint)return;

    float dx = m_dx;
    float dy = m_dy;
    if(width() > m_image->width()*m_scale)
        dx = -width()*0.5f + m_image->width()*m_scale*0.5f;
    if(height() > m_image->height()*m_scale)
        dy = -height()*0.5f + m_image->height()*m_scale*0.5f;

    if(m_showThumbnails)
    {
        m_vaoThumb->bind();
        m_thumbnailTexture->bind(1);
        m_thumbnailProgram->bind();
        f->glUniform3i(m_thumbnailProgram->uniformLocation("viewport_row"), width(), height(), width()/THUMB_SIZE_BORDER);
        m_thumbnailProgram->setUniformValue("mtf_param", m_low, m_mid, m_high);
        m_thumbnailProgram->setUniformValue("invert", m_invert);
        m_thumbnailProgram->setUniformValue("offset", 0, m_dy);
        f3->glVertexAttribDivisor(m_thumbnailProgram->attributeLocation("imageSize_num"), 1);
        QMatrix4x4 mvp;
        mvp.ortho(rect());
        m_thumbnailProgram->setUniformValue("mvp", mvp);
        if(f3)f3->glDrawArraysInstanced(GL_TRIANGLE_STRIP, 0, 4, m_thumbnailCount);

        QPainter painter(this);
        int w = width()/THUMB_SIZE_BORDER;
        for(int i=0; i < m_thumbnailCount; i++)
        {
            float x = (i % w) * THUMB_SIZE_BORDER;
            float y = i / w * THUMB_SIZE_BORDER_Y + THUMB_SIZE - m_dy;
            QRectF rect(x, y, THUMB_SIZE_BORDER, 32);
            painter.drawText(rect, Qt::AlignCenter | Qt::TextWrapAnywhere, QString(m_Names.at(i)));
            if(m_currentImg == i)
            {
                int off = (THUMB_SIZE_BORDER - THUMB_SIZE) / 2;
                painter.save();
                painter.setPen(QPen(Qt::red, 2.0));
                painter.drawRect((i % w) * THUMB_SIZE_BORDER + off, i / w * THUMB_SIZE_BORDER_Y - m_dy + off, THUMB_SIZE, THUMB_SIZE);
                painter.restore();
            }
        }
    }
    else
    {
        m_vao->bind();
        m_image->bind(0);
        m_program->bind();
        m_program->setUniformValue("viewport", (float)width(), (float)height());
        m_program->setUniformValue("offset", dx, dy);
        m_program->setUniformValue("mtf_param", m_low, m_mid, m_high);
        m_program->setUniformValue("zoom", 1.0f/m_scale);
        m_program->setUniformValue("bw", m_bwImg);
        m_program->setUniformValue("invert", m_invert);
        f->glDrawArrays(GL_TRIANGLE_STRIP, 0, 4);
    }

}

void ImageWidget::resizeGL(int w, int h)
{
    m_width = w;
    m_height = h;
    f->glViewport(0, 0, w, h);
}

void ImageWidget::initializeGL()
{
    f = context()->functions();
    f->glClearColor(0.5f, 0.5f, 0.5f, 1.0f);
    f3 = context()->versionFunctions<QOpenGLFunctions_3_3_Core>();

    if(f3 == nullptr)
        QMessageBox::critical(this, tr("OpenGL error"), tr("Could not initialize OpenGL 3.3 context. Ensure that proper GPU driver is installed."));

    m_vao = std::unique_ptr<QOpenGLVertexArrayObject>(new QOpenGLVertexArrayObject);
    m_vaoThumb = std::unique_ptr<QOpenGLVertexArrayObject>(new QOpenGLVertexArrayObject);
    m_vao->create();
    m_vaoThumb->create();
    m_vao->bind();

    QOpenGLDebugLogger *logger = new QOpenGLDebugLogger(this);
    logger->initialize();
    logger->startLogging();
    connect(logger, &QOpenGLDebugLogger::messageLogged, [](const QOpenGLDebugMessage &message)
        {
            qDebug() << message;
        });

    qDebug() << (char*)f->glGetString(GL_VENDOR);
    qDebug() << (char*)f->glGetString(GL_RENDERER);
    qDebug() << (char*)f->glGetString(GL_VERSION);

    qDebug() << context()->format();

    // each vertex is  x,y 2D position and s,t texture coordinates
    float vertexs[] = {-1.0f, -1.0f, 0.0f, 1.0f,
                        1.0f, -1.0f, 1.0f, 1.0f,
                       -1.0f,  1.0f, 0.0f, 0.0f,
                        1.0f,  1.0f, 1.0f, 0.0f,};
    m_buffer = std::unique_ptr<QOpenGLBuffer>(new QOpenGLBuffer(QOpenGLBuffer::VertexBuffer));
    m_buffer->setUsagePattern(QOpenGLBuffer::StaticDraw);
    m_buffer->create();
    m_buffer->bind();
    m_buffer->allocate(vertexs, sizeof(vertexs));
    // f->glVertexAttribPointer(0, 2, GL_FLOAT, false, sizeof(float)*4, 0);

    m_program = std::unique_ptr<QOpenGLShaderProgram>(new QOpenGLShaderProgram);
    m_program->addShaderFromSourceFile(QOpenGLShader::Vertex, ":/shaders/image.vert");
    m_program->addShaderFromSourceFile(QOpenGLShader::Fragment, ":/shaders/image.frag");
    if(f3)f3->glBindFragDataLocation(m_program->programId(), 0, "color");
    if(!m_program->link())
    {
        qDebug() << "Link failed" << m_program->log();
    }

    m_program->bind();
    m_program->enableAttributeArray("qt_Vertex");
    m_program->setAttributeBuffer("qt_Vertex", GL_FLOAT, 0, 2, sizeof(float)*4);
    m_program->enableAttributeArray("qt_MultiTexCoord0");
    m_program->setAttributeBuffer("qt_MultiTexCoord0", GL_FLOAT, sizeof(float)*2, 2, sizeof(float)*4);
    m_program->setUniformValue("qt_Texture0", (GLuint)0);
    m_program->setUniformValue("scale", 1.0f, 0.0f);

    m_vaoThumb->bind();

    m_thumbnailProgram = std::unique_ptr<QOpenGLShaderProgram>(new QOpenGLShaderProgram);
    m_thumbnailProgram->addShaderFromSourceFile(QOpenGLShader::Vertex, ":/shaders/thumb.vert");
    m_thumbnailProgram->addShaderFromSourceFile(QOpenGLShader::Fragment, ":/shaders/thumb.frag");
    if(f3)f3->glBindFragDataLocation(m_program->programId(), 0, "color");
    m_thumbnailProgram->bind();
    m_thumbnailProgram->enableAttributeArray("qt_Vertex");
    m_thumbnailProgram->setAttributeBuffer("qt_Vertex", GL_FLOAT, 0, 2, sizeof(float)*4);
    m_thumbnailProgram->enableAttributeArray("qt_MultiTexCoord0");
    m_thumbnailProgram->setAttributeBuffer("qt_MultiTexCoord0", GL_FLOAT, sizeof(float)*2, 2, sizeof(float)*4);
    if(!m_thumbnailProgram->link())
    {
        qDebug() << "Link failed" << m_thumbnailProgram->log();
    }
    m_thumbnailProgram->setUniformValue("qt_Texture0", (GLuint)1);

    m_bufferSizes = std::unique_ptr<QOpenGLBuffer>(new QOpenGLBuffer(QOpenGLBuffer::VertexBuffer));
    m_bufferSizes->setUsagePattern(QOpenGLBuffer::StaticDraw);
    m_bufferSizes->create();
    m_bufferSizes->bind();

    m_bufferSizes->allocate(12);
    m_thumbnailProgram->enableAttributeArray("imageSize_num");
    m_thumbnailProgram->setAttributeBuffer("imageSize_num", GL_FLOAT, 0, 3);
    f3->glVertexAttribDivisor(m_thumbnailProgram->attributeLocation("imageSize_num"), 1);

    m_image = std::unique_ptr<QOpenGLTexture>(new QOpenGLTexture(QOpenGLTexture::Target2D));
    m_image->setFormat(QOpenGLTexture::RGB8U);
    m_image->allocateStorage();
    m_image->bind(0);
    m_image->setMinificationFilter(QOpenGLTexture::LinearMipMapLinear);
    m_image->setMagnificationFilter(QOpenGLTexture::Linear);

    m_thumbnailTexture = std::unique_ptr<QOpenGLTexture>(new QOpenGLTexture(QOpenGLTexture::Target2DArray));
    m_thumbnailTexture->setFormat(QOpenGLTexture::RGB16_UNorm);
    m_thumbnailTexture->setSize(THUMB_SIZE, THUMB_SIZE);
    m_thumbnailTexture->setLayers(1);
    m_thumbnailTexture->allocateStorage();
    m_thumbnailTexture->bind(1);
    m_thumbnailTexture->setMinificationFilter(QOpenGLTexture::Linear);
    m_thumbnailTexture->setMinificationFilter(QOpenGLTexture::Linear);

    m_transferOptions = std::unique_ptr<QOpenGLPixelTransferOptions>(new QOpenGLPixelTransferOptions);
    m_transferOptions->setAlignment(1);
}

void ImageWidget::dragEnterEvent(QDragEnterEvent *event)
{
    if(event->mimeData()->hasUrls() && event->proposedAction() & (Qt::CopyAction | Qt::MoveAction))
        event->acceptProposedAction();
}

void ImageWidget::dropEvent(QDropEvent *event)
{
    if(event->mimeData()->hasUrls() && event->proposedAction() & (Qt::CopyAction | Qt::MoveAction))
    {
        for(const QUrl &url : event->mimeData()->urls())
        {
            if(url.isLocalFile())
            {
                emit fileDropped(url.toLocalFile());
                event->accept();
                return;
            }
        }
    }
    event->ignore();
}

ImageScrollAreaGL::ImageScrollAreaGL(QWidget *parent) : QWidget(parent)
{
    QGridLayout *layout = new QGridLayout(this);
    setLayout(layout);

    m_imageWidget = new ImageWidget(this);

    m_verticalScrollBar = new QScrollBar(Qt::Vertical, this);
    m_horizontalScrollBar = new QScrollBar(Qt::Horizontal, this);
    m_scale = 1.0f;
    m_bestFit = false;
    m_thumbCount = 0;

    layout->setSpacing(0);
    layout->addWidget(m_imageWidget, 0, 0);
    layout->addWidget(m_verticalScrollBar, 0, 1);
    layout->addWidget(m_horizontalScrollBar, 1, 0);

    connect(m_verticalScrollBar, SIGNAL(valueChanged(int)), this, SLOT(scrollEvent()));
    connect(m_horizontalScrollBar, SIGNAL(valueChanged(int)), this, SLOT(scrollEvent()));
}

ImageScrollAreaGL::~ImageScrollAreaGL()
{

}

void ImageScrollAreaGL::setImage(RawImage *image, int index)
{
    m_imageWidget->setImage(image, index);
    m_imgWidth = image->width();
    m_imgHeight = image->height();
    if(m_bestFit)bestFit();
    updateScrollbars();
}

ImageWidget *ImageScrollAreaGL::imageWidget()
{
    return m_imageWidget;
}

void ImageScrollAreaGL::setThumbnails(int count)
{
    m_thumbCount = count;
    if(m_thumbCount)
    {
        m_verticalScrollBar->setRange(0, m_thumbCount / (m_imageWidget->width() / THUMB_SIZE_BORDER) * THUMB_SIZE_BORDER_Y);
        m_verticalScrollBar->setPageStep(THUMB_SIZE_BORDER_Y);
    }
    else
    {
        m_verticalScrollBar->setPageStep(m_imageWidget->height());
        m_horizontalScrollBar->setPageStep(m_imageWidget->width());
    }
    updateScrollbars();
}

void ImageScrollAreaGL::updateScrollbars(bool zoom)
{
    if(m_thumbCount)
    {
        m_horizontalScrollBar->hide();
        m_verticalScrollBar->show();
        m_verticalScrollBar->setRange(0, m_thumbCount / (m_imageWidget->width() / THUMB_SIZE_BORDER) * THUMB_SIZE_BORDER_Y);
        m_verticalScrollBar->setPageStep(THUMB_SIZE_BORDER_Y);
    }
    else
    {
        if(zoom)
        {
            setScrollRange(m_verticalScrollBar, m_imgHeight*m_scale);
            setScrollRange(m_horizontalScrollBar, m_imgWidth*m_scale);
        }
        else
        {
            m_verticalScrollBar->setRange(0, m_imgHeight*m_scale - m_verticalScrollBar->pageStep());
            m_horizontalScrollBar->setRange(0, m_imgWidth*m_scale - m_horizontalScrollBar->pageStep());
        }
    }
}

void ImageScrollAreaGL::resizeEvent(QResizeEvent *event)
{
    QWidget::resizeEvent(event);
    if(m_thumbCount)
    {
        m_verticalScrollBar->setRange(0, m_thumbCount / (m_imageWidget->width() / THUMB_SIZE_BORDER) * THUMB_SIZE_BORDER_Y);
        m_verticalScrollBar->setPageStep(THUMB_SIZE_BORDER_Y);
    }
    else
    {
        m_verticalScrollBar->setPageStep(m_imageWidget->height());
        m_horizontalScrollBar->setPageStep(m_imageWidget->width());
    }
    updateScrollbars();
}

void ImageScrollAreaGL::mouseMoveEvent(QMouseEvent *event)
{
    QPoint delta = m_lastPos - event->pos();
    if(m_thumbCount == 0)m_horizontalScrollBar->setValue(m_horizontalScrollBar->value() + delta.x());
    m_verticalScrollBar->setValue(m_verticalScrollBar->value() + delta.y());
    m_lastPos = event->pos();
}

void ImageScrollAreaGL::mousePressEvent(QMouseEvent *event)
{
    m_lastPos = event->pos();
}

void ImageScrollAreaGL::wheelEvent(QWheelEvent *event)
{
    if(m_thumbCount)
    {
        m_verticalScrollBar->setValue(m_verticalScrollBar->value() - event->angleDelta().y());
    }
    else
    {
        m_bestFit = false;
        if(event->angleDelta().y() != 0)
            zoom(event->angleDelta().y() / 1200.0f);
    }
}

void ImageScrollAreaGL::zoom(float delta)
{
    if((m_scale >= 8.0f && delta > 0) || (m_scale <= 0.1f && delta < 0))return;

    m_scale += delta;
    m_imageWidget->blockRepaint(true);
    m_imageWidget->setScale(m_scale);
    updateScrollbars(true);
    m_imageWidget->blockRepaint(false);
}

void ImageScrollAreaGL::zoomIn()
{
    zoom(0.1f);
    m_bestFit = false;
}

void ImageScrollAreaGL::zoomOut()
{
    zoom(-0.1f);
    m_bestFit = false;
}

void ImageScrollAreaGL::bestFit()
{
    m_bestFit = true;
    m_scale = std::min((float)m_imageWidget->width()/m_imgWidth, (float)m_imageWidget->height()/m_imgHeight);
    zoom(0.0f);
}

void ImageScrollAreaGL::oneToOne()
{
    m_scale = 1.0f;
    zoom(0.0f);
    m_bestFit = false;
}

void ImageScrollAreaGL::scrollEvent()
{
    m_imageWidget->setOffset(m_horizontalScrollBar->value(), m_verticalScrollBar->value());
}