// This file is part of OpenCV project.
// It is subject to the license terms in the LICENSE file found in the top-level directory
// of this distribution and at http://opencv.org/license.html.
// Copyright Amir Hassan (kallaballa) <amir@viel-zu.org>

#include "opencv2/viz2d/viz2d.hpp"
#include "opencv2/viz2d/nvg.hpp"

#include <vector>
#include <string>

#include <opencv2/dnn.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/face.hpp>
#include <opencv2/stitching/detail/blenders.hpp>
#include <opencv2/tracking.hpp>

using std::cerr;
using std::endl;
using std::vector;
using std::string;

#ifdef __EMSCRIPTEN__
//enables KCF tracking instead of continuous detection.
#define USE_TRACKER 1;
#endif

/** Application parameters **/
constexpr unsigned int WIDTH = 1920;
constexpr unsigned int HEIGHT = 1080;
constexpr double SCALE = 0.125;
constexpr bool OFFSCREEN = false;
constexpr const char *OUTPUT_FILENAME = "beauty-demo.mkv";
const unsigned long DIAG = hypot(double(WIDTH), double(HEIGHT));

/** Effect parameters **/
constexpr int BLUR_DIV = 500;
int blur_skin_kernel_size = std::max(int(DIAG / BLUR_DIV % 2 == 0 ? DIAG / BLUR_DIV + 1 : DIAG / BLUR_DIV), 1);
float eyes_and_lips_saturation = 1.5f; //factor
float skin_saturation = 1.2f; //factor
float skin_contrast = 0.6f; //0.0-1.0
#ifndef __EMSCRIPTEN__
bool side_by_side = true;
bool stretch = true;
#else
bool side_by_side = false;
bool stretch = false;
#endif

static cv::Ptr<cv::viz::Viz2D> v2d = cv::viz::Viz2D::make(cv::Size(WIDTH, HEIGHT), cv::Size(WIDTH, HEIGHT), OFFSCREEN, "Beauty Demo");
static cv::Ptr<cv::face::Facemark> facemark = cv::face::createFacemarkLBF();
#ifdef USE_TRACKER
static cv::Ptr<cv::Tracker> tracker = cv::TrackerKCF::create();
#endif

struct FaceFeatures {
    cv::Rect faceRect_;
    vector<cv::Point2f> chin_;
    vector<cv::Point2f> top_nose_;
    vector<cv::Point2f> bottom_nose_;
    vector<cv::Point2f> left_eyebrow_;
    vector<cv::Point2f> right_eyebrow_;
    vector<cv::Point2f> left_eye_;
    vector<cv::Point2f> right_eye_;
    vector<cv::Point2f> outer_lips_;
    vector<cv::Point2f> inside_lips_;

    FaceFeatures(const cv::Rect &faceRect, const vector<cv::Point2f> &shape, double scale) {
        faceRect_ = cv::Rect(faceRect.x / scale, faceRect.y / scale, faceRect.width / scale, faceRect.height / scale);
        size_t i = 0;
        // Around Chin. Ear to Ear
        for (i = 0; i <= 16; ++i)
            chin_.push_back(shape[i] / scale);
        // left eyebrow
        for (; i <= 21; ++i)
            left_eyebrow_.push_back(shape[i] / scale);
        // Right eyebrow
        for (; i <= 26; ++i)
            right_eyebrow_.push_back(shape[i] / scale);
        // Line on top of nose
        for (; i <= 30; ++i)
            top_nose_.push_back(shape[i] / scale);
        // Bottom part of the nose
        for (; i <= 35; ++i)
            bottom_nose_.push_back(shape[i] / scale);
        // Left eye
        for (; i <= 41; ++i)
            left_eye_.push_back(shape[i] / scale);
        // Right eye
        for (; i <= 47; ++i)
            right_eye_.push_back(shape[i] / scale);
        // Lips outer part
        for (; i <= 59; ++i)
            outer_lips_.push_back(shape[i] / scale);
        // Lips inside part
        for (; i <= 67; ++i)
            inside_lips_.push_back(shape[i] / scale);
    }

    vector<cv::Point2f> points() const {
        vector<cv::Point2f> allPoints;
        allPoints.insert(allPoints.begin(), chin_.begin(), chin_.end());
        allPoints.insert(allPoints.begin(), top_nose_.begin(), top_nose_.end());
        allPoints.insert(allPoints.begin(), bottom_nose_.begin(), bottom_nose_.end());
        allPoints.insert(allPoints.begin(), left_eyebrow_.begin(), left_eyebrow_.end());
        allPoints.insert(allPoints.begin(), right_eyebrow_.begin(), right_eyebrow_.end());
        allPoints.insert(allPoints.begin(), left_eye_.begin(), left_eye_.end());
        allPoints.insert(allPoints.begin(), right_eye_.begin(), right_eye_.end());
        allPoints.insert(allPoints.begin(), outer_lips_.begin(), outer_lips_.end());
        allPoints.insert(allPoints.begin(), inside_lips_.begin(), inside_lips_.end());

        return allPoints;
    }

    vector<vector<cv::Point2f>> features() const {
        return {chin_,
            top_nose_,
            bottom_nose_,
            left_eyebrow_,
            right_eyebrow_,
            left_eye_,
            right_eye_,
            outer_lips_,
            inside_lips_};
    }

    size_t empty() const {
        return points().empty();
    }
};

void draw_face_oval_mask(const vector<FaceFeatures> &lm) {
    using namespace cv::viz::nvg;
    for (size_t i = 0; i < lm.size(); i++) {
        vector<vector<cv::Point2f>> features = lm[i].features();
        cv::RotatedRect rotRect = cv::fitEllipse(features[0]);

        beginPath();
        fillColor(cv::Scalar(255, 255, 255, 255));
        ellipse(rotRect.center.x, rotRect.center.y * 1, rotRect.size.width / 2, rotRect.size.height / 2.5);
        rotate(rotRect.angle);
        fill();
    }
}

void draw_face_eyes_and_lips_mask(const vector<FaceFeatures> &lm) {
    using namespace cv::viz::nvg;
    for (size_t i = 0; i < lm.size(); i++) {
        vector<vector<cv::Point2f>> features = lm[i].features();
        for (size_t j = 5; j < 8; ++j) {
            beginPath();
            fillColor(cv::Scalar(255, 255, 255, 255));
            moveTo(features[j][0].x, features[j][0].y);
            for (size_t k = 1; k < features[j].size(); ++k) {
                lineTo(features[j][k].x, features[j][k].y);
            }
            closePath();
            fill();
        }

        beginPath();
        fillColor(cv::Scalar(0, 0, 0, 255));
        moveTo(features[8][0].x, features[8][0].y);
        for (size_t k = 1; k < features[8].size(); ++k) {
            lineTo(features[8][k].x, features[8][k].y);
        }
        closePath();
        fill();
    }
}

void adjust_saturation(const cv::UMat &srcBGR, cv::UMat &dstBGR, float factor) {
    static vector<cv::UMat> channels;
    static cv::UMat hls;

    cvtColor(srcBGR, hls, cv::COLOR_BGR2HLS);
    split(hls, channels);
    cv::multiply(channels[2], factor, channels[2]);
    merge(channels, hls);
    cvtColor(hls, dstBGR, cv::COLOR_HLS2BGR);
}

void setup_gui(cv::Ptr<cv::viz::Viz2D> v2d) {
    v2d->nanogui([&](cv::viz::FormHelper& form){
        form.makeDialog(5, 30, "Effect");

        form.makeGroup("Display");
        form.makeFormVariable("Side by side", side_by_side, "Enable or disable side by side view");
        form.makeFormVariable("Stretch", stretch, "Enable or disable stetching to the window size");
    #ifndef __EMSCRIPTEN__
        form.makeButton("Fullscreen", [=]() {
            v2d->setFullscreen(!v2d->isFullscreen());
        });
    #endif
        form.makeButton("Offscreen", [=]() {
            v2d->setOffscreen(!v2d->isOffscreen());
        });

        form.makeGroup("Face Skin");
        auto* kernelSize = form.makeFormVariable("Blur", blur_skin_kernel_size, 0, 256, true, "", "use this kernel size to blur the face skin");
        kernelSize->set_callback([=](const int& k) {
            static int lastKernelSize = blur_skin_kernel_size;

            if(k == lastKernelSize)
                return;

            if(k <= lastKernelSize) {
                blur_skin_kernel_size = std::max(int(k % 2 == 0 ? k - 1 : k), 1);
            } else if(k > lastKernelSize)
                blur_skin_kernel_size = std::max(int(k % 2 == 0 ? k + 1 : k), 1);

            lastKernelSize = k;
            kernelSize->set_value(blur_skin_kernel_size);
        });
        form.makeFormVariable("Saturation", skin_saturation, 0.0f, 100.0f, true, "", "adjust the skin saturation by this amount");
        form.makeFormVariable("Contrast", skin_contrast, 0.0f, 1.0f, true, "", "contrast amount of the face skin");

        form.makeGroup("Eyes and Lips");
        form.makeFormVariable("Saturation", eyes_and_lips_saturation, 0.0f, 100.0f, true, "", "adjust the saturation of the eyes and the lips by this amount");
    });
}

bool iteration() {
    try {
#ifdef USE_TRACKER
        static bool trackerInitalized = false;
#endif
        static cv::Ptr<cv::FaceDetectorYN> detector = cv::FaceDetectorYN::create("assets/face_detection_yunet_2022mar.onnx", "", cv::Size(v2d->getFrameBufferSize().width * SCALE, v2d->getFrameBufferSize().height * SCALE), 0.9, 0.3, 5000, cv::dnn::DNN_BACKEND_OPENCV, cv::dnn::DNN_TARGET_OPENCL);
        static cv::detail::MultiBandBlender blender(false, 5);
        blender.prepare(cv::Rect(0, 0, WIDTH, HEIGHT));
        //BGR
        static cv::UMat input, down, blurred, contrast, faceOval, eyesAndLips, skin;
        static cv::UMat frameOut(HEIGHT, WIDTH, CV_8UC3);
        static cv::UMat lhalf(HEIGHT * SCALE, WIDTH * SCALE, CV_8UC3);
        static cv::UMat rhalf(lhalf.size(), lhalf.type());
        //GREY
        static cv::UMat faceSkinMaskGrey, eyesAndLipsMaskGrey, backgroundMaskGrey;
        //BGR-Float
        static cv::UMat frameOutFloat;

        static cv::Mat faces;
        static cv::Rect trackedFace;
        static vector<cv::Rect> faceRects;
        static vector<vector<cv::Point2f>> shapes;
        static vector<FaceFeatures> featuresList;

        if (!v2d->capture())
            return false;

        v2d->fb([&](cv::UMat &frameBuffer) {
            cvtColor(frameBuffer, input, cv::COLOR_BGRA2BGR);
        });

        cv::resize(input, down, cv::Size(0, 0), SCALE, SCALE);

        shapes.clear();
        faceRects.clear();
#ifdef USE_TRACKER
        if (trackerInitalized && tracker->update(down, trackedFace)) {
            faceRects.push_back(trackedFace);
        } else
#endif
        {
            detector->detect(down, faces);

            for (int i = 0; i < faces.rows; i++) {
                faceRects.push_back(cv::Rect(int(faces.at<float>(i, 0)), int(faces.at<float>(i, 1)), int(faces.at<float>(i, 2)), int(faces.at<float>(i, 3))));
            }
        }

        if (!faceRects.empty() && facemark->fit(down, faceRects, shapes)) {
#ifdef USE_TRACKER
            if(!trackerInitalized) {
                tracker->init(down, faceRects[0]);
                trackerInitalized = true;
            }
#endif
            featuresList.clear();
            for (size_t i = 0; i < faceRects.size(); ++i) {
                featuresList.push_back(FaceFeatures(faceRects[i], shapes[i], float(down.size().width) / WIDTH));
            }

            v2d->nvg([&](const cv::Size &sz) {
                v2d->clear();
                //Draw the face oval
                draw_face_oval_mask(featuresList);
            });

            v2d->fb([&](cv::UMat &frameBuffer) {
                //Convert/Copy the mask
                cvtColor(frameBuffer, faceOval, cv::COLOR_BGRA2GRAY);
            });

            v2d->nvg([&](const cv::Size &sz) {
                v2d->clear();
                //Draw eyes eyes and lips
                draw_face_eyes_and_lips_mask(featuresList);
            });

            v2d->fb([&](cv::UMat &frameBuffer) {
                //Convert/Copy the mask
                cvtColor(frameBuffer, eyesAndLipsMaskGrey, cv::COLOR_BGRA2GRAY);
            });

            cv::subtract(faceOval, eyesAndLipsMaskGrey, faceSkinMaskGrey);
            cv::bitwise_not(eyesAndLipsMaskGrey,backgroundMaskGrey);

            //boost saturation of eyes and lips
            adjust_saturation(input,eyesAndLips, eyes_and_lips_saturation);
            //reduce skin contrast
            multiply(input, cv::Scalar::all(skin_contrast), contrast);
            //fix skin brightness
            add(contrast, cv::Scalar::all((1.0 - skin_contrast) / 2.0) * 255.0, contrast);
            //blur the skin
            cv::boxFilter(contrast, blurred, -1, cv::Size(blur_skin_kernel_size, blur_skin_kernel_size), cv::Point(-1, -1), true, cv::BORDER_REPLICATE);
            //boost skin saturation
            adjust_saturation(blurred,skin, skin_saturation);

            //piece it all together
            blender.feed(skin, faceSkinMaskGrey, cv::Point(0, 0));
            blender.feed(input, backgroundMaskGrey, cv::Point(0, 0));
            blender.feed(eyesAndLips, eyesAndLipsMaskGrey, cv::Point(0, 0));
            blender.blend(frameOutFloat, cv::UMat());
            frameOutFloat.convertTo(frameOut, CV_8U, 1.0);

            if (side_by_side) {
                cv::resize(input, lhalf, cv::Size(0, 0), 0.5, 0.5);
                cv::resize(frameOut, rhalf, cv::Size(0, 0), 0.5, 0.5);

                frameOut = cv::Scalar::all(0);
                lhalf.copyTo(frameOut(cv::Rect(0, 0, lhalf.size().width, lhalf.size().height)));
                rhalf.copyTo(frameOut(cv::Rect(rhalf.size().width, 0, rhalf.size().width, rhalf.size().height)));
            }

            v2d->fb([&](cv::UMat &frameBuffer) {
                cvtColor(frameOut, frameBuffer, cv::COLOR_BGR2BGRA);
            });
        } else {
            if (side_by_side) {
                frameOut = cv::Scalar::all(0);
                cv::resize(input, lhalf, cv::Size(0, 0), 0.5, 0.5);
                lhalf.copyTo(frameOut(cv::Rect(0, 0, lhalf.size().width, lhalf.size().height)));
                lhalf.copyTo(frameOut(cv::Rect(lhalf.size().width, 0, lhalf.size().width, lhalf.size().height)));
            } else {
                input.copyTo(frameOut);
            }

            v2d->fb([&](cv::UMat &frameBuffer) {
                cvtColor(frameOut, frameBuffer, cv::COLOR_BGR2BGRA);
            });
        }
        updateFps(v2d, true);

#ifndef __EMSCRIPTEN__
        v2d->write();
#endif

        //If onscreen rendering is enabled it displays the framebuffer in the native window. Returns false if the window was closed.
        if (!v2d->display())
            return false;
    } catch (std::exception &ex) {
        cerr << ex.what() << endl;
        return false;
    }
    return true;
}

int main(int argc, char **argv) {
    using namespace cv::viz;
#ifndef __EMSCRIPTEN__
    if (argc != 2) {
        std::cerr << "Usage: beauty-demo <input-video-file>" << endl;
        exit(1);
    }
#endif
    facemark->loadModel("assets/lbfmodel.yaml");

    printSystemInfo();

    v2d->setStretching(stretch);

    if (!v2d->isOffscreen()) {
        setup_gui(v2d);
        v2d->setVisible(true);
    }

#ifndef __EMSCRIPTEN__
    Source src = makeCaptureSource(argv[1]);
    v2d->setSource(src);
    Sink sink = makeWriterSink(OUTPUT_FILENAME, cv::VideoWriter::fourcc('V', 'P', '9', '0'), src.fps(), cv::Size(WIDTH, HEIGHT));
    v2d->setSink(sink);
#endif

    v2d->run(iteration);

    return 0;
}