/*
 Yliluoma's ordered dithering algorithm 3 a.
 Usage: ./yliluoma_ordered_dither_3a.exe scene.png scenedither11-gamma.png

 Arbitrary-palette positional dithering algorithm
 https://bisqwit.iki.fi/story/howto/dither/jy/
 */

#include <gd.h>
#include <stdio.h>
#include <math.h>
#include <algorithm> /* For std::sort() */
#include <vector>
#include <map> /* For associative container, std::map<> */

#define COMPARE_RGB 1

/* 8x8 threshold map */
static const unsigned char map[8 * 8] = {
    0, 48, 12, 60, 3, 51, 15, 63,
    32, 16, 44, 28, 35, 19, 47, 31,
    8, 56, 4, 52, 11, 59, 7, 55,
    40, 24, 36, 20, 43, 27, 39, 23,
    2, 50, 14, 62, 1, 49, 13, 61,
    34, 18, 46, 30, 33, 17, 45, 29,
    10, 58, 6, 54, 9, 57, 5, 53,
    42, 26, 38, 22, 41, 25, 37, 21};

static const double Gamma = 2.2; // Gamma correction we use.

double GammaCorrect(double v) { return pow(v, Gamma); }
double GammaUncorrect(double v) { return pow(v, 1.0 / Gamma); }

/* CIE C illuminant */
static const double illum[3 * 3] =
    {0.488718, 0.176204, 0.000000,
     0.310680, 0.812985, 0.0102048,
     0.200602, 0.0108109, 0.989795};
struct LabItem // CIE L*a*b* color value with C and h added.
{
    double L, a, b, C, h;

    LabItem() {}
    LabItem(double R, double G, double B) { Set(R, G, B); }
    void Set(double R, double G, double B)
    {
        const double *const i = illum;
        double X = i[0] * R + i[3] * G + i[6] * B, x = X / (i[0] + i[1] + i[2]);
        double Y = i[1] * R + i[4] * G + i[7] * B, y = Y / (i[3] + i[4] + i[5]);
        double Z = i[2] * R + i[5] * G + i[8] * B, z = Z / (i[6] + i[7] + i[8]);
        const double threshold1 = (6 * 6 * 6.0) / (29 * 29 * 29.0);
        const double threshold2 = (29 * 29.0) / (6 * 6 * 3.0);
        double x1 = (x > threshold1) ? pow(x, 1.0 / 3.0) : (threshold2 * x) + (4 / 29.0);
        double y1 = (y > threshold1) ? pow(y, 1.0 / 3.0) : (threshold2 * y) + (4 / 29.0);
        double z1 = (z > threshold1) ? pow(z, 1.0 / 3.0) : (threshold2 * z) + (4 / 29.0);
        L = (29 * 4) * y1 - (4 * 4);
        a = (500 * (x1 - y1));
        b = (200 * (y1 - z1));
        C = sqrt(a * a + b + b);
        h = atan2(b, a);
    }
    LabItem(unsigned rgb) { Set(rgb); }
    void Set(unsigned rgb)
    {
        Set((rgb >> 16) / 255.0, ((rgb >> 8) & 0xFF) / 255.0, (rgb & 0xFF) / 255.0);
    }
};

/* From the paper "The CIEDE2000 Color-Difference Formula: Implementation Notes, */
/* Supplementary Test Data, and Mathematical Observations", by */
/* Gaurav Sharma, Wencheng Wu and Edul N. Dalal, */
/* Color Res. Appl., vol. 30, no. 1, pp. 21-30, Feb. 2005. */
/* Return the CIEDE2000 Delta E color difference measure squared, for two Lab values */
static double ColorCompare(const LabItem &lab1, const LabItem &lab2)
{
#define RAD2DEG(xx) (180.0 / M_PI * (xx))
#define DEG2RAD(xx) (M_PI / 180.0 * (xx))
    /* Compute Cromanance and Hue angles */
    double C1, C2, h1, h2;
    {
        double Cab = 0.5 * (lab1.C + lab2.C);
        double Cab7 = pow(Cab, 7.0);
        double G = 0.5 * (1.0 - sqrt(Cab7 / (Cab7 + 6103515625.0)));
        double a1 = (1.0 + G) * lab1.a;
        double a2 = (1.0 + G) * lab2.a;
        C1 = sqrt(a1 * a1 + lab1.b * lab1.b);
        C2 = sqrt(a2 * a2 + lab2.b * lab2.b);

        if (C1 < 1e-9)
            h1 = 0.0;
        else
        {
            h1 = RAD2DEG(atan2(lab1.b, a1));
            if (h1 < 0.0)
                h1 += 360.0;
        }

        if (C2 < 1e-9)
            h2 = 0.0;
        else
        {
            h2 = RAD2DEG(atan2(lab2.b, a2));
            if (h2 < 0.0)
                h2 += 360.0;
        }
    }

    /* Compute delta L, C and H */
    double dL = lab2.L - lab1.L, dC = C2 - C1, dH;
    {
        double dh;
        if (C1 < 1e-9 || C2 < 1e-9)
        {
            dh = 0.0;
        }
        else
        {
            dh = h2 - h1;
            /**/ if (dh > 180.0)
                dh -= 360.0;
            else if (dh < -180.0)
                dh += 360.0;
        }

        dH = 2.0 * sqrt(C1 * C2) * sin(DEG2RAD(0.5 * dh));
    }

    double h;
    double L = 0.5 * (lab1.L + lab2.L);
    double C = 0.5 * (C1 + C2);
    if (C1 < 1e-9 || C2 < 1e-9)
    {
        h = h1 + h2;
    }
    else
    {
        h = h1 + h2;
        if (fabs(h1 - h2) > 180.0)
        {
            /**/ if (h < 360.0)
                h += 360.0;
            else if (h >= 360.0)
                h -= 360.0;
        }
        h *= 0.5;
    }
    double T = 1.0 - 0.17 * cos(DEG2RAD(h - 30.0)) + 0.24 * cos(DEG2RAD(2.0 * h)) + 0.32 * cos(DEG2RAD(3.0 * h + 6.0)) - 0.2 * cos(DEG2RAD(4.0 * h - 63.0));
    double hh = (h - 275.0) / 25.0;
    double ddeg = 30.0 * exp(-hh * hh);
    double C7 = pow(C, 7.0);
    double RC = 2.0 * sqrt(C7 / (C7 + 6103515625.0));
    double L50sq = (L - 50.0) * (L - 50.0);
    double SL = 1.0 + (0.015 * L50sq) / sqrt(20.0 + L50sq);
    double SC = 1.0 + 0.045 * C;
    double SH = 1.0 + 0.015 * C * T;
    double RT = -sin(DEG2RAD(2 * ddeg)) * RC;
    double dLsq = dL / SL, dCsq = dC / SC, dHsq = dH / SH;
    return dLsq * dLsq + dCsq * dCsq + dHsq * dHsq + RT * dCsq * dHsq;
#undef RAD2DEG
#undef DEG2RAD
}

static double ColorCompare(int r1, int g1, int b1, int r2, int g2, int b2)
{
    double luma1 = (r1 * 299 + g1 * 587 + b1 * 114) / (255.0 * 1000);
    double luma2 = (r2 * 299 + g2 * 587 + b2 * 114) / (255.0 * 1000);
    double lumadiff = luma1 - luma2;
    double diffR = (r1 - r2) / 255.0, diffG = (g1 - g2) / 255.0, diffB = (b1 - b2) / 255.0;
    return (diffR * diffR * 0.299 + diffG * diffG * 0.587 + diffB * diffB * 0.114) * 0.75 + lumadiff * lumadiff;
}

/* Palette */
static const unsigned palettesize = 16;
static const unsigned pal[palettesize] =
    {0x080000, 0x201A0B, 0x432817, 0x492910, 0x234309, 0x5D4F1E, 0x9C6B20, 0xA9220F,
     0x2B347C, 0x2B7409, 0xD0CA40, 0xE8A077, 0x6A94AB, 0xD5C4B3, 0xFCE76E, 0xFCFAE2};

/* Luminance for each palette entry, to be initialized as soon as the program begins */
static unsigned luma[palettesize];
static LabItem meta[palettesize];
static double pal_g[palettesize][3]; // Gamma-corrected palette entry

inline bool PaletteCompareLuma(unsigned index1, unsigned index2)
{
    return luma[index1] < luma[index2];
}

typedef std::vector<unsigned> MixingPlan;

MixingPlan DeviseBestMixingPlan(unsigned color, size_t limit)
{
    // Input color in RGB
    int input_rgb[3] = {(int)((color >> 16) & 0x00FF),
                        (int)((color >> 8) & 0x00FF),
                        (int)(color & 0x00FF)};
    // Input color in CIE L*a*b*
    LabItem input(color);

    std::map<unsigned, unsigned> Solution;

    // The penalty of our currently "best" solution.
    double current_penalty = -1;

    // First, find the closest color to the input color.
    // It is our seed.
    if (1)
    {
        unsigned chosen = 0;
        for (unsigned index = 0; index < palettesize; ++index)
        {
            const unsigned color = pal[index];
#if COMPARE_RGB
            unsigned r = color >> 16, g = (color >> 8) & 0xFF, b = color & 0xFF;
            double penalty = ColorCompare(
                input_rgb[0], input_rgb[1], input_rgb[2],
                r, g, b);
#else
            LabItem test_lab(color);
            double penalty = ColorCompare(input, test_lab);
#endif
            if (penalty < current_penalty || current_penalty < 0)
            {
                current_penalty = penalty;
                chosen = index;
            }
        }

        Solution[chosen] = limit;
    }

    double dbllimit = 1.0 / limit;
    while (current_penalty != 0.0)
    {
        // Find out if there is a region in Solution that
        // can be split in two for benefit.
        double best_penalty = current_penalty;
        unsigned best_splitfrom = ~0u;
        unsigned best_split_to[2] = {0, 0};

        for (std::map<unsigned, unsigned>::iterator
                 i = Solution.begin();
             i != Solution.end();
             ++i)
        {
            // if(i->second <= 1) continue;
            unsigned split_color = i->first;
            unsigned split_count = i->second;
            // Tally the other colors
            double sum[3] = {0, 0, 0};
            for (std::map<unsigned, unsigned>::iterator
                     j = Solution.begin();
                 j != Solution.end();
                 ++j)
            {
                if (j->first == split_color)
                    continue;
                sum[0] += pal_g[j->first][0] * j->second * dbllimit;
                sum[1] += pal_g[j->first][1] * j->second * dbllimit;
                sum[2] += pal_g[j->first][2] * j->second * dbllimit;
            }
            double portion1 = (split_count / 2) * dbllimit;
            double portion2 = (split_count - split_count / 2) * dbllimit;
            for (unsigned a = 0; a < palettesize; ++a)
            {
                // if(a != split_color && Solution.find(a) != Solution.end()) continue;
                unsigned firstb = 0;
                if (portion1 == portion2)
                    firstb = a + 1;
                for (unsigned b = firstb; b < palettesize; ++b)
                {
                    if (a == b)
                        continue;
                    // if(b != split_color && Solution.find(b) != Solution.end()) continue;
                    int lumadiff = int(luma[a]) - int(luma[b]);
                    if (lumadiff < 0)
                        lumadiff = -lumadiff;
                    if (lumadiff > 80000)
                        continue;

                    double test[3] =
                        {GammaUncorrect(sum[0] + pal_g[a][0] * portion1 + pal_g[b][0] * portion2),
                         GammaUncorrect(sum[1] + pal_g[a][1] * portion1 + pal_g[b][1] * portion2),
                         GammaUncorrect(sum[2] + pal_g[a][2] * portion1 + pal_g[b][2] * portion2)};
                    // Figure out if this split is better than what we had
#if COMPARE_RGB
                    double penalty = ColorCompare(
                        input_rgb[0], input_rgb[1], input_rgb[2],
                        test[0] * 255, test[1] * 255, test[2] * 255);
#else
                    LabItem test_lab(test[0], test[1], test[2]);
                    double penalty = ColorCompare(input, test_lab);
#endif
                    if (penalty < best_penalty)
                    {
                        best_penalty = penalty;
                        best_splitfrom = split_color;
                        best_split_to[0] = a;
                        best_split_to[1] = b;
                    }
                    if (portion2 == 0)
                        break;
                }
            }
        }
        if (best_penalty == current_penalty)
            break; // No better solution was found.

        std::map<unsigned, unsigned>::iterator i = Solution.find(best_splitfrom);
        unsigned split_count = i->second, split1 = split_count / 2, split2 = split_count - split1;
        Solution.erase(i);
        if (split1 > 0)
            Solution[best_split_to[0]] += split1;
        if (split2 > 0)
            Solution[best_split_to[1]] += split2;
        current_penalty = best_penalty;
    }

    // Sequence the solution.
    MixingPlan result;
    for (std::map<unsigned, unsigned>::iterator
             i = Solution.begin();
         i != Solution.end(); ++i)
    {
        result.resize(result.size() + i->second, i->first);
    }
    // Sort the colors according to luminance
    std::sort(result.begin(), result.end(), PaletteCompareLuma);
    return result;
}

int main(int argc, char **argv)
{
    if (argc != 3)
    {
        printf("Usage: %s INPUT.png OUTPUT.png\n", argv[0]);
        return 0;
    }

    FILE *fp = fopen(argv[1], "rb");
    gdImagePtr srcim = gdImageCreateFromPng(fp);
    fclose(fp);

    unsigned w = gdImageSX(srcim), h = gdImageSY(srcim);
    gdImagePtr im = gdImageCreate(w, h);
    for (unsigned c = 0; c < palettesize; ++c)
    {
        unsigned r = pal[c] >> 16, g = (pal[c] >> 8) & 0xFF, b = pal[c] & 0xFF;
        gdImageColorAllocate(im, r, g, b);
        luma[c] = r * 299 + g * 587 + b * 114;
        meta[c].Set(pal[c]);
        pal_g[c][0] = GammaCorrect(r / 255.0);
        pal_g[c][1] = GammaCorrect(g / 255.0);
        pal_g[c][2] = GammaCorrect(b / 255.0);
    }
#pragma omp parallel for
    for (unsigned y = 0; y < h; ++y)
        for (unsigned x = 0; x < w; ++x)
        {
            unsigned color = gdImageGetTrueColorPixel(srcim, x, y);
            unsigned map_value = map[(x & 7) + ((y & 7) << 3)];
            MixingPlan plan = DeviseBestMixingPlan(color, 16);
            map_value = map_value * plan.size() / 64;
            gdImageSetPixel(im, x, y, plan[map_value]);
        }
    fp = fopen(argv[2], "wb");
    gdImagePng(im, fp);
    fclose(fp);
    gdImageDestroy(im);
    gdImageDestroy(srcim);
}