#include "EyeImageProcessor.h"
#include <opencv2/highgui.hpp>

using namespace cv;

EyeImageProcessor::EyeImageProcessor(QString id, QObject* parent)
    : QObject(parent)
    , cameraCalibration(nullptr)
    , id(id)
    , pupilDetectionMethod(nullptr)
    , pupilTrackingMethod(nullptr)
{
    availablePupilDetectionMethods.push_back(std::make_shared<PuRe>());
    availablePupilDetectionMethods.push_back(std::make_shared<ElSe>());
    availablePupilDetectionMethods.push_back(std::make_shared<ExCuSe>());

    availablePupilTrackingMethods.push_back(std::make_shared<PuReST>());

#ifdef STARBURST
    availablePupilDetectionMethods.push_back(std::make_shared<Starburst>());
#endif
#ifdef SWIRSKI
    availablePupilDetectionMethods.push_back(std::make_shared<Swirski>());
#endif
    settings = new QSettings(gCfgDir + "/" + id + " Image Processor.ini", QSettings::IniFormat);
    updateConfig();

    pmIdx = gPerformanceMonitor.enrol(id, "Image Processor");
}

void EyeImageProcessor::updateConfig()
{
    QMutexLocker locker(&cfgMutex);
    cfg.load(settings);

    pupilDetectionMethod = nullptr;
    for (auto method : availablePupilDetectionMethods)
        if (cfg.pupilDetectionMethod == QString(method->description().c_str()))
            pupilDetectionMethod = method;

    pupilTrackingMethod = nullptr;
    for (auto method : availablePupilTrackingMethods)
        if (cfg.pupilTrackingMethod == QString(method->description().c_str()))
            pupilTrackingMethod = method;
}

EyeImageProcessor::~EyeImageProcessor()
{
    availablePupilDetectionMethods.clear();
    availablePupilTrackingMethods.clear();

    if (settings)
        settings->deleteLater();
}

void EyeImageProcessor::process(Timestamp timestamp, const Mat& frame)
{
    // TODO: parametrize frame drop due to lack of processing power
    if (gPerformanceMonitor.shouldDrop(pmIdx, timestamp, 100))
        return;

    QMutexLocker locker(&cfgMutex);
    Timestamp processingStart = gTimer.elapsed();

    data.timestamp = timestamp;

    Q_ASSERT_X(frame.data != data.input.data, "Eye Image Processing", "Previous and current input image matches!");
    if (cfg.inputSize.width > 0 && cfg.inputSize.height > 0) {
        data.input = Mat(cfg.inputSize, frame.type());
        resize(frame, data.input, cfg.inputSize);
    } else {
        data.input = frame;
    }

    if (cfg.flip != CV_FLIP_NONE)
        flip(data.input, data.input, cfg.flip);

    if (data.input.channels() > 1) // TODO: make it algorithm dependent
        cvtColor(data.input, data.input, cv::COLOR_BGR2GRAY);

    float minPupilDiameterPx = cfg.minPupilDiameterRatio * data.input.rows;
    float maxPupilDiameterPx = cfg.maxPupilDiameterRatio * data.input.rows;

    data.pupil = Pupil();
    data.validPupil = false;
    if (pupilDetectionMethod != nullptr) {
        Rect userROI = {
            static_cast<int>(cfg.roi.x * data.input.cols),
            static_cast<int>(cfg.roi.y * data.input.rows),
            static_cast<int>(cfg.roi.width * data.input.cols),
            static_cast<int>(cfg.roi.height * data.input.rows)
        };

        float scalingFactor = 1;
        if (cfg.processingDownscalingFactor > 1)
            scalingFactor = static_cast<float>(1.0 / cfg.processingDownscalingFactor);

        /*
		 *  From here on, our reference frame is the scaled user ROI
		 */
        Mat downscaled;
        resize(data.input(userROI), downscaled, Size(),
            scalingFactor, scalingFactor,
            INTER_AREA);
        Rect coarseROI = { 0, 0, downscaled.cols, downscaled.rows };

        // Rescale pupil size limits as well
        minPupilDiameterPx *= scalingFactor;
        maxPupilDiameterPx *= scalingFactor;

        // If the user wants a coarse location and the method has none embedded,
        // we further constrain the search using the generic one
        if (!pupilDetectionMethod->hasCoarseLocation() && cfg.coarseDetection) {
            coarseROI = PupilDetectionMethod::coarsePupilDetection(downscaled, 0.5f, 60, 40);
            data.coarseROI = Rect(
                userROI.tl() + coarseROI.tl() / scalingFactor,
                userROI.tl() + coarseROI.br() / scalingFactor);
        } else
            data.coarseROI = Rect();

        if (pupilTrackingMethod)
            pupilTrackingMethod->detectAndTrack(timestamp, downscaled, coarseROI, data.pupil, pupilDetectionMethod, minPupilDiameterPx, maxPupilDiameterPx);
        else
            data.pupil = pupilDetectionMethod->detectWithConfidence(downscaled, coarseROI, minPupilDiameterPx, maxPupilDiameterPx);

        if (data.pupil.center.x > 0 && data.pupil.center.y > 0) {
            // Upscale
            data.pupil.resize(1.0f / scalingFactor);

            // User region shift
            data.pupil.shift(userROI.tl());
            data.validPupil = true;
        }
    }

    data.modelData = EyeModelData();

    data.cameraCalibration = cameraCalibration;
    data.processingTimestamp = gTimer.elapsed() - processingStart;

    emit newData(data);
}

void EyeImageProcessor::newROI(QPointF sROI, QPointF eROI)
{
    QMutexLocker locker(&cfgMutex);
    if (sROI.isNull() || eROI.isNull()) {
        cfg.roi = { 0.0f, 0.0f, 1.0f, 1.0f };
    } else {
        cfg.roi = {
            static_cast<float>(min(sROI.x(), eROI.x())),
            static_cast<float>(min(sROI.y(), eROI.y())),
            static_cast<float>(abs(eROI.x() - sROI.x())),
            static_cast<float>(abs(eROI.y() - sROI.y()))
        };
    }
}

void EyeImageProcessor::newMaxRadius(double maxRadius)
{
    cfg.maxPupilDiameterRatio = 2 * maxRadius;
}