#include "mainwindow.h"
#include "./ui_mainwindow.h"
#include "QFileDialog"
#include "QMessageBox"
#include <opencv2/opencv.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <iostream>
#include <string>
#include <vector>
#include <fstream>
#include "loguru.hpp"
#include "detector.h"
#include "QMessageBox"
using namespace cv;
using namespace std;

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

    connect(ui->inputBtn, SIGNAL(clicked()), this, SLOT(selectImgVid()));
    connect(ui->modelBtn, SIGNAL(clicked()), this, SLOT(selectModel()));
    connect(ui->saveBtn, SIGNAL(clicked()), this, SLOT(selectSavePath()));
    connect(ui->okBtn, SIGNAL(clicked()), this, SLOT(objectDetection()));
}

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

bool isImg(QString s){
    QStringList imgSuffix = {".png", ".jpg", ".jpeg", ".bmp"};
    for(int i=0; i<imgSuffix.size(); i++){
        if (s.endsWith(imgSuffix.at(i)))
            return true;
    }
    return false;
}

bool isVid(QString s){
    QStringList imgSuffix = {".mp4", ".avi"};
    for(int i=0; i<imgSuffix.size(); i++){
        if (s.endsWith(imgSuffix.at(i)))
            return true;
    }
    return false;
}

static QImage cvMat2Qimage(const cv::Mat &mat ) {
  switch ( mat.type() )
  {
     // 8-bit, 4 channel
     case CV_8UC4:
     {
        QImage image( mat.data, mat.cols, mat.rows, mat.step, QImage::Format_RGB32 );
        return image;
     }

     // 8-bit, 3 channel
     case CV_8UC3:
     {
        QImage image( mat.data, mat.cols, mat.rows, mat.step, QImage::Format_RGB888 );
        return image.rgbSwapped();
     }

     // 8-bit, 1 channel
     case CV_8UC1:
     {
        static QVector<QRgb>  sColorTable;
        // only create our color table once
        if ( sColorTable.isEmpty() )
        {
           for ( int i = 0; i < 256; ++i )
              sColorTable.push_back( qRgb( i, i, i ) );
        }
        QImage image( mat.data, mat.cols, mat.rows, mat.step, QImage::Format_Indexed8 );
        image.setColorTable( sColorTable );
        return image;
     }

     default:
        qDebug("Image format is not supported: depth=%d and %d channels\n", mat.depth(), mat.channels());
        break;
  }
  return QImage();
}

void MainWindow::selectImgVid(){
    path = QFileDialog::getOpenFileName(
                this,
                "Select an image or a video.",
                "./images",
                "images(*.png *.jpg *jpeg *bmp);;videos(*.mp4 *.avi)");
    if (!path.isEmpty()) {
        // show input image file name
        ui->inputPathEdit->setText(path);
        // image or video
        if (isImg(path)){
            // show input image
            QPixmap pixmap(path);
            pixmap = pixmap.scaled(ui->inputViewLabel->size(), Qt::KeepAspectRatio);
            ui->inputViewLabel->setPixmap(pixmap);
            ui->inputViewLabel->show();
        }
        if (isVid(path)){
            // show first frame of video
            QPixmap pixmap(path);
            // open video file
            cap = VideoCapture(ui->inputPathEdit->text().toStdString());
            if (!cap.isOpened()){
                QMessageBox::critical(this, "error", "Invalid video file!");
                return;
            }
            if (cap.read(frame)){
                isShown = false;
                ui->inputViewLabel->setPixmap(
                            QPixmap::fromImage(
                                cvMat2Qimage(frame)).scaled(ui->inputViewLabel->size(),
                                Qt::KeepAspectRatio));
                return;
            }
            else{
                QMessageBox::critical(this, "error", "Can not read frames from video!");
                return;
            }
        }
    }
}

void MainWindow::selectModel(){
    modelPath = QFileDialog::getOpenFileName(
                this,
                "Select an model.",
                "./models",
                "*.onnx");
    if (!modelPath.isEmpty()) {
        // show input image file name
        ui->modelPathEdit->setText(modelPath);
    }
}

void MainWindow::selectSavePath(){
    savePath = QFileDialog::getSaveFileName(
                this,
                "Select a location to save output file.",
                "./runs",
                "images(*.png *.jpg *jpeg *bmp)");
    if (!savePath.isEmpty()){
        // show save path
        ui->savePathEdit->setText(savePath);
    }
}

void MainWindow::objectDetection(){
    Config config = {0.25f, 0.45f, modelPath.toStdString(), "./classes.names", Size(640, 640),false};
    Detector detector(config);
    if (isImg(path)){
        // object detection on image
        Mat img = imread(path.toStdString(), IMREAD_COLOR);
        Detection detection = detector.detect(img);
        Colors cl = Colors();
        String summary = detector.postProcess(img, detection,cl);


        Mat imgColorConverted;
        cv::cvtColor(img, imgColorConverted, COLOR_BGR2RGB);
        QImage qimg = QImage(
                    (unsigned char *)imgColorConverted.data,
                    imgColorConverted.cols,
                    imgColorConverted.rows,
                    imgColorConverted.step,
                    QImage::Format_RGB888);
        QPixmap qpix = QPixmap::fromImage(qimg).scaled(ui->outputViewLabel->size(), Qt::KeepAspectRatio);
        ui->outputViewLabel->setPixmap(qpix);
        ui->statusBrowser->setText(QString::fromStdString(summary));
        if (!savePath.isEmpty()){
            imwrite(savePath.toStdString(), img);
        }
        return;
    }
    if (isVid(path)){
        // object detection on video
        while (true){
            // show origin video frame
            ui->inputViewLabel->setPixmap(
                        QPixmap::fromImage(
                            cvMat2Qimage(frame)).scaled(ui->inputViewLabel->size(),
                                         Qt::KeepAspectRatio));

            // show labeled frame
            Detection detection = detector.detect(frame);
            Colors cl = Colors();
            String summary = detector.postProcess(frame, detection,cl);

            Mat imgColorConverted;
            cv::cvtColor(frame, imgColorConverted, COLOR_BGR2RGB);
            QImage qimg = QImage(
                        (unsigned char *)imgColorConverted.data,
                        imgColorConverted.cols,
                        imgColorConverted.rows,
                        imgColorConverted.step,
                        QImage::Format_RGB888);
            QPixmap qpix = QPixmap::fromImage(qimg).scaled(ui->outputViewLabel->size(), Qt::KeepAspectRatio);
            ui->outputViewLabel->setPixmap(qpix);
            ui->statusBrowser->setText(QString::fromStdString(summary));

            waitKey(1);

            // get next frame
            if (!cap.read(frame)){
                break;
            }
        }
    }
}