/* test shit */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdlib.h>

#include <math.h>

#include <SDL_image.h>

#define MAXWIDTH 512
#define MAXHEIGHT 512

int WIDTH, HEIGHT;

static inline float get(float *src, int x, int y)
{
    return src[y * WIDTH + x];
}

static inline void put(float *src, int x, int y, float p)
{
    src[y * WIDTH + x] = p;
}

static float *new(void)
{
    return malloc(WIDTH * HEIGHT * sizeof(float));
}

static float *copy(float *src)
{
    float *dest = malloc(WIDTH * HEIGHT * sizeof(float));
    memcpy(dest, src, WIDTH * HEIGHT * sizeof(float)); 
    return dest;
}

#define N 5
#define NN ((N * 2 + 1))

static void makegauss(float mat[NN][NN], float dx, float dy)
{
    float t = 0;
    int i, j;

    for(j = 0; j < NN; j++)
        for(i = 0; i < NN; i++)
        {
            float a = dx + i - N;
            float b = dy + j - N;
            mat[i][j] = pow(M_E, - (a * a + b * b) / (2. * 1.2 * 1.2));
            t += mat[i][j];
        }

    for(j = 0; j < NN; j++)
        for(i = 0; i < NN; i++)
            mat[i][j] /= t;
}

static float *gauss(float *src, float mat[NN][NN])
{
    float *dest = new();
    int x, y, i, j;

    for(y = N; y < HEIGHT - N; y++)
        for(x = N; x < WIDTH - N; x++)
        {
            float p = 0;

            for(j = 0; j < NN; j++)
                for(i = 0; i < NN; i++)
                    p += get(src, x + i - 5, y + j - 5) * mat[i][j];

            put(dest, x, y, p);
        }

    return dest;
}

static float *fullgauss(float *src, float mat[NN][NN])
{
    float *dest = new();
    int x, y, i, j;

    for(y = 0; y < HEIGHT; y++)
        for(x = 0; x < WIDTH; x++)
        {
            float p = 0;

            for(j = 0; j < NN; j++)
                for(i = 0; i < NN; i++)
                    if(x + i >= 5 && x + i < WIDTH + 5
                       && y + j >= 5 && y + j < HEIGHT + 5)
                        p += get(src, x + i - 5, y + j - 5) * mat[i][j];

            put(dest, x, y, p);
        }

    return dest;
}

static float fulldist(float *p1, float *p2)
{
    float error = 0.;
    int x, y;

    for(y = 0; y < HEIGHT; y++)
        for(x = 0; x < WIDTH; x++)
        {
            float t = get(p1, x, y) - get(p2, x, y);
            error += t * t;
        }

    return error / (WIDTH * HEIGHT);
}

static float dist(float *p1, float *p2)
{
    float error = 0.;
    int x, y;

    for(y = N; y < HEIGHT - N; y++)
        for(x = N; x < WIDTH - N; x++)
        {
            float t = get(p1, x, y) - get(p2, x, y);
            error += t * t;
        }

    return error / ((WIDTH - N) * (HEIGHT - N));
}

static float *load(char *name)
{
    SDL_Surface *tmp, *surface;
    uint32_t *pixels;
    float *floats;
    int x, y;

    tmp = IMG_Load(name);
    if(!tmp)
        return NULL;

    WIDTH = tmp->w > MAXWIDTH ? MAXWIDTH : tmp->w;
    HEIGHT = tmp->h > MAXHEIGHT ? MAXHEIGHT : tmp->h;
    floats = malloc(WIDTH * HEIGHT * sizeof(float));
    if(!floats)
        return NULL;

    surface = SDL_CreateRGBSurface(SDL_SWSURFACE, WIDTH, HEIGHT, 32,
                                   0xff0000, 0xff00, 0xff, 0x0);
    pixels = (uint32_t *)surface->pixels;
    SDL_BlitSurface(tmp, NULL, surface, NULL);
    SDL_FreeSurface(tmp);

    for(y = 0; y < HEIGHT; y++)
        for(x = 0; x < WIDTH; x++)
        {
            int green = (pixels[y * surface->pitch / 4 + x] >> 8) & 0xff;
            put(floats, x, y, (float)green / 0xff);
        }

    return floats;
}

static void save(float *src, char *name)
{
    SDL_Surface *surface;
    uint32_t *pixels;
    int x, y;

    surface = SDL_CreateRGBSurface(SDL_SWSURFACE, WIDTH, HEIGHT, 32,
                                   0xff0000, 0xff00, 0xff, 0x0);
    pixels = (uint32_t *)surface->pixels;

    for(y = 0; y < HEIGHT; y++)
        for(x = 0; x < WIDTH; x++)
            pixels[surface->pitch / 4 * y + x]
                         = get(src, x, y) > 0.5 ? 0xffffff : 0x0;

    SDL_SaveBMP(surface, name);
}

/* Dither using error diffusion and Floyd-Steinberg-like coefficients:
         X a b
     c d e f g
     h i j k l
*/
static float *ed(float *src, int a, int b, int c, int d, int e, int f,
                             int g, int h, int i, int j, int k, int l)
{
    float *dest = new();
    float *tmp = copy(src);
    int x, y, n;

    n = a + b + c + d + e + f + g + h + i + j + k + l;

    for(y = 0; y < HEIGHT; y++)
    {
        for(x = 0; x < WIDTH; x++)
        {
            float p = get(tmp, x, y);
            float q = p > 0.5 ? 1. : 0.;
            float error = (p - q) / n;

            put(dest, x, y, q);

            if(x < WIDTH - 1)
                put(tmp, x + 1, y, get(tmp, x + 1, y) + error * a);
            if(x < WIDTH - 2)
                put(tmp, x + 2, y, get(tmp, x + 2, y) + error * b);
            if(y < HEIGHT - 1)
            {
                if(x > 0)
                {
                    put(tmp, x - 1, y + 1,
                        get(tmp, x - 1, y + 1) + error * d);
                    if(x > 1)
                        put(tmp, x - 2, y + 1,
                            get(tmp, x - 2, y + 1) + error * c);
                }
                put(tmp, x, y + 1, get(tmp, x, y + 1) + error * e);
                if(x < WIDTH - 1)
                {
                    put(tmp, x + 1, y + 1,
                        get(tmp, x + 1, y + 1) + error * f);
                    if(x < WIDTH - 2)
                        put(tmp, x + 2, y + 1,
                            get(tmp, x + 2, y + 1) + error * g);
                }
            }
            if(y < HEIGHT - 2)
            {
                if(x > 0)
                {
                    put(tmp, x - 1, y + 2,
                        get(tmp, x - 1, y + 2) + error * h);
                    if(x > 1)
                        put(tmp, x - 2, y + 2,
                            get(tmp, x - 2, y + 2) + error * i);
                }
                put(tmp, x, y + 2, get(tmp, x, y + 2) + error * j);
                if(x < WIDTH - 1)
                {
                    put(tmp, x + 1, y + 2,
                        get(tmp, x + 1, y + 2) + error * k);
                    if(x < WIDTH - 2)
                        put(tmp, x + 2, y + 2,
                            get(tmp, x + 2, y + 2) + error * l);
                }
            }
        }
    }

    free(tmp);

    return dest;
}

/* XXX */
static float *dbs(float *src, float *orig, float dx, float dy)
{
    float mat[NN][NN];
    float *dest, *tmp, *tmp2;
    float error;

    makegauss(mat, 0., 0.);
    tmp = fullgauss(src, mat);

    makegauss(mat, dx, dy);
    dest = copy(orig);
    tmp2 = fullgauss(dest, mat);

    error = dist(tmp, tmp2);

    for(;;)
    {
        int changes = 0;
        int x, y, i, j, n;

        for(y = 0; y < HEIGHT; y++)
        for(x = 0; x < WIDTH; x++)
        {
            float d, d2, e, best = 0.;
            int opx = -1, opy = -1;

            d = get(dest, x, y);

            /* Compute the effect of a toggle */
            e = 0.;
            for(j = -N; j < N + 1; j++)
            {
                if(y + j < 0 || y + j >= HEIGHT)
                    continue;

                for(i = -N; i < N + 1; i++)
                {
                    float m, p, q1, q2;

                    if(x + i < 0 || x + i >= WIDTH)
                        continue;
                    
                    m = mat[i + N][j + N];
                    p = get(tmp, x + i, y + j);
                    q1 = get(tmp2, x + i, y + j);
                    q2 = q1 - m * d + m * (1. - d);
                    e += (q1 - p) * (q1 - p) - (q2 - p) * (q2 - p);
                }
            }
            if(e > best)
            {
                best = e;
                opx = opy = 0;
            }

            /* Compute the effect of swaps */
            for(n = 0; n < 8; n++)
            {
                static int const step[] =
                    { 0, 1, 0, -1, -1, 0, 1, 0, -1, -1, -1, 1, 1, -1, 1, 1 };
                int idx = step[n * 2], idy = step[n * 2 + 1];
                if(y + idy < 0 || y + idy >= HEIGHT
                     || x + idx < 0 || x + idx >= WIDTH)
                    continue;
                d2 = get(dest, x + idx, y + idy);
                if(d2 == d)
                    continue;
                e = 0.;
                for(j = -N; j < N + 1; j++)
                {
                    if(y + j < 0 || y + j >= HEIGHT)
                        continue;
                    if(j - idy + N < 0 || j - idy + N >= NN)
                        continue;
                    for(i = -N; i < N + 1; i++)
                    {
                        float ma, mb, p, q1, q2;
                        if(x + i < 0 || x + i >= WIDTH)
                            continue;
                        if(i - idx + N < 0 || i - idx + N >= NN)
                            continue;
                        ma = mat[i + N][j + N];
                        mb = mat[i - idx + N][j - idy + N];
                        p = get(tmp, x + i, y + j);
                        q1 = get(tmp2, x + i, y + j);
                        q2 = q1 - ma * d + ma * d2 - mb * d2 + mb * d;
                        e += (q1 - p) * (q1 - p) - (q2 - p) * (q2 - p);
                    }
                }
                if(e > best)
                {
                    best = e;
                    opx = idx;
                    opy = idy;
                }
            }

            /* Apply the change if interesting */
            if(best <= 0.)
                continue;
            if(opx || opy)
            {
                d2 = get(dest, x + opx, y + opy);
                put(dest, x + opx, y + opy, d);
            }
            else
                d2 = 1. - d;
            put(dest, x, y, d2);
            for(j = -N; j < N + 1; j++)
            for(i = -N; i < N + 1; i++)
            {
                float m = mat[i + N][j + N];
                if(y + j >= 0 && y + j < HEIGHT
                    && x + i >= 0 && x + i < WIDTH)
                {
                    float t = get(tmp2, x + i, y + j);
                    put(tmp2, x + i, y + j, t + m * (d2 - d));
                }
                if((opx || opy) && y + opy + j >= 0 && y + opy + j < HEIGHT
                                && x + opx + i >= 0 && x + opx + i < WIDTH)
                {
                    float t = get(tmp2, x + opx + i, y + opy + j);
                    put(tmp2, x + opx + i, y + opy + j, t + m * (d - d2));
                }
            }

            changes++;
        }

        printf("did %i changes\n", changes);

        if(changes == 0)
            break;
    }

    free(tmp);
    free(tmp2);

    return dest;
}

static void study(float *src, float *dest, float precision)
{
#   define Z 3
    float mat[NN][NN];
    float *tmp, *tmp2;
    float e0, best, fx = -1., fy = -1., step = 2.;
    int dx, dy;

    makegauss(mat, 0., 0.);
    tmp = gauss(src, mat);

    makegauss(mat, fx, fy);
    tmp2 = gauss(dest, mat);
    e0 = dist(tmp, tmp2);
    free(tmp2);

    while(step > precision)
    {
        float bfx = 0., bfy = 0., e;

        best = 9999.;

        for(dy = 0; dy <= Z; dy++)
            for(dx = 0; dx <= Z; dx++)
            {
                makegauss(mat, fx + step * dx / Z, fy + step * dy / Z);
                tmp2 = gauss(dest, mat);
                e = dist(tmp, tmp2);
                free(tmp2);
                if(e < best)
                {
                    best = e;
                    bfx = fx + step * dx / Z;
                    bfy = fy + step * dy / Z;
                }
            }

        fx = bfx - step / Z;
        fy = bfy - step / Z;
        step = step * 2 / Z;
    }

    free(tmp);

    printf("%g -> %g for %g %g\n", 1000 * e0, 1000 * best, fx, fy);
    fflush(stdout);
}

int main(int argc, char *argv[])
{
    float *src, *dest, *tmp, *tmp2;
    int a, b, c, d;

    if(argc < 2)
        return 1;

    src = load(argv[1]);
    if(!src)
        return 2;

#if 0
    tmp = ed(src, 7, 0, 0, 3, 5, 1, 0, 0, 0, 0, 0, 0);
    //dest = dbs(src, tmp, 0., 0.);
    dest = dbs(src, tmp, 0.20, 0.30);
    //dest = dbs(src, tmp, 0.158718, 0.283089);
    //dest = copy(tmp);
    free(tmp);
    study(src, dest, 0.00001);
    save(dest, "output.bmp");
    free(dest);
#endif

#if 1
    //dest = ed(src, 5, 0, 0, 3, 5, 3, 0, 0, 0, 0, 0, 0);
    //dest = ed(src, 7, 0, 0, 3, 5, 1, 0, 0, 0, 0, 0, 0);
    dest = ed(src, 7, 5, 3, 5, 7, 5, 3, 1, 3, 5, 3, 1);
    //dest = ed(src, 7, 0, 1, 3, 5, 0, 0, 0, 0, 0, 0, 0);
    //dest = ed(src, 2, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0);
    printf("%s: ", argv[1]);
    study(src, dest, 0.001);
    save(dest, "output.bmp");
    free(dest);
#endif

#if 0
#   define STEP 32
    dest = ed(src, 7, 0, 0, 3, 5, 1, 0, 0, 0, 0, 0, 0);
    tmp = gauss(src, 0., 0.);
    for(dy = 0; dy < STEP; dy++)
    {
        for(dx = 0; dx < STEP; dx++)
        {
            float fy = 2. / STEP * (dy - STEP / 2.);
            float fx = 2. / STEP * (dx - STEP / 2.);

            tmp2 = gauss(dest, fx, fy);
            printf("%g %g %g\n", fy, fx, 1000. * dist(tmp, tmp2));
            fflush(stdout);
            free(tmp2);
        }
        printf("\n");
    }

    save(dest, "output.bmp");
#endif

#if 0
    for(a = 0; a <= 16; a++)
    for(b = 0; b <= 16; b++)
    for(c = 0; c <= 16; c++)
    for(d = 0; d <= 16; d++)
    {
        float *tmp;
        if(a + b + c + d != 16)
            continue;
#if 0
        if(b > c) continue;
        if(d > c) continue;
        if(d > a) continue;
#endif
        printf("%02d %02d %02d %02d: ", a, b, c, d);
        tmp = ed(src, a, 0, 0, b, c, d, 0, 0, 0, 0, 0, 0);
        study(src, tmp, 0.01);
        free(tmp);
    }
#endif

#if 0
    tmp = gauss(src, 0., 0.);
    for(a = 0; a < 16; a++)
    for(b = 0; b < 16; b++)
    for(c = 0; c < 16; c++)
    for(d = 0; d < 16; d++)
    {
        if(a + b + c + d != 16)
            continue;
        dest = ed(src, a, 0, 0, b, c, d, 0, 0, 0, 0, 0, 0);
        tmp2 = gauss(dest, 0., 0.);
        printf("%.5g: (%02d %02d %02d %02d)\n",
               1000. * dist(tmp, tmp2), a, b, c, d);
        free(dest);
        free(tmp2);
    }

    save(dest, "output.bmp");
#endif

    return 0;
}