/*
 Yliluoma's ordered dithering algorithm 3 a.

 Usage:
     ./yodither3.exe -i INPUT.png -o OUTPUT.png [-d N] [-p PALETTE] [--help]

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

 Compile:
     g++ -O3 yodither3.cpp -fopenmp -o yodither3.exe -static -lstdc++ -lgcc -lwinpthread -lm

 */

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

#define _USE_MATH_DEFINES
#include <math.h>

/*
 use stb library

 nothings/stb: stb single-file public domain libraries for C/C++
 https://github.com/nothings/stb
 */
#define STB_IMAGE_STATIC
#define STB_IMAGE_IMPLEMENTATION
#include "stb_image.h"

#define STB_IMAGE_WRITE_STATIC
#define STB_IMAGE_WRITE_IMPLEMENTATION
#include "stb_image_write.h"

#define CBUF_N 1024
#define TMP_PAL_NUM 1024

/* 2x2 threshold map */
static const unsigned char map2x2[2 * 2] = {
    0, 2,
    3, 1};

/* 4x4 threshold map */
static const unsigned char map4x4[4 * 4] = {
    0, 8, 2, 10,
    12, 4, 14, 6,
    3, 11, 1, 9,
    15, 7, 13, 5};

/* 8x8 threshold map */
static const unsigned char map8x8[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};

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

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};

// CIE L*a*b* color value with C and h added.
struct LabItem
{
    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);
        C = sqrt(a * a + b * b);
        h = atan2(b, a);
    }

    LabItem(unsigned rgb)
    {
        Set(rgb);
    }

    void Set(unsigned rgb)
    {
        Set(((rgb >> 16) & 0x0ff) / 255.0,
            ((rgb >> 8) & 0x0ff) / 255.0,
            (rgb & 0x0ff) / 255.0);
    }
};

/* Luminance for each palette entry, to be initialized as soon as the program begins */
static unsigned int *luma;

static double rad2deg(double x) { return x * 180.0 / M_PI; }
static double deg2rad(double x) { return x * M_PI / 180.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 ColorCompareCiede2000(const LabItem &lab1, const LabItem &lab2)
{

    /* 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;
}

// Compare the difference of two RGB values, weigh by CCIR 601 luminosity:
static double ColorCompareCcir601(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 ld = luma1 - luma2;
    double r = (r1 - r2) / 255.0;
    double g = (g1 - g2) / 255.0;
    double b = (b1 - b2) / 255.0;
    return (r * r * 0.299 + g * g * 0.587 + b * b * 0.114) * 0.75 + ld * ld;
}

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

unsigned int DeviseBestMixingPlan(unsigned int color,
                                  unsigned int *pal, int len_pal,
                                  unsigned int *luma,
                                  unsigned char map_value,
                                  int drange, int ciede2000,
                                  LabItem *meta, double pal_g[][3])
{
    double limit = double(len_pal);

    // Input color in RGB
    int src[3] = {(int)((color >> 16) & 0x0ff),
                  (int)((color >> 8) & 0x0ff),
                  (int)(color & 0x0ff)};

    // 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 int chosen = 0;
        for (unsigned int index = 0; index < len_pal; ++index)
        {
            const unsigned col = pal[index];
            double penalty;

            if (ciede2000)
            {
                // CIEDE2000
                LabItem test_lab(col);
                penalty = ColorCompareCiede2000(input, test_lab);
            }
            else
            {
                unsigned int r, g, b;
                r = (col >> 16) & 0x0ff;
                g = (col >> 8) & 0x0ff;
                b = col & 0x0ff;
                penalty = ColorCompareCcir601(src[0], src[1], src[2],
                                              r, g, b);
            }

            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 < len_pal; ++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 < len_pal; ++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 t[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

                    double penalty;
                    if (ciede2000)
                    {
                        // CIEDE2000
                        LabItem test_lab(t[0], t[1], t[2]);
                        penalty = ColorCompareCiede2000(input, test_lab);
                    }
                    else
                    {
                        penalty = ColorCompareCcir601(src[0], src[1], src[2],
                                                      t[0] * 255, t[1] * 255, t[2] * 255);
                    }

                    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 int split_count = i->second;
        unsigned int split1 = split_count / 2;
        unsigned int 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.
    std::vector<unsigned int> 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);

    int idx = map_value * result.size() / drange;
    return pal[result[idx]];
}

unsigned int *load_palette_from_gpl(char *infile, int *len_pal)
{
    unsigned int *pal = nullptr;
    unsigned int tmp_pal[TMP_PAL_NUM];
    FILE *fp;
    char str[CBUF_N];
    unsigned int r, g, b;
    int idx = 0;

    for (int i = 0; i < TMP_PAL_NUM; i++)
        tmp_pal[i] = 0;

    fp = fopen(infile, "r");
    if (fp == NULL)
    {
        printf("Error : Can not open %s\n", infile);
        *len_pal = 0;
        return nullptr;
    }

    while (fgets(str, CBUF_N, fp) != NULL)
    {
        str[strlen(str) - 1] = '\0';
        if (strncmp("GIMP Palette", str, 12) == 0)
            continue;
        if (strncmp("Name:", str, 5) == 0)
            continue;
        if (strncmp("Columns", str, 7) == 0)
            continue;
        if (strncmp("#", str, 1) == 0)
            continue;

        sscanf(str, "%d %d %d", &r, &g, &b);
        // printf("%s", str);
        // printf("\t\tRGB : (%d, %d, %d)\n", r, g, b);

        tmp_pal[idx] = (r << 16) | (g << 8) | b;
        idx++;
    }
    fclose(fp);

    *len_pal = idx;
    pal = new unsigned int[idx];
    for (int i = 0; i < idx; i++)
        pal[i] = tmp_pal[i];

    return pal;
}

unsigned int *load_palette_from_img(char *infile, int *len_pal)
{
    unsigned int *pal = nullptr;
    std::map<unsigned int, unsigned int> pal_map;

    // load image
    unsigned char *pixels;
    int w;
    int h;
    int bpp;

    pixels = stbi_load(infile, &w, &h, &bpp, 0);
    if (pixels == NULL)
    {
        printf("Error : Can not load %s\n", infile);
        *len_pal = 0;
        return nullptr;
    }

    for (int y = 0; y < h; y++)
    {
        int yofs = y * w * bpp;
        for (int x = 0; x < w; x++)
        {
            int i = x * bpp + yofs;
            unsigned int col = (pixels[i + 0] << 16) | (pixels[i + 1] << 8) | pixels[i + 2];
            pal_map[col] += 1;
        }
    }
    stbi_image_free(pixels);

    int n = pal_map.size();

    // pal = (unsigned int *)malloc(sizeof(unsigned int) * n);
    pal = new unsigned int[n];

    int j = 0;
    for (auto i = pal_map.begin(); i != pal_map.end(); i++)
    {
        pal[j] = i->first;
        j++;
    }

    // printf("Count palette = %d\n", j);
    *len_pal = j;
    return pal;
}

unsigned int *load_palette(char *infile, int *len_pal)
{
    unsigned int *pal = nullptr;
    int n;

    char *ext = strrchr(infile, '.');
    if (stricmp(".gpl", ext) == 0)
    {
        // load from .gpl (GIMP palette file)
        pal = load_palette_from_gpl(infile, &n);
    }
    else
    {
        // load from .png
        pal = load_palette_from_img(infile, &n);
    }
    *len_pal = n;
    return pal;
}

unsigned char *set_dither_map(int dither, int *dw, int *dh, int *drange, int verbose)
{
    unsigned char *dmap = NULL;

    switch (dither)
    {
    case 2:
        dmap = (unsigned char *)map2x2;
        *dw = 2;
        *dh = 2;
        *drange = 4;
        if (verbose)
            printf("Select Dither map : 2x2\n");
        break;
    case 4:
        dmap = (unsigned char *)map4x4;
        *dw = 4;
        *dh = 4;
        *drange = 16;
        if (verbose)
            printf("Select Dither map : 4x4\n");
        break;
    case 8:
        dmap = (unsigned char *)map8x8;
        *dw = 8;
        *dh = 8;
        *drange = 64;
        if (verbose)
            printf("Select Dither map : 8x8\n");
        break;
    default:
        break;
    }

    return dmap;
}

void usage(char *name)
{
    printf("Usage:\n");
    printf("  %s -i IN.png -o OUT.png [-d N] [-p PALETTE] [-c] [-v] [-h]\n", name);
    printf("\n");
    printf("  -i FILE, --input FILE    input png image file.\n");
    printf("  -o FILE, --output FILE   output png image file.\n");
    printf("  -d N, --dither N         dither map. N = 2/4/8(2x2,4x4,8x8) [default=8]\n");
    printf("  -p FILE, --palette FILE  Set palette (.gpl or .png)\n");
    printf("  -c, --ciede2000          Enable CIEDE2000\n");
    printf("  -v, --verbose            Print verbose\n");
    printf("  -h, --help               print this message\n");
    printf("\n");
}

int main(int argc, char **argv)
{
    // parse option
    char *infile = NULL;
    char *outfile = NULL;
    char *palfile = NULL;
    int dither = 8;
    int ciede2000 = 0;
    int verbose = 0;

    struct option longopts[] = {
        {"help", no_argument, NULL, 'h'},
        {"input", required_argument, NULL, 'i'},
        {"output", required_argument, NULL, 'o'},
        {"dither", required_argument, NULL, 'd'},
        {"palette", required_argument, NULL, 'p'},
        {"ciede2000", no_argument, NULL, 'c'},
        {"verbose", no_argument, NULL, 'v'},
        {0, 0, 0, 0}};

    int opt;
    int longindex;
    while ((opt = getopt_long(argc, argv, "hi:o:d:p:cv", longopts, &longindex)) != -1)
    {
        switch (opt)
        {
        case 'i': // input image filename
            infile = optarg;
            break;

        case 'o': // output image filename
            outfile = optarg;
            break;

        case 'p': // palette filename
            palfile = optarg;
            break;

        case 'd': // dither map size, 2 or 4 or 8
            switch (optarg[0])
            {
            case '2':
                dither = 2;
                break;

            case '4':
                dither = 4;
                break;

            case '8':
                dither = 8;
                break;

            default:
                printf("Error : Unknown dither mode = %s\n\n", optarg);
                dither = -1;
                break;
            }
            break;

        case 'c': // enable CIEDE2000
            ciede2000 = 1;
            break;

        case 'v': // enable verbose
            verbose = 1;
            break;

        case 'h': // help
            usage(argv[0]);
            return 0;

        default:
            break;
        }
    }

    if (optind < argc)
    {
        for (; optind < argc; optind++)
            printf("Unknown option : %s\n", argv[optind]);
        usage(argv[0]);
        return 1;
    }

    if (infile == NULL || outfile == NULL || dither == -1)
    {
        usage(argv[0]);
        return 1;
    }

    // ----------------------------------------
    // set dither map
    unsigned char *dmap;
    int dw, dh;
    int drange;

    dmap = set_dither_map(dither, &dw, &dh, &drange, verbose);
    if (dmap == NULL)
    {
        printf("Error : Unknown dither mode = %d\n", dither);
        return 1;
    }

    // ----------------------------------------
    // set pallete
    unsigned int *pal = def_pal;
    int len_pal = 16;

    if (palfile == NULL)
    {
        // use default palette
        if (verbose)
            printf("Use default palette.\n");
    }
    else
    {
        // load palette file
        if (verbose)
            printf("Set palette file : %s\n", palfile);

        pal = load_palette(palfile, &len_pal);
        if (pal == nullptr)
        {
            printf("Error : Load failure palette [%s]\n", palfile);
            return 1;
        }
    }

    if (verbose)
    {
        // print palette information
        printf("Palette length : %d\n", len_pal);
        printf("Palette data : ");
        for (int i = 0; i < len_pal; i++)
        {
            unsigned int col, r, g, b;
            col = pal[i];
            r = (col >> 16) & 0x0ff;
            g = (col >> 8) & 0x0ff;
            b = col & 0x0ff;
            printf("{%d, %d, %d}, ", r, g, b);
        }
        printf("\n");
    }

    // ----------------------------------------
    // load image
    unsigned char *pixels;
    int w;
    int h;
    int bpp;

    pixels = stbi_load(infile, &w, &h, &bpp, 0);
    if (pixels == NULL)
    {
        printf("Error : Not load %s\n", infile);
        return 1;
    }

    // ----------------------------------------
    // calc Luminance value
    luma = new unsigned int[len_pal];
    LabItem *meta = new LabItem[len_pal];

    // Gamma-corrected palette entry
    double pal_g[len_pal][3];

    for (int i = 0; i < len_pal; i++)
    {
        unsigned int col, r, g, b;
        col = pal[i];
        r = (col >> 16) & 0x0ff;
        g = (col >> 8) & 0x0ff;
        b = col & 0x0ff;
        luma[i] = r * 299 + g * 587 + b * 114;

        meta[i].Set(col);
        pal_g[i][0] = GammaCorrect(double(r) / 255.0);
        pal_g[i][1] = GammaCorrect(double(g) / 255.0);
        pal_g[i][2] = GammaCorrect(double(b) / 255.0);
    }

#pragma omp parallel for
    for (int y = 0; y < h; ++y)
        for (int x = 0; x < w; ++x)
        {
            int idx = (x + y * w) * bpp;
            unsigned char r = pixels[idx + 0];
            unsigned char g = pixels[idx + 1];
            unsigned char b = pixels[idx + 2];
            unsigned int color = (r << 16) | (g << 8) | b;
            unsigned char map_value = dmap[(x % dw) + (y % dh) * dw];

            unsigned int c = DeviseBestMixingPlan(color, pal, len_pal,
                                                  luma, map_value, drange,
                                                  ciede2000, meta, pal_g);

            pixels[idx + 0] = (c >> 16) & 0x0ff;
            pixels[idx + 1] = (c >> 8) & 0x0ff;
            pixels[idx + 2] = c & 0x0ff;
        }

    // write image
    int fg = stbi_write_png(outfile, w, h, bpp, pixels, 0);
    if (!fg)
    {
        printf("Error : Not save %s\n", outfile);
    }

    stbi_image_free(pixels);
    delete luma;
    delete meta;

    return 0;
}