#include "stackercpp.h"

Stacker::Stacker()
{
    Stacker::layers = 0;
}

void
Stacker::add_frame(unsigned char *data, int width, int height)
{
    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);
    int number_of_iterations = 5000;
    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);
        cv::findTransformECC(grayscale, Stacker::reference, warp_matrix, MOTION_HOMOGRAPHY, criteria);
        warpPerspective(mat, warped, warp_matrix, warped.size(), INTER_LINEAR);

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

char *
Stacker::get_result()
{
    Stacker::stacked.convertTo(Stacker::stacked, CV_8U, 1. / Stacker::layers);
    size_t size = Stacker::stacked.total() * Stacker::stacked.elemSize();
    char *data = (char *) malloc(size);
    std::memcpy(data, Stacker::stacked.data, size * sizeof(char));
    return data;
}