/*
 Yliluoma's ordered dithering algorithm 1

 Usage:
     ./yodither1.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 yodither1.cpp -fopenmp -o yodither1.exe -static -lstdc++ -lgcc -lwinpthread -lm

 */

#include <stdio.h>
#include <getopt.h>
#include <map>

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

// Compare the difference of two RGB values, weigh by CCIR 601 luminosity:
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 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;
}

double EvaluateMixingError(int r, int g, int b,    // Desired color
                           int r0, int g0, int b0, // Mathematical mix product
                           int r1, int g1, int b1, // Mix component 1
                           int r2, int g2, int b2, // Mix component 2
                           double ratio)           // Mixing ratio
{
    return ColorCompare(r, g, b, r0, g0, b0) + ColorCompare(r1, g1, b1, r2, g2, b2) * 0.1 * (fabs(ratio - 0.5) + 0.5);
}

unsigned int DeviseBestMixingPlan(unsigned color, unsigned int *pal,
                                  double map_value, double drange, int len_pal)
{
    unsigned int r, g, b;
    r = (color >> 16) & 0x0ff;
    g = (color >> 8) & 0x0ff;
    b = color & 0x0ff;

    unsigned int r_colors[2] = {0, 0};

    /* 0 = always index1, 1 = always index2, 0.5 = 50% of both */
    double r_ratio = 0.5;

    double least_penalty = 1e99;

    /*
       Loop through every unique combination of two colors from the palette,
       and through each possible way to mix those two colors.
       They can be mixed in exactly 64 ways, when the threshold matrix is 8x8.
     */

    for (int i1 = 0; i1 < len_pal; i1++)
        for (int i2 = i1; i2 < len_pal; i2++)
            for (int ratio = 0; ratio < drange;  ratio++)
            {
                if (i1 == i2 && ratio != 0)
                    break;

                // Determine the two component colors
                unsigned int color1, r1, g1, b1;
                unsigned int color2, r2, g2, b2;

                color1 = pal[i1];
                color2 = pal[i2];

                r1 = (color1 >> 16) & 0x0ff;
                g1 = (color1 >> 8) & 0x0ff;
                b1 = color1 & 0x0ff;

                r2 = (color2 >> 16) & 0x0ff;
                g2 = (color2 >> 8) & 0x0ff;
                b2 = color2 & 0x0ff;

                // Determine what mixing them in this proportion will produce
                unsigned int r0, g0, b0;
                r0 = r1 + ratio * int(r2 - r1) / drange;
                g0 = g1 + ratio * int(g2 - g1) / drange;
                b0 = b1 + ratio * int(b2 - b1) / drange;

                // Determine how well that matches what we want to accomplish
                double penalty = EvaluateMixingError(r, g, b, r0, g0, b0,
                                                     r1, g1, b1, r2, g2, b2,
                                                     ratio / drange);

                if (penalty < least_penalty)
                {
                    // Keep the result that has the smallest error
                    least_penalty = penalty;
                    r_colors[0] = i1;
                    r_colors[1] = i2;
                    r_ratio = ratio / drange;
                }
            }

    return pal[r_colors[(map_value / drange) < r_ratio ? 1 : 0]];
}

unsigned int DeviseBestMixingPlanTri(unsigned color, unsigned int *pal,
                                     double map_value, double drange, int len_pal,
                                     int submap_value)
{
    unsigned int r, g, b;
    r = (color >> 16) & 0x0ff;
    g = (color >> 8) & 0x0ff;
    b = color & 0x0ff;

    int r_colors[4] = {0, 0, 0, 0};

    /* 0 = always index1, 1 = always index2, 0.5 = 50% of both */
    /* 4 = special three or four-color dither */
    double r_ratio = 0.5;

    double least_penalty = 1e99;
    for (int i1 = 0; i1 < len_pal; ++i1)
        for (int i2 = i1; i2 < len_pal; ++i2)
        // for(int ratio=0; ratio<drange; ++ratio)
        {
            // Determine the two component colors
            unsigned int color1, r1, g1, b1;
            unsigned int color2, r2, g2, b2;
            color1 = pal[i1];
            color2 = pal[i2];

            r1 = (color1 >> 16) & 0x0ff;
            g1 = (color1 >> 8) & 0x0ff;
            b1 = color1 & 0x0ff;

            r2 = (color2 >> 16) & 0x0ff;
            g2 = (color2 >> 8) & 0x0ff;
            b2 = color2 & 0x0ff;

            int ratio = drange / 2;
            if (color1 != color2)
            {
                /*
                   Determine the ratio of mixing for each channel.
                     solve(r1 + ratio*(r2-r1)/64 = r, ratio)
                   Take a weighed average of these three ratios according
                   to the perceived luminosity of each channel (according to CCIR 601).
                 */
                ratio = ((r2 != r1 ? 299 * 64 * int(r - r1) / int(r2 - r1) : 0) + (g2 != g1 ? 587 * 64 * int(g - g1) / int(g2 - g1) : 0) + (b1 != b2 ? 114 * 64 * int(b - b1) / int(b2 - b1) : 0)) / ((r2 != r1 ? 299 : 0) + (g2 != g1 ? 587 : 0) + (b2 != b1 ? 114 : 0));
                if (ratio < 0)
                    ratio = 0;
                else if (ratio > (drange - 1))
                    ratio = (drange - 1);
            }

            // Determine what mixing them in this proportion will produce
            unsigned int r0, g0, b0;
            r0 = r1 + ratio * int(r2 - r1) / drange;
            g0 = g1 + ratio * int(g2 - g1) / drange;
            b0 = b1 + ratio * int(b2 - b1) / drange;

            double penalty = EvaluateMixingError(r, g, b, r0, g0, b0,
                                                 r1, g1, b1, r2, g2, b2,
                                                 ratio / drange);

            if (penalty < least_penalty)
            {
                least_penalty = penalty;
                r_colors[0] = i1;
                r_colors[1] = i2;
                r_ratio = ratio / drange;
            }

            if (i1 != i2)
                for (int i3 = 0; i3 < len_pal; ++i3)
                {
                    if (i3 == i2 || i3 == i1)
                        continue;
                    // 50% index3, 25% index2, 25% index1
                    unsigned int color3, r3, g3, b3;
                    color3 = pal[i3];
                    r3 = (color3 >> 16) & 0x0ff;
                    g3 = (color3 >> 8) & 0x0ff;
                    b3 = color3 & 0x0ff;
                    r0 = (r1 + r2 + r3 * 2) / 4;
                    g0 = (g1 + g2 + g3 * 2) / 4;
                    b0 = (b1 + b2 + b3 * 2) / 4;

                    penalty = ColorCompare(r, g, b, r0, g0, b0) + ColorCompare(r1, g1, b1, r2, g2, b2) * 0.025 + ColorCompare((r1 + g1) / 2, (g1 + g2) / 2, (b1 + b2) / 2, r3, g3, b3) * 0.025;

                    if (penalty < least_penalty)
                    {
                        least_penalty = penalty;
                        r_colors[0] = i3; // (0,0) index3 occurs twice
                        r_colors[1] = i1; // (0,1)
                        r_colors[2] = i2; // (1,0)
                        r_colors[3] = i3; // (1,1)
                        r_ratio = 4.0;
                    }
                }
        }

    if (r_ratio == 4.0)
    {
        // Tri-tone or quad-tone dithering
        return pal[r_colors[submap_value]];
    }

    return pal[r_colors[(map_value / drange) < r_ratio ? 1 : 0]];
}

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] [-t] [-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("  -t, --trimode            Enable tri-mode.\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 trimode = 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'},
        {"trimode", no_argument, NULL, 't'},
        {"verbose", no_argument, NULL, 'v'},
        {0, 0, 0, 0}};

    int opt;
    int longindex;
    while ((opt = getopt_long(argc, argv, "hi:o:d:p:tv", 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 't': // enable tri-mode
            trimode = 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 *my_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);

        my_pal = load_palette(palfile, &len_pal);
        if (my_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 = my_pal[i];
            unsigned int r = (col >> 16) & 0x0ff;
            unsigned int g = (col >> 8) & 0x0ff;
            unsigned int 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;
    }

#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;
            double map_value = dmap[(x % dw) + (y % dh) * dw];
            int submap_value = (y & 1) * 2 + (x & 1);
            unsigned int c;

            if (trimode)
                c = DeviseBestMixingPlanTri(color, my_pal, map_value,
                                            drange, len_pal, submap_value);
            else
                c = DeviseBestMixingPlan(color, my_pal, map_value,
                                         drange, len_pal);

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

    return 0;
}