/*
 Pattern dither (Adobe Photoshop like)

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

 */

#include <stdio.h>
#include <getopt.h>
#include <math.h>
#include <map>
#include <algorithm> /* For std::sort() */

/*
 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 (note: the patented pattern dithering algorithm uses 4x4) */
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 pal[16] =
    {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 int *luma;

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

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

unsigned int DeviseBestMixingPlan(unsigned int color, const unsigned int *pal,
                                  unsigned int *luma, unsigned char map_value,
                                  int drange, int len_pal)
{
    // unsigned int *rcolors = (unsigned int *)malloc(sizeof(unsigned int) * drange);
    unsigned int *rcolors = new unsigned int[drange];
    for (int i = 0; i < drange; i++)
        rcolors[i] = 0;

    const int src[3] = {(int)((color >> 16) & 0x0ff),
                        (int)((color >> 8) & 0x0ff),
                        (int)(color & 0x0ff)};

    const double X = 0.09; // Error multiplier
    int e[3] = {0, 0, 0};  // Error accumulator

    for (unsigned c = 0; c < drange; ++c)
    {
        // Current temporary value
        int t[3] = {(int)(src[0] + e[0] * X),
                    (int)(src[1] + e[1] * X),
                    (int)(src[2] + e[2] * X)};

        // Clamp it in the allowed RGB range
        if (t[0] < 0)
            t[0] = 0;
        else if (t[0] > 255)
            t[0] = 255;

        if (t[1] < 0)
            t[1] = 0;
        else if (t[1] > 255)
            t[1] = 255;

        if (t[2] < 0)
            t[2] = 0;
        else if (t[2] > 255)
            t[2] = 255;

        // Find the closest color from the palette
        double least_penalty = 1e99;
        unsigned chosen = c % len_pal;
        for (int i = 0; i < len_pal; i++)
        {
            const unsigned int color = pal[i];
            const int pc[3] = {(int)((color >> 16) & 0x0ff),
                               (int)((color >> 8) & 0x0ff),
                               (int)(color & 0x0ff)};
            double penalty = ColorCompare(pc[0], pc[1], pc[2], t[0], t[1], t[2]);
            if (penalty < least_penalty)
            {
                least_penalty = penalty;
                chosen = i;
            }
        }
        // Add it to candidates and update the error
        rcolors[c] = chosen;
        unsigned color = pal[chosen];
        const int pc[3] = {(int)((color >> 16) & 0x0ff),
                           (int)((color >> 8) & 0x0ff),
                           (int)(color & 0x0ff)};
        e[0] += src[0] - pc[0];
        e[1] += src[1] - pc[1];
        e[2] += src[2] - pc[2];
    }

    // Sort the colors according to luminance
    std::sort(rcolors, rcolors + drange, PaletteCompareLuma);

    unsigned int rcol = pal[rcolors[map_value]];
    // free(rcolors);
    delete rcolors;

    return rcol;
}

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] [-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=4]\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 = 4;
    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'},
        {"verbose", no_argument, NULL, 'v'},
        {0, 0, 0, 0}};

    int opt;
    int longindex;
    while ((opt = getopt_long(argc, argv, "hi:o:d:p:v", 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 '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 = 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;
    }

    // ----------------------------------------
    // calc Luminance value
    luma = new unsigned int[len_pal];
    for (int i = 0; i < len_pal; i++)
    {
        unsigned int r = (my_pal[i] >> 16) & 0x0FF;
        unsigned int g = (my_pal[i] >> 8) & 0x0FF;
        unsigned int b = my_pal[i] & 0x0FF;
        luma[i] = r * 299 + g * 587 + b * 114;
    }

#pragma omp parallel for
    for (unsigned y = 0; y < h; ++y)
        for (unsigned 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, my_pal, luma, 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);
    delete luma;

    return 0;
}