/**
 * @file        main.cpp
 * @author      Ondrej Klima, BUT FIT Brno, iklima@fit.vutbr.cz
 * @version     1.0
 * @date        15 June 2020
 *
 * @brief       2D-3D Registration of product / package to single image
 *              Usage:
 *
 *              register.exe <inputMask> <model> <cameraCalibration> <startingPose> <outputImage> <resultPose> <saveImagesPath> <verbose> <inputImage>
 *
 *              <inputMask> - Input binary mask obtained from segmentation by Mask-RCNN
 *              <model> - Polygonal model in .ply format and milimeter scale
 *              <cameraCalibration> - Camera calibration for the registration
 *              <startingPose> - The initial pose for the registration to start with
 *              <outputImage> - Image of the product in standardized pose and resolution. If set to "none", no file is stored.
 *              <resultPose> - File name for the result csv containig target rotation and translation
 *              <saveImagesPath> - Folder for saving the image of each iteration. If set to "none", no images are stored.
 *              <verbose> - If *not* set to "0", current value of the objective function is diplayed at each iteration
 *              <inputImage> - Input image to be standardized
 *
 * @license     BUT Licence 1.0
 *
 */

#include <QApplication>
#include <QDebug>
#include <QTextStream>
#include <QFile>
#include <QProcess>

#define _USE_MATH_DEFINES
#include <cmath>
#include <iostream>
#include <stdio.h>
#include <stdlib.h>

#include <ssimrenderer.h>
#include <libmultifragmentregister.h>

#include <ImageMetric/tsimplemetric.h>
#include <ImageMetric/tsimplemetricmask.h>
#include <ImageMetric/topenglsquareddifferencesmetric.h>
#include <VertexMetric/tsimplevertexmetric.h>
#include <Observer/tdefaultobserver.h>

#include "../include/standardizePose.h"

int main(int argc, char *argv[])
{
    QApplication a(argc, argv);
    a.addLibraryPath("plugins");


    if(argc != 10)
    {
        qDebug() << "Error: Wrong number of input arguments.";
        qDebug() << "Usage:";
        qDebug() << "register.exe <inputMask> <model> <cameraCalibration> <startingPose> <outputImage> <resultPose> <saveImagesPath> <verbose> <inputImage>";

        return EXIT_FAILURE;
    }

    // Parsing input parameters
    QString inputMaskFileName = argv[1];
    QString modelFileName = argv[2];
    QString cameraCalibrationFileName = argv[3];
    QString startingPoseFileName = argv[4];
    QString outputImageFileName = argv[5];
    QString resultPoseFileName = argv[6];
    QString saveImagesPath = argv[7];
    QString verbose = argv[8];
    QString inputImageFileName = argv[9];


    typedef LibMultiFragmentRegister<TOpenGLSquaredDifferencesMetric> RegistrationType;
    RegistrationType * registration = RegistrationType::New(1, 1, new TSimpleVertexMetric);

    SSIMRenderer::PLYMeshFile * meshFile;
    QVector<float> poseFile;

    try {
        meshFile         = new SSIMRenderer::PLYMeshFile(modelFileName);
        poseFile         = RegistrationType::loadDataFromCSVFile(startingPoseFileName);
    }
    catch (std::exception &e) {
        // Wrong file
        qFatal(e.what());
        exit(EXIT_FAILURE);
    }

    //QVector<QVector3D> rotation(   1, QVector3D(0, 270, 270));
    //QVector<QVector3D> translation(1, QVector3D(16, 0, -9));

    QVector<QVector3D> rotation(   1, QVector3D(poseFile.at(0), poseFile.at(1), poseFile.at(2)));
    QVector<QVector3D> translation(1, QVector3D(poseFile.at(3), poseFile.at(4), poseFile.at(5)));

    registration->enableDensity(false);
    registration->setMeshModel(meshFile);
    registration->enableMirroring(false);

    registration->setRotations(rotation);
    registration->setTranslations(translation);

    QVector<SSIMRenderer::Pyramid> perspectives(1);
    perspectives[0] = SSIMRenderer::CSVPyramidFile(cameraCalibrationFileName);

    registration->setPerspectives(perspectives);

    QVector<QSize> sizes(1, QSize(5755, 5755));
    registration->setSizes(sizes);

    QVector<QRect> crops(1, QRect(2496,1791,768,1229));
    registration->setCrops(crops);

    QVector<QImage> images(1);
    // Produkt
    images[0] = QImage(inputMaskFileName);
    registration->setImages(images);

    TDefaultObserver * observer = new TDefaultObserver;
    observer->setVerbose(verbose != "0");
    observer->enableImages(saveImagesPath != "none");
    observer->setImagesPath(saveImagesPath);
    observer->setRefImages(images);
    observer->setWholeBones(true);
    registration->setObserver(observer);

    registration->setRotations(rotation);
    registration->setTranslations(translation);

    registration->setPoseEps(10);
    registration->optimizePose();

    if(observer->getValues().last() > 0)
    {
        registration->setPoseEps(1);
        registration->optimizePose();
    }
    if(observer->getValues().last() > 0)
    {
        registration->setPoseEps(0.1);
        registration->optimizePose();
    }
    if(observer->getValues().last() > 0)
    {
        registration->setPoseEps(0.01);
        registration->optimizePose();
    }
    if(observer->getValues().last() > 0)
    {
        registration->setPoseEps(0.001);
        registration->optimizePose();
    }

    QFile results(resultPoseFileName);
    results.open(QIODevice::Append);
    QTextStream out(&results);

    QVector3D resultRotation = registration->getRotations().at(0);
    QVector3D resultTranslation = registration->getRotations().at(0);

    out << resultRotation.x() << ';' << resultRotation.y() << ';' << resultRotation.z() << "\r\n";
    out << resultTranslation.x() << ';' << resultTranslation.y() << ';' << resultTranslation.z() << "\r\n";
    out.flush();
    results.close();

    if(outputImageFileName != "none")
    {
        QVector<QImage> registratedImages = registration->getImages();
        standardizePose(registratedImages.at(0),
                        inputImageFileName,
                        outputImageFileName);
    }

    delete observer;
    delete registration;
    delete meshFile;

    return EXIT_SUCCESS;
}
