#include "mainwindow.h"
#include "./ui_mainwindow.h"
#include <QFileDialog>
#include <QMessageBox>
#include <QtConcurrent/QtConcurrent>
#include <QPromise>
#include "laplacian.h"

MainWindow::MainWindow(QWidget *parent) : QMainWindow(parent)
    , ui(new Ui::MainWindow)
{
    ui->setupUi(this);

    connect(ui->actionOpen, &QAction::triggered, this, &MainWindow::openFile);
    connect(ui->actionSave, &QAction::triggered, this, &MainWindow::saveImage);
    connect(ui->analyzeButton, &QPushButton::clicked, this, &MainWindow::analyze);
    connect(ui->stackButton, &QPushButton::clicked, this, &MainWindow::stack);
    connect(ui->cancelButton, &QPushButton::clicked, this, [this](){
        if(_stack.isRunning())_stack.cancel();
        if(_analyze.isRunning())_analyze.cancel();
    });

    ui->progressBar->hide();

    connect(&_watcherAnalyze, &QFutureWatcher<void>::progressValueChanged, [this](int val){ ui->progressBar->setValue(val); });
    connect(&_watcherAnalyze, &QFutureWatcher<void>::finished, [this](){ ui->progressBar->hide(); ui->stackButton->setEnabled(true); });
    connect(&_watcherAnalyze, &QFutureWatcher<void>::canceled, [this](){ ui->progressBar->hide(); });
    connect(&_watcherAnalyze, &QFutureWatcher<void>::started, [this](){ ui->progressBar->show(); });
    connect(&_watcherStack, &QFutureWatcher<void>::progressValueChanged, [this](int val){ ui->progressBar->setValue(val); });
    connect(&_watcherStack, &QFutureWatcher<void>::finished, [this](){ ui->progressBar->hide(); stackFinish(); });
    connect(&_watcherStack, &QFutureWatcher<void>::canceled, [this](){ ui->progressBar->hide(); });
    connect(&_watcherStack, &QFutureWatcher<void>::started, [this](){ ui->progressBar->show(); });

}

MainWindow::~MainWindow()
{
    delete ui;
}

void MainWindow::openFile()
{
    QString fileName = QFileDialog::getOpenFileName(this, tr("Open file"), QString(), tr("SER video (*.ser)"));
    _variance.clear();
    ui->stackButton->setEnabled(false);
    ui->analyzeButton->setEnabled(false);
    if(!fileName.isEmpty())
    {
        if(!_ser.open(fileName))
        {
            QMessageBox::critical(this, tr("Failed to open file"), tr("Failed to open %1").arg(fileName));
            _ser.close();
            return;
        }

        if(_ser.pixelDepth() < 9 || _ser.pixelDepth() > 16 || _ser.pixelFormat() != SERPixelFormat::MONO)
        {
            QMessageBox::critical(this, tr("Invalid format"), tr("Only 9-16 bit mono video file supported"));
            ui->stackButton->setEnabled(false);
            _ser.close();
            return;
        }

        ui->analyzeButton->setEnabled(true);
    }
}

void MainWindow::saveImage()
{
    if(_stack.isValid())
    {
        QString fileName = QFileDialog::getSaveFileName(this, tr("Save image"), QString(), tr("Image (*.png)"));
        if(!fileName.isEmpty())
        {
            cv::Mat stack = _stack.result();
            stack.convertTo(stack, CV_16U, 65355.0);
            QImage img(stack.data, stack.cols, stack.rows, stack.cols * 2, QImage::Format_Grayscale16);
            img.save(fileName, "PNG");
        }
    }
}

void MainWindow::analyze()
{
    if(_ser.frameCount() && !_analyze.isRunning() && !_stack.isRunning())
    {
        _analyze = QtConcurrent::run([this](QPromise<int> &promise) {
            promise.start();
            cv::Mat img(_ser.height(), _ser.width(), CV_16U);
            _variance.clear();
            promise.setProgressRange(0, 100);
            for(uint32_t i = 0; i < _ser.frameCount() && !promise.isCanceled(); i++)
            {
                _ser.getFrame(i, (char*)img.data);
                double var = laplacian((uint16_t*)img.data, nullptr, img.cols, img.rows);
                _variance.push_back({i, var});
                promise.setProgressValue(i * 100.0 / _ser.frameCount());
            }
            std::sort(_variance.begin(), _variance.end(), [](const std::pair<int, double> &a, const std::pair<int, double> &b){
                return a.second > b.second;
            });

            promise.finish();
        });
        _watcherAnalyze.setFuture(_analyze);
    }
}

void MainWindow::stack()
{
    if(!_variance.empty() && !_analyze.isRunning() && !_stack.isRunning())
    {
        _stack = QtConcurrent::run([this](QPromise<cv::Mat> &promise, uint32_t frames) {
            promise.start();
            cv::Mat img(_ser.height(), _ser.width(), CV_16U);
            cv::Mat ref(_ser.height(), _ser.width(), CV_16U);
            cv::Mat stack(_ser.height(), _ser.width(), CV_32F);
            _ser.getFrame(_variance[0].first, (char*)ref.data);
            stack += ref;
            cv::Mat warp;
            promise.setProgressRange(0, 100);
            for(uint32_t i = 1; i < frames && !promise.isCanceled(); i++)
            {
                _ser.getFrame(_variance[i].first, (char*)img.data);
                reflow(ref, img, warp);
                stack += warp;
                promise.setProgressValue(i * 100.0 / frames);
            }
            cv::Mat norm;
            cv::normalize(stack, norm, 0.0, 1.0, cv::NORM_MINMAX);
            promise.addResult(norm);
            promise.finish();
        }, ui->frameTop->value() / 100.0 * _ser.frameCount());
        _watcherStack.setFuture(_stack);
    }
}

void MainWindow::stackFinish()
{
    cv::Mat stack = _stack.result();
    cv::Mat scaled;
    QSize s = ui->image->size();
    cv::resize(stack, scaled, cv::Size(s.width(), s.height()));
    scaled.convertTo(scaled, CV_8U, 255.0);
    QImage img(scaled.data, s.width(), s.height(), scaled.cols, QImage::Format_Grayscale8);
    ui->image->setPixmap(QPixmap::fromImage(img));
}