#include "stackercpp.h"

Stacker::Stacker(bool verbose)
{
    Stacker::verbose = verbose;
    Stacker::export_width = 0;
    Stacker::export_height = 0;
    Stacker::layers = 0;
    Stacker::trimratio = 0;
    cv::setNumThreads(0);
    Stacker::stopwatch = clock();
}

void
Stacker::add_frame(unsigned char *data, int width, int height)
{
    stopwatch_start();
    Mat warp_matrix = Mat::eye(3, 3, CV_32F);
    Mat mat = Mat(height, width, CV_8UC3, data);
    Mat grayscale = Mat(height, width, CV_8UC1);
    cv::cvtColor(mat, grayscale, cv::COLOR_BGR2GRAY);
    stopwatch_mark("grayscale input");

    stopwatch_start();
    Scalar mean, stddev;
    meanStdDev(grayscale, mean, stddev);
    printf("mean: %f, dev: %f\n", mean[0], stddev[0]);
    if (mean[0] < 10) {
        return;
    }
    stopwatch_mark("filter");

    int number_of_iterations = 5;
    double termination_eps = 1e-10;
    TermCriteria criteria(TermCriteria::COUNT + TermCriteria::EPS, number_of_iterations, termination_eps);
    if (layers == 0) {
        // First image in the stack is used as the reference to align the next frames to
        Stacker::reference = grayscale;

        // Create black image to accumulate the average
        Stacker::stacked = Mat(height, width, CV_64FC3);
        Stacker::stacked.setTo(Scalar(0, 0, 0, 0));

        // Add first frame to the stack
        Stacker::stacked += mat;
        Stacker::layers += 1;
    } else {
        // All frames after the initial one are stacked
        Mat warped = Mat(Stacker::stacked.rows, Stacker::stacked.cols, CV_32FC1);

        stopwatch_start();
        cv::findTransformECC(grayscale, Stacker::reference, warp_matrix, MOTION_HOMOGRAPHY, criteria);
        stopwatch_mark("find alignment");

        stopwatch_start();
        warpPerspective(mat, warped, warp_matrix, warped.size(), INTER_LINEAR);
        stopwatch_mark("warp image");

        // Check how much the image should be cropped to hide the warped edges
        float current_trimratio = cv::videostab::estimateOptimalTrimRatio(warp_matrix, mat.size());
        Stacker::trimratio = std::max(Stacker::trimratio, current_trimratio);

        // Add the warped image to the stack
        Stacker::stacked += warped;
        Stacker::layers += 1;
    }
}

Mat
Stacker::postprocess_mat(Mat input)
{
    stopwatch_start();
    int h_crop = (int) ((float) input.cols * Stacker::trimratio);
    int v_crop = (int) ((float) input.rows * Stacker::trimratio);
    Mat cropped;
    input(Rect(h_crop, v_crop, input.cols - h_crop - h_crop, input.rows - v_crop - v_crop)).copyTo(input);
    stopwatch_mark("trim");

    stopwatch_start();
    Mat blur;
    GaussianBlur(input, blur, Size(0, 0), 1.8);
    std::vector<cv::Mat> rgb_planes(3);
    cv::split(blur, rgb_planes);
    double min_r, max_r, min_g, max_g, min_b, max_b;
    minMaxIdx(rgb_planes[0], &min_b, &max_b);
    minMaxIdx(rgb_planes[1], &min_g, &max_g);
    minMaxIdx(rgb_planes[2], &min_r, &max_r);
    input = input - Scalar(min_b + 5, min_g + 5, min_r + 5);
    double scale_r, scale_g, scale_b;
    scale_r = 255 / (max_r - min_r + 5);
    scale_g = (255 / (max_g - min_g + 5));
    scale_b = 255 / (max_b - min_b + 5);
    multiply(input, Scalar(scale_b, scale_g, scale_r), input);
    stopwatch_mark("levels");

    stopwatch_start();
    float gamma = 1 / 0.9;
    Mat table(1, 256, CV_8U);
    uchar *p = table.ptr();
    for (int i = 0; i < 256; ++i) {
        p[i] = (uchar) (pow(i / 255.0, gamma) * 255);
    }
    LUT(input, table, input);
    stopwatch_mark("gamma");

    stopwatch_start();
    Mat sharpened;
    GaussianBlur(input, sharpened, Size(0, 0), 1.7);
    addWeighted(input, 2.5, sharpened, -1.5, 0, sharpened);
    stopwatch_mark("sharpen");

    /* Disabled CLAHE local contrast, it's a bit to overpronounced and doesn't
     * seem really necessary at this point

    stopwatch_start();
    Mat lab;
    cvtColor(sharpened, lab, COLOR_BGR2Lab);
    std::vector<cv::Mat> lab_planes(3);
    cv::split(lab, lab_planes);
    stopwatch_mark("to Lab");

    stopwatch_start();
    cv::Ptr<cv::CLAHE> clahe = cv::createCLAHE();
    clahe->setClipLimit(1);
    clahe->setTilesGridSize(Size(8, 8));
    cv::Mat dst;
    clahe->apply(lab_planes[0], dst);
    dst.copyTo(lab_planes[0]);
    stopwatch_mark("clahe");

    stopwatch_start();
    Mat result;
    cv::merge(lab_planes, lab);
    cvtColor(lab, result, COLOR_Lab2BGR);
    stopwatch_mark("to RGB");
    */

    return sharpened;
}

char *
Stacker::get_result()
{
    // Complete the averaging and go back to an 8-bit image
    stopwatch_start();
    Stacker::stacked.convertTo(Stacker::stacked, CV_8U, 1. / Stacker::layers);
    stopwatch_mark("average");

    // Run the final-frame postprocessing
    Mat result = postprocess_mat(Stacker::stacked);
    Stacker::export_width = result.cols;
    Stacker::export_height = result.rows;

    // Convert mat to bytes
    size_t size = result.total() * result.elemSize();
    char *data = (char *) malloc(size);
    std::memcpy(data, result.data, size * sizeof(char));
    return data;
}

char *
Stacker::postprocess(unsigned char *data, int width, int height)
{
    // Convert bytes to mat
    Mat mat = Mat(height, width, CV_8UC3, data);

    // Run the final-frame postprocessing
    Mat result = postprocess_mat(mat);
    Stacker::export_width = result.cols;
    Stacker::export_height = result.rows;

    // Convert mat to bytes
    size_t size = result.total() * result.elemSize();
    char *outdata = (char *) malloc(size);
    std::memcpy(outdata, result.data, size * sizeof(char));
    return outdata;
}

void
Stacker::stopwatch_start()
{
    Stacker::stopwatch = clock();
}

void
Stacker::stopwatch_mark(const char *name)
{
    if (Stacker::verbose) {
        printf("[%.1fms] %s\n", float(clock() - Stacker::stopwatch) / CLOCKS_PER_SEC * 1000, name);
    }
}

int
Stacker::get_width()
{
    return Stacker::export_width;
}

int
Stacker::get_height()
{
    return Stacker::export_height;
}