#include <opencv2/core/core.hpp>
#include <opencv2/imgcodecs/imgcodecs.hpp>
#include <vector>

using namespace cv;
using namespace std;


Vec3b find_nearest_color(vector<Vec3b>& palette, Vec3b& color, Vec3f& weights) {
    Vec3b palette_color;
    Vec3b nearest_color;
    float color_dist;
    float min_color_dist = 3 * 255;
    Vec3b diff;

    vector<Vec3b>::iterator it;
    for (it = palette.begin(); it != palette.end(); ++it) {
        palette_color = *it;
        absdiff(color, palette_color, diff);
        color_dist = weights[0] * diff[0]
            + weights[1] * diff[1]
            + weights[2] * diff[2];
        if (color_dist < min_color_dist) {
            min_color_dist = color_dist;
            nearest_color = palette_color;
        }
    }
    return nearest_color;
}


void dither(Mat& _img, vector<Vec3b>& palette) {
    // accept only char type matrices
    CV_Assert(_img.depth() == CV_8U);

    Mat_<Vec3b> img = _img;

    Vec3b color;
    Vec3b nearest_color;
    Vec3i err, d0, d1, d2, d3;
    Vec3f weights = Vec3f(0.11, 0.59, 0.30);

    int i, j;
    uchar *p, *p1;
    for (i = 0; i < img.rows - 1; ++i) {
        p = img.ptr<uchar>(i);
        p1 = img.ptr<uchar>(i+1);
        for (j = 1; j < img.cols - 1; ++j) {
            //color = img(i, j);
            memcpy(&color, p+j*3, 3);
            nearest_color = find_nearest_color(palette, color, weights);
            err = (Vec3i)color - (Vec3i)nearest_color;
            //img(i, j) = nearest_color;
            memcpy(p+j*3, &nearest_color, 3);

            //d0 = (Vec3i)img(i, j + 1) +     Vec3i(err[0] * 7 >> 4, err[1] * 7 >> 4, err[2] * 7 >> 4);
            d0 = Vec3i(*(p + (j + 1) * 3), *(p + 1 + (j + 1) * 3), *(p + 2 + (j + 1) * 3)) +    Vec3i(err[0] * 7 >> 4, err[1] * 7 >> 4, err[2] * 7 >> 4);
            //d1 = (Vec3i)img(i + 1, j - 1) + Vec3i(err[0] * 3 >> 4, err[1] * 3 >> 4, err[2] * 3 >> 4);
            d1 = Vec3i(*(p1 + (j - 1) * 3), *(p1 + 1 + (j - 1) * 3), *(p1 + 2 + (j - 1) * 3)) + Vec3i(err[0] * 3 >> 4, err[1] * 3 >> 4, err[2] * 3 >> 4);
            //d2 = (Vec3i)img(i + 1, j) +     Vec3i(err[0] * 5 >> 4, err[1] * 5 >> 4, err[2] * 5 >> 4);
            d2 = Vec3i(*(p1 + j * 3), *(p1 + 1 + j * 3), *(p1 + 2 + j * 3)) +                   Vec3i(err[0] * 5 >> 4, err[1] * 5 >> 4, err[2] * 5 >> 4);
            //d3 = (Vec3i)img(i + 1, j + 1) + Vec3i(err[0]     >> 4, err[1]     >> 4, err[2]     >> 4);
            d3 = Vec3i(*(p1 + (j + 1) * 3), *(p1 + 1 + (j + 1) * 3), *(p1 + 2 + (j + 1) * 3)) + Vec3i(err[0]     >> 4, err[1]     >> 4, err[2]     >> 4);

            //img(i, j + 1) = Vec3b(saturate_cast<uchar>(d0[0]), saturate_cast<uchar>(d0[1]), saturate_cast<uchar>(d0[2]));
            memcpy(p + (j + 1) * 3, &Vec3b(saturate_cast<uchar>(d0[0]), saturate_cast<uchar>(d0[1]), saturate_cast<uchar>(d0[2])), 3);
            //img(i + 1, j - 1) = Vec3b(saturate_cast<uchar>(d1[0]), saturate_cast<uchar>(d1[1]), saturate_cast<uchar>(d1[2]));
            memcpy(p1 + (j - 1) * 3, &Vec3b(saturate_cast<uchar>(d1[0]), saturate_cast<uchar>(d1[1]), saturate_cast<uchar>(d1[2])), 3);
            //img(i + 1, j) = Vec3b(saturate_cast<uchar>(d2[0]), saturate_cast<uchar>(d2[1]), saturate_cast<uchar>(d2[2]));
            memcpy(p1 + j * 3, &Vec3b(saturate_cast<uchar>(d2[0]), saturate_cast<uchar>(d2[1]), saturate_cast<uchar>(d2[2])), 3);
            //img(i + 1, j + 1) = Vec3b(saturate_cast<uchar>(d3[0]), saturate_cast<uchar>(d3[1]), saturate_cast<uchar>(d3[2]));
            memcpy(p1 + (j + 1) * 3, &Vec3b(saturate_cast<uchar>(d3[0]), saturate_cast<uchar>(d3[1]), saturate_cast<uchar>(d3[2])), 3);
        }
    }
    _img = img;

}

int main() {

    Mat img;
    img = imread("source.jpg", IMREAD_COLOR);

    vector<Vec3b> palette;
    palette.push_back(Vec3b(189, 203, 189));
    palette.push_back(Vec3b(148, 93, 0));
    palette.push_back(Vec3b(8, 12, 123));
    palette.push_back(Vec3b(8, 60, 123));
    palette.push_back(Vec3b(8, 125, 123));
    palette.push_back(Vec3b(8, 125, 8));
    palette.push_back(Vec3b(123, 125, 8));
    palette.push_back(Vec3b(123, 12, 8));
    palette.push_back(Vec3b(123, 12, 57));
    palette.push_back(Vec3b(123, 12, 123));
    palette.push_back(Vec3b(123, 125, 123));
    palette.push_back(Vec3b(8, 12, 8));
    palette.push_back(Vec3b(99, 101, 99));
    palette.push_back(Vec3b(49, 48, 49));

    dither(img, palette);
    imwrite("dithered.png", img);

    return 0;
}