#include <iostream>
#include <random>
#include <cstdint>
#include <limits>
#include <cmath>

#include <opencv2/core/core.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/highgui/highgui.hpp>

using namespace std;

cv::Mat VoronoiImage(const cv::Mat& image);

int main(int argc, char* argv[])
{
    if (argc != 2) {
        cout << "This program takes one argument." << endl;
        return 1;
    }

    const string INPUT = argv[1];
    const string OUTPUT = "output.png";

    cv::Mat input = cv::imread(INPUT);

    if (!input.data) {
        cerr << "Error while opening input image" << endl;
        return 2;
    }

    cv::Mat image = VoronoiImage(input);

    cv::imwrite(OUTPUT, image);

    return 0;
}

struct WorleyResult
{
    double x;
    double y;
    double dist_sq;

    WorleyResult() : x(0.0), y(0.0), dist_sq(numeric_limits<double>::max()) {}

    WorleyResult(double _x, double _y, double _dist_sq) : x(_x), y(_y), dist_sq(_dist_sq) {}
};

int GenerateSeed(int i, int j)
{
    return (541 * i + 79 * j) % numeric_limits<int>::max();
}

tuple<WorleyResult, WorleyResult> Worley(double x, double y)
{
    default_random_engine generator;
    uniform_real_distribution<double> distribution(0.0, 1.0);

    // In which cell is the point
    const double int_x = floor(x);
    const double int_y = floor(y);

    const double u = x - int_x;
    const double v = y - int_y;

    const int cx = int(int_x);
    const int cy = int(int_y);

    // Distance to the first and second nearest neighbors
    WorleyResult f1, f2;

    // Exploring neighboring cells
    for (int i = cy - 1; i <= cy + 1; i++)
    {
        for (int j = cx - 1; j <= cx + 1; j++)
        {
            // Fixed seed for internal consistency
            generator.seed(GenerateSeed(i, j));

            // How many points in this cell
            int n = 1;
            for (int p = 0; p < n; p++)
            {
                // Generate each points
                double px = distribution(generator);
                double py = distribution(generator);

                double dist_sq = hypot(j + px - x, i + py - y);

                if (dist_sq < f1.dist_sq)
                {
                    f2 = f1;
                    f1 = WorleyResult(j + px, i + py, dist_sq);
                }
                else if (dist_sq < f2.dist_sq)
                {
                    f2 = WorleyResult(j + px, i + py, dist_sq);
                }
            }
        }
    }
    
    return make_tuple(f1, f2);
}

cv::Mat VoronoiImage(const cv::Mat& image)
{
    const int INPUT_WIDTH = image.cols;
    const int INPUT_HEIGHT = image.rows;

    // Size of the output image
    const int OUTPUT_WIDTH = image.cols;
    const int OUTPUT_HEIGHT = image.rows;

    // Number of lattice cells along each axis
    const int CELLS = 64;

    cv::Mat result(OUTPUT_HEIGHT, OUTPUT_WIDTH, CV_8UC3);

#pragma omp parallel for shared(image)
    for (int i = 0; i < OUTPUT_HEIGHT; i++) {
        for (int j = 0; j < OUTPUT_WIDTH; j++) {
            const double x = double(j) * CELLS / OUTPUT_WIDTH;
            const double y = double(i) * CELLS / OUTPUT_HEIGHT;
            
            // Identify the nearest neighbor in the lattice
            WorleyResult f1, f2;
            tie(f1, f2) = Worley(x, y);

            // Assign the color of the nearest neighbor
            const int nearest_i = int(round(f1.y * INPUT_HEIGHT / CELLS));
            const int nearest_j = int(round(f1.x * INPUT_WIDTH / CELLS));
            
            if (nearest_i >= 0 && nearest_i < INPUT_HEIGHT && nearest_j >= 0 && nearest_j < INPUT_WIDTH)
            {
                result.at<cv::Vec3b>(i, j) = image.at<cv::Vec3b>(nearest_i, nearest_j);
            }
        }
    }

    return result;
}