/*
Author: Vargha Csongor Csaba
Created: 2023-06-25 10:10:13

Description:
    This file contains the main function for the TV image denoising cli tool.
    It reads an image file, denoises it using the TV denoising algorithm,
    and saves the denoised image to a file.
    You can run it with the following command:
        ./TV_Denoising_CUDA <input_image> <output_image> <lambda> <iterations>
    where:
        - <input_image> is the path to the input image file you want to denoise.
        - <output_image> is the path to the output denoised image file.
        - <lambda> is the regularization parameter for TV denoising (optional, default: 0.02).
        - <iterations> is the number of iterations for TV denoising (optional, default: 10).
*/
#include <iostream>
#include <QApplication>
#include <QVBoxLayout>
#include <QHBoxLayout>
#include <QLabel>
#include <QDebug>
#include <QLineEdit>
#include <QPushButton>
#include <QSlider>
#include <QFileDialog>
#include <QMessageBox>
#include <QImageReader>
#include <QImage>
#include <opencv2/opencv.hpp>
#include "tv_denoising.hpp"

// Global variables for GUI elements
QLabel *originalImageLabel;
QLabel *resultImageLabel;
QLineEdit *inputImagePathEdit;
QLineEdit *outputImagePathEdit;
QSlider *lambdaSlider;
QSlider *iterationsSlider;

void updateDenoisedImage()
{
    std::string inputImagePath = inputImagePathEdit->text().toStdString();
    std::string outputImagePath = outputImagePathEdit->text().toStdString();
    float lambda = lambdaSlider->value() / 100.0;
    int iterations = iterationsSlider->value();

    // Read the input image
    cv::Mat image = cv::imread(inputImagePath, cv::IMREAD_GRAYSCALE);

    // Check if the image was successfully loaded
    if (image.empty())
    {
        QMessageBox::critical(nullptr, "Error", "Failed to read the input image.");
        return;
    }

    // Perform TV denoising
    TVDenoising(image, lambda, iterations);

    // Display the denoised image
    QImage resultImage(image.data, image.cols, image.rows, QImage::Format_Grayscale8);
    resultImageLabel->setPixmap(QPixmap::fromImage(resultImage));

    // Save the denoised image
    cv::imwrite(outputImagePath, image);
}

void openInputImage()
{
    QString filePath = QFileDialog::getOpenFileName(nullptr, "Select Input Image");
    if (!filePath.isEmpty())
        inputImagePathEdit->setText(filePath);
}

void openOutputImage()
{
    QString filePath = QFileDialog::getSaveFileName(nullptr, "Select Output Image");
    if (!filePath.isEmpty())
        outputImagePathEdit->setText(filePath);
}

int main(int argc, char **argv)
{
    QApplication app(argc, argv);

    // Create the main window
    QWidget window;
    window.setWindowTitle("TV Image Denoising");
    window.resize(800, 600);

    // Create the layout
    QGridLayout *layout = new QGridLayout(&window);

    // Create the original image label
    QLabel *originalImageLabel = new QLabel("Original Image");
    layout->addWidget(originalImageLabel, 0, 0, 1, 2); // Row 0, spanning 1 row, 2 columns

    // Create the original image display widget
    QLabel *originalImageDisplay = new QLabel();
    originalImageDisplay->setAlignment(Qt::AlignCenter);
    layout->addWidget(originalImageDisplay, 1, 0, 1, 2); // Row 1, spanning 1 row, 2 columns

    // Create the result image label
    QLabel *resultImageLabel = new QLabel("Result Image");
    layout->addWidget(resultImageLabel, 0, 2, 1, 2); // Row 0, spanning 1 row, 2 columns

    // Create the result image display widget
    QLabel *resultImageDisplay = new QLabel();
    resultImageDisplay->setAlignment(Qt::AlignCenter);
    layout->addWidget(resultImageDisplay, 1, 2, 1, 2); // Row 1, spanning 1 row, 2 columns

    // Create the input image path input field
    QHBoxLayout *inputLayout = new QHBoxLayout;
    QLabel *inputLabel = new QLabel("Input Image:");
    QLineEdit *inputImagePathEdit = new QLineEdit;
    QPushButton *inputBrowseButton = new QPushButton("Browse");
    QObject::connect(inputBrowseButton, &QPushButton::clicked, openInputImage);
    inputLayout->addWidget(inputLabel);
    inputLayout->addWidget(inputImagePathEdit);
    inputLayout->addWidget(inputBrowseButton);
    layout->addLayout(inputLayout, 2, 0, 1, 4); // Row 2, spanning 1 row, 4 columns

    // Create the output image path input field
    QHBoxLayout *outputLayout = new QHBoxLayout;
    QLabel *outputLabel = new QLabel("Output Image:");
    QLineEdit *outputImagePathEdit = new QLineEdit;
    QPushButton *outputBrowseButton = new QPushButton("Browse");
    QObject::connect(outputBrowseButton, &QPushButton::clicked, openOutputImage);
    outputLayout->addWidget(outputLabel);
    outputLayout->addWidget(outputImagePathEdit);
    outputLayout->addWidget(outputBrowseButton);
    layout->addLayout(outputLayout, 3, 0, 1, 4); // Row 3, spanning 1 row, 4 columns

    // Create the lambda slider
    QLabel *lambdaLabel = new QLabel("Lambda:");
    QSlider *lambdaSlider = new QSlider(Qt::Horizontal);
    lambdaSlider->setObjectName("LambdaSlider"); // Set object name for styling
    lambdaSlider->setMinimum(0);
    lambdaSlider->setMaximum(100);

    // Create the iterations slider
    QLabel *iterationsLabel = new QLabel("Iterations:");
    QSlider *iterationsSlider = new QSlider(Qt::Horizontal);
    iterationsSlider->setObjectName("IterationsSlider"); // Set object name for styling
    iterationsSlider->setMinimum(1);
    iterationsSlider->setMaximum(100);

    // Create labels to display slider values
    QLabel *lambdaValueLabel = new QLabel();
    QLabel *iterationsValueLabel = new QLabel();

    // Set stylesheet for dark design
    QString styleSheet = "QLabel { color: white; }"                                                                // Set label text color to white
                         "QSlider::groove:horizontal { background-color: #555555; height: 6px; }"                  // Set groove color and height
                         "QSlider::handle:horizontal { background-color: #ffffff; width: 10px; margin: -6px 0; }"; // Set handle color and size

    lambdaSlider->setStyleSheet(styleSheet);
    iterationsSlider->setStyleSheet(styleSheet);

    // Set size policies for sliders
    lambdaSlider->setSizePolicy(QSizePolicy::Expanding, QSizePolicy::Preferred);
    iterationsSlider->setSizePolicy(QSizePolicy::Expanding, QSizePolicy::Preferred);

    // Create the denoise button
    QPushButton *denoiseButton = new QPushButton("Denoise");
    QObject::connect(denoiseButton, &QPushButton::clicked, updateDenoisedImage);

    // Add the lambda slider and its label
    layout->addWidget(lambdaLabel, 4, 0);      // Row 4, column 0
    layout->addWidget(lambdaSlider, 4, 1);     // Row 4, column 1
    layout->addWidget(lambdaValueLabel, 5, 0); // Row 5, column 0

    // Add the iterations slider and its label
    layout->addWidget(iterationsLabel, 4, 2);      // Row 4, column 2
    layout->addWidget(iterationsSlider, 4, 3);     // Row 4, column 3
    layout->addWidget(iterationsValueLabel, 5, 2); // Row 5, column 2

    // Add the denoise button
    layout->addWidget(denoiseButton, 6, 0, 1, 4); // Row 6, spanning 1 row, 4 columns

    QList<QByteArray> supportedFormats = QImageReader::supportedImageFormats();
    qDebug() << "Supported image formats:" << supportedFormats;

    // Load and display the input image
    QString inputImagePath = inputImagePathEdit->text();
    QImage inputImage(inputImagePath);
    if (inputImage.isNull())
    {
        qDebug() << "Failed to load input image";
    }
    else
    {
        qDebug() << "Input image loaded successfully";
    }
    QPixmap inputPixmap = QPixmap::fromImage(inputImage);
    originalImageDisplay->setPixmap(inputPixmap.scaled(originalImageDisplay->size(), Qt::KeepAspectRatio));

    // Load and display the output image
    QString outputImagePath = outputImagePathEdit->text();
    QImage outputImage(outputImagePath);
    if (outputImage.isNull())
    {
        qDebug() << "Failed to load input image";
    }
    else
    {
        qDebug() << "Input image loaded successfully";
    }
    QPixmap outputPixmap = QPixmap::fromImage(outputImage);
    resultImageDisplay->setPixmap(outputPixmap.scaled(resultImageDisplay->size(), Qt::KeepAspectRatio));

    // Show the main window
    window.show();

    return app.exec();
}