#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <stdbool.h>
#include <time.h>
#include <math.h>


struct pixel {
    uint8_t R;
    uint8_t G;
    uint8_t B;
};
typedef struct pixel pixel;

struct dpixel {
    double R;
    double G;
    double B;
};
typedef struct dpixel dpixel;

dpixel dpixel_of_pixel(pixel p)
{
    return (dpixel){.R=(double)p.R, .G=(double)p.G, .B=(double)p.B};
}

pixel pixel_of_dpixel(dpixel p)
{
    return (pixel){.R=floor(p.R), .G=floor(p.G), .B=floor(p.B)};
}

struct image {
    int height;
    int width;
    pixel** bitmap;
};
typedef struct image image;


image* create_image(int height, int width)
{
    image* img = malloc(sizeof(image));
    img->height = height;
    img->width = width;

    img->bitmap = malloc(sizeof(pixel*) * height);
    for (int i=0; i<height; i++)
    {
        img->bitmap[i] = malloc(sizeof(pixel) * width);
    }

    return img;
}

void destroy_image(image *img)
{
    for (int i=0; i<img->height; i++)
    {
        free(img->bitmap[i]);
    }
    free(img->bitmap);
    free(img);
}

int read_integer(FILE* file, int pos, int size)
{
    int n=0;
    fseek(file, pos, SEEK_SET);
    for (int i=0; i<size; i++)
    {
        n += (getc(file)) << (i*8);
    }
    return n;
}

image* load_bmp(char* path)
{
    FILE* file = fopen(path, "rb");

    if (fgetc(file) != 0x42 || fgetc(file) != 0x4d)
    {
        fprintf(stderr, "Did not read a BMP file\n");
        return NULL;
    }
    printf("Read a BMP file\n");
    
    int width = read_integer(file, 0x12, 4);
    int height = read_integer(file, 0x16, 4);

    printf("Image size : w=%d, h=%d\n", width, height);

    image *img = create_image(height, width);
    int offset = read_integer(file, 0x0A, 4);

    int rowsize = width * 3;

    for (int i=0; i<height; i++)
    {
        for (int j=0; j<width; j++)
        {
            int delta = offset + (i*rowsize + 3*j);
            fseek(file, delta, SEEK_SET);
            uint8_t b = fgetc(file);
            uint8_t g = fgetc(file);
            uint8_t r = fgetc(file);
            pixel px = {.R=r, .G=g, .B=b};
            img->bitmap[height-1-i][j] = px;
            //printf("%d,%d: %d\n", i, j, img->bitmap[height-1-i][j].R);
        }
    }

    fclose(file);

    return img;
}

void fprintint(FILE *file, int N, int size)
{
    for (int i=0; i<size; i++)
    {
        fputc((N>>(8*i))&255, file);
    }
}

void write_bmp(image *img, char *path)
{
    int height = img->height;
    int width = img->width;

    FILE* file = fopen(path, "wb");

    /* print header */
    fprintf(file, "BM");
    fseek(file, 0x0E, SEEK_SET);
    fprintint(file, 40, 4); // header size
    fprintint(file, width, 4); // width
    fprintint(file, height, 4); // height
    fprintint(file, 1, 2); // number of color planes
    fprintint(file, 24, 2); // bits per pixel
    fprintint(file, 0, 4); // compression
    fprintint(file, width*height, 2); // bits per pixel
    fprintint(file, 100, 4); // résolution horizontale
    fprintint(file, 100, 4); // résolution verticale
    fprintint(file, 0, 4); // nb de couleurs (0 => 2**n)
    fprintint(file, 0, 4); // nb de couleurs importantes (ignoré)

    int offset = ftell(file);
    fseek(file, 0x02, SEEK_SET);
    fprintint(file, offset + width*height, 4);
    fseek(file, 0x0A, SEEK_SET);
    fprintint(file, offset, 4);

    fseek(file, offset, SEEK_SET);

    for (int i=0; i<height; i++)
    {
        for (int j=0; j<width; j++)
        {
            pixel px = img->bitmap[height-1-i][j];
            //printf("(%d, %d, %d)\n", px.B, px.G, px.R);
            fputc(px.B, file);
            fputc(px.G, file);
            fputc(px.R, file);
        }
    }

    fclose(file);
}


double distance(pixel px1, dpixel px2)
{
    double Dr = px2.R - (double)px1.R; 
    double Dg = px2.G - (double)px1.G; 
    double Db = px2.B - (double)px1.B;
    return sqrt(Dr*Dr + Dg*Dg + Db*Db);
}


image* k_means_reduce(image* input, int k)
{
    int width = input->width;
    int height = input->height;

    dpixel* centroids = malloc(sizeof(dpixel) * k);
    long* counts = calloc(k, sizeof(long));

    int** classes = malloc(sizeof(int*) * height);
    for (int i=0; i<height; i++)
    {
        classes[i] = calloc(width, sizeof(int));
    }

    /* I. On choisit les centroides au hasard */
    bool** chosen = malloc(sizeof(bool*) * height);
    for (int i=0; i<height; i++)
    {
        chosen[i] = calloc(width, sizeof(bool));
    }
    for (int p=0; p<k; p++)
    {
        int i=0, j=0;
        do {
            i = rand() % height;
            j = rand() % width;
        } while (chosen[i][j] != false);

        chosen[i][j] = true;
        centroids[p] = dpixel_of_pixel(input->bitmap[i][j]);
    }
    for (int i=0; i<height; i++)
        free(chosen[i]);
    free(chosen);

    bool changed = true;
    int iteration=0;
    while (changed)
    {
        changed = false;
        
        printf("\rProcessing image, step %d... ", iteration);
        fflush(stdout);

        /*
        printf("[");
        for (int p=0; p<k; p++)
        {
            printf("(%f, %f, %f -> %ld); ", centroids[p].R, centroids[p].G, centroids[p].B, counts[p]);
        }
        printf("]\n");
        */

        for (int i=0; i<height; i++)
        {
            //printf("[");
            for (int j=0; j<width; j++)
            {
                // on calcule le centroide le plus proche
                
                int kmin=0;
                for (int p=1; p<k; p++)
                {
                    if (distance(input->bitmap[i][j], centroids[p]) < distance(input->bitmap[i][j], centroids[kmin]))
                    {
                        kmin = p;
                    }
                }

                if (classes[i][j] != kmin)
                {
                    classes[i][j] = kmin;
                    changed = true;
                }
                //printf("%d  ", classes[i][j]);
            }
            //printf("\n");
        }

        // maintenant il faut recalculer les centroides selon les barycentres
        for (int p=0; p<k; p++)
        {
            centroids[p].R = 0.;
            centroids[p].G = 0.;
            centroids[p].B = 0.;
            counts[p] = 0;
        }

        for (int i=0; i<height; i++)
        {
            for (int j=0; j<width; j++)
            {
                int c = classes[i][j];
                counts[c]++;
                centroids[c].R += (double)input->bitmap[i][j].R;
                centroids[c].G += (double)input->bitmap[i][j].G;
                centroids[c].B += (double)input->bitmap[i][j].B;
            }
        }

        for (int p=0; p<k; p++)
        {
            if (counts[p] != 0)
            {
                centroids[p].R /= (double)counts[p];
                centroids[p].G /= (double)counts[p];
                centroids[p].B /= (double)counts[p];
            }
        }
        iteration++;
    }
    printf("\n");

        
    image* img = create_image(height, width);
    for (int i=0; i<height; i++)
    {
        for (int j=0; j<width; j++) {
            img->bitmap[i][j] = pixel_of_dpixel(centroids[classes[i][j]]);
        }
    }
    
    free(counts);

    free(centroids);
    
    for (int i=0; i<height; i++)
    {
        free(classes[i]);
    }
    free(classes);
    
    return img;
}


int main(int argc, char **argv)
{
    printf("Simple image reducer (C) 2024 Valentin Moguérou. This is free software with ABSOLUTELY NO WARRANTY.\n");

    if (argc != 4)
    {
        fprintf(stderr, "Usage: %s <input_file> <output_file> <n_colors>\n", argv[0]);
        exit(1);
    }
    int k = strtol(argv[3], NULL, 10);

    srand(time(NULL));

    image *img = load_bmp(argv[1]);

    image *reduce = k_means_reduce(img, k);
    write_bmp(reduce, argv[2]);

    destroy_image(img);
    destroy_image(reduce);

    return 0;
}