#include "graph.h"
#include "naive_ppm.h"

#include <assert.h>
#include <math.h>
#include <stdio.h>
#include <stdlib.h>

#define COLOR_BITS 3
#define NB_COLORS (1 << (3 * COLOR_BITS))

size_t color_code(Pixel p)
{
    size_t r = p.R >> (8 - COLOR_BITS);
    size_t g = p.G >> (8 - COLOR_BITS);
    size_t b = p.B >> (8 - COLOR_BITS);
    return r << (2 * COLOR_BITS) | (g << COLOR_BITS) | b;
}

void free_graph(graph* G)
{
    free(G->degree);
    free(G->Adj);
    free(G->_label);
    free(G->_arcs);
    free(G->_source_arcs);
    free(G->_sink_arcs);
    free(G);
}

int capacity_neighboors(Pixel p1, Pixel p2, double sigma, int coeff, double dist)
{
    double dR = (p1.R - p2.R) / 255.0;
    double dG = (p1.G - p2.G) / 255.0;
    double dB = (p1.B - p2.B) / 255.0;
    double r = exp(-(dR * dR + dG * dG + dB * dB) / (2 * sigma * sigma));
    return (int)(0.5 + coeff * r / dist);
}

int capacity_external(Pixel p1, double* P, int max_capacity, double lambda, int coeff)
{
    size_t color = color_code(p1);
    int weight = P[color] == 0.0 ? max_capacity : (int)(0.5 - coeff * log(P[color]));
    return (int)(0.5 + lambda * weight);
}

void* realloc_or_exit(void* ptr, size_t size, char* msg)
{
    void* q = realloc(ptr, size);
    if (q == NULL && size != 0) {
        fprintf(stderr, "error during realloc (%s)\n", msg);
        exit(2);
    }
    return q;
}

void add_inner_arc(graph* G, size_t from, size_t to, int cap, int nb_neighboors)
{
    if (cap == 0) {
        return;
    }
    size_t d = G->degree[from];
    G->_arcs[(nb_neighboors + 3) * from + d].target = to;
    G->_arcs[(nb_neighboors + 3) * from + d].capacity = cap;
    G->_arcs[(nb_neighboors + 3) * from + d].flow = 0;
    G->_arcs[(nb_neighboors + 3) * from + d].opposite = NULL;
    G->degree[from] += 1;
}

// check to see if node k is surrounded by nodes with the same label
int surrounded(graph* G, char* label, size_t k, int nb_neighboors)
{
    size_t x = k % G->_width;
    size_t y = k / G->_width;
    char l = label[k];
    assert(l == 'R' || l == 'G');
    if (x > 0 && label[x - 1 + y * G->_width] != l) {
        return 0;
    }
    if (x < G->_width - 1 && label[x + 1 + y * G->_width] != l) {
        return 0;
    }
    if (y > 0 && label[x + (y - 1) * G->_width] != l) {
        return 0;
    }
    if (y < G->_height - 1 && label[x + (y + 1) * G->_width] != l) {
        return 0;
    }
    if (nb_neighboors == 4) {
        return 1;
    }

    if (x > 0 && y > 0 && label[x - 1 + (y - 1) * G->_width] != l) {
        return 0;
    }
    if (x > 0 && y < G->_height - 1 && label[x - 1 + (y + 1) * G->_width] != l) {
        return 0;
    }
    if (x < G->_width - 1 && y > 0 && label[x + 1 + (y - 1) * G->_width] != l) {
        return 0;
    }
    if (x < G->_width - 1 && y < G->_height - 1 && label[x + 1 + (y + 1) * G->_width] != l) {
        return 0;
    }

    return 1;
}

typedef struct {
    arc* data;
    size_t len;
    size_t alloc;
} Vec;

void append_arc(Vec* V, size_t node, int cap)
{
    if (cap == 0) {
        return;
    }
    if (!(V->len < V->alloc)) {
        V->alloc *= 2;
        V->data = realloc_or_exit(V->data, V->alloc * sizeof(arc), "append_arc");
    }
    assert(V->len < V->alloc);
    V->data[V->len].target = node;
    V->data[V->len].capacity = cap;
    V->data[V->len].flow = 0;
    V->data[V->len].opposite = NULL;
    V->len += 1;
}

graph* image_to_graph(PPM* img, PPM* hard_seeds, PPM* soft_seeds, double sigma, double lambda,
    int coeff, int nb_neighboors, int opt)
{
    assert(nb_neighboors == 4 || nb_neighboors == 8);

    // image and seeds must have the same size!
    assert(img->width == hard_seeds->width && img->height == hard_seeds->height);
    assert(soft_seeds == NULL
        || (img->width == soft_seeds->width && img->height == soft_seeds->height));

    char* label = malloc(img->width * img->height * sizeof(char));
    int nb_hard_sources = 0;
    int nb_hard_sinks = 0;
    int nb_hard_errors = 0;
    for (size_t k = 0; k < img->width * img->height; k++) {
        Pixel p = hard_seeds->P[k];
        if (p.R >> 4 == 0xf && p.G >> 4 == 0 && p.B >> 4 == 0) {
            // red seed: sink node (background)
            label[k] = 'R';
            nb_hard_sinks += 1;
        } else if (p.R >> 4 == 0 && p.G >> 4 == 0xf && p.B >> 4 == 0) {
            // green seed: source node (object)
            label[k] = 'G';
            nb_hard_sources += 1;
        } else {
            // neutral
            if (p.R >> 4 != p.G >> 4 || p.G >> 4 != p.B >> 4) {
                nb_hard_errors += 1;
            }
            label[k] = 'N';
        }
    }

    graph* G = malloc(sizeof(graph));
    G->_label = label;
    G->_width = img->width;
    G->_height = img->height;
    G->nb_nodes = img->width * img->height + 2; // "+2" for the SOURCE and
                                                // SINK nodes
    G->source = G->nb_nodes - 2;
    G->sink = G->nb_nodes - 1;

    G->degree = malloc(G->nb_nodes * sizeof(size_t)); // initialized later
    for (size_t k = 0; k < G->nb_nodes; k++) {
        G->degree[k] = 0;
    }

    // most of the inner nodes have 4, 5 or 6 neighboors (or 9, 10, 11 when nb_neighboors==8); we
    // use a linear array with unused cells unused cell have a target of (size_t)-1 and a weight of
    // 0
    size_t nb_inner_arcs = (nb_neighboors + 3) * G->_width * G->_height;
    // FIXME: I don't need to allocate (nb_neighboors + 3) * G->nb_nodes arcs for the inner nodes,
    // as the sink / source arcs are allocated elsewhere.
    // But without the next check, gcc raises the warning alloc-size-larger-than!?!?!
    // What is strange is that defining nb_inner_arcs = (nb_neighboors+3) * G->nb_nodes doesn't
    // require the check!
    /* if (nb_inner_arcs > PTRDIFF_MAX) { */
    /*     exit(17); */
    /* } */
    // for some unknown reason, it now works!
    G->_arcs = malloc(nb_inner_arcs * sizeof(arc));

    G->_nb_source_arcs = 0;
    G->_nb_sink_arcs = 0;
    int max_capacity = 0; // max capacity of inner nodes

    double proba_S[NB_COLORS];
    double proba_T[NB_COLORS];
    for (size_t i = 0; i < NB_COLORS; i++) {
        proba_S[i] = proba_T[i] = 0.0;
    }

    size_t nb_soft_sources = 0;
    size_t nb_soft_sinks = 0;
    int nb_soft_errors = 0;
    for (size_t i = 0; i < G->_width; i++) {
        for (size_t j = 0; j < G->_height; j++) {
            size_t k = i + j * G->_width;

            if (soft_seeds != NULL) {
                Pixel p = soft_seeds->P[k];
                if (p.R >> 4 == 0xf && p.G >> 4 == 0 && p.B >> 4 == 0) {
                    // red seed: sink node (background)
                    nb_soft_sinks += 1;
                    p = img->P[k];
                    uint32_t color = color_code(p);
                    proba_T[color] += 1.0;
                } else if (p.R >> 4 == 0 && p.G >> 4 == 0xf && p.B >> 4 == 0) {
                    // green seed: source node (object)
                    nb_soft_sources += 1;
                    p = img->P[k];
                    uint32_t color = color_code(p);
                    proba_S[color] += 1.0;
                } else if (p.R >> 4 != p.G >> 4 || p.G >> 4 != p.B >> 4) {
                    nb_soft_errors += 1;
                }
            }

            int node_capacity = 0; // capacity of current node

            if (i > 0) {
                // there is a node on the left
                size_t other = k - 1;
                if (!(opt & 1) || label[k] == 'N' || label[other] == 'N'
                    || label[k] != label[other]) {
                    int cap = capacity_neighboors(img->P[k], img->P[other], sigma, coeff, 1.0);
                    add_inner_arc(G, k, other, cap, nb_neighboors);
                    node_capacity += cap;
                }
            }
            if (i < G->_width - 1) {
                // there is a node on the right
                size_t other = k + 1;
                if (!(opt & 1) || label[k] == 'N' || label[other] == 'N'
                    || label[k] != label[other]) {
                    int cap = capacity_neighboors(img->P[k], img->P[other], sigma, coeff, 1.0);
                    add_inner_arc(G, k, other, cap, nb_neighboors);
                    node_capacity += cap;
                }
            }
            if (j > 0) {
                // there is a node above
                size_t other = k - G->_width;
                if (!(opt & 1) || label[k] == 'N' || label[other] == 'N'
                    || label[k] != label[other]) {
                    int cap = capacity_neighboors(img->P[k], img->P[other], sigma, coeff, 1.0);
                    add_inner_arc(G, k, other, cap, nb_neighboors);
                    node_capacity += cap;
                }
            }
            if (j < G->_height - 1) {
                // there is a node below
                size_t other = k + G->_width;
                if (!(opt & 1) || label[k] == 'N' || label[other] == 'N'
                    || label[k] != label[other]) {
                    int cap = capacity_neighboors(img->P[k], img->P[other], sigma, coeff, 1.0);
                    add_inner_arc(G, k, other, cap, nb_neighboors);
                    node_capacity += cap;
                }
            }
            if (nb_neighboors == 8) {
                if (i > 0 && j > 0) {
                    // there is a node up-left
                    size_t other = k - 1 - G->_width;
                    if (!(opt & 1) || label[k] == 'N' || label[other] == 'N'
                        || label[k] != label[other]) {
                        int cap
                            = capacity_neighboors(img->P[k], img->P[other], sigma, coeff, sqrt(2));
                        add_inner_arc(G, k, other, cap, nb_neighboors);
                        node_capacity += cap;
                    }
                }
                if (i < G->_width - 1 && j > 0) {
                    // there is a node on up right
                    size_t other = k + 1 - G->_width;
                    if (!(opt & 1) || label[k] == 'N' || label[other] == 'N'
                        || label[k] != label[other]) {
                        int cap
                            = capacity_neighboors(img->P[k], img->P[other], sigma, coeff, sqrt(2));
                        add_inner_arc(G, k, other, cap, nb_neighboors);
                        node_capacity += cap;
                    }
                }
                if (i > 0 && j < G->_height - 1) {
                    // there is a node bottom-left
                    size_t other = k - 1 + G->_width;
                    if (!(opt & 1) || label[k] == 'N' || label[other] == 'N'
                        || label[k] != label[other]) {
                        int cap
                            = capacity_neighboors(img->P[k], img->P[other], sigma, coeff, sqrt(2));
                        add_inner_arc(G, k, other, cap, nb_neighboors);
                        node_capacity += cap;
                    }
                }
                if (i < G->_width - 1 && j < G->_height - 1) {
                    // there is a node bottom-right
                    size_t other = k + 1 + G->_width;
                    if (!(opt & 1) || label[k] == 'N' || label[other] == 'N'
                        || label[k] != label[other]) {
                        int cap
                            = capacity_neighboors(img->P[k], img->P[other], sigma, coeff, sqrt(2));
                        add_inner_arc(G, k, other, cap, nb_neighboors);
                        node_capacity += cap;
                    }
                }
            }
            max_capacity = node_capacity > max_capacity ? node_capacity : max_capacity;
        }
    }

    if (nb_hard_errors > 0) {
        printf("\n");
        printf("*** WARNING %d invalid pixels in seed file\n", nb_hard_errors);
        printf("*** WARNING %d object pixels and %d background pixels were found\n",
            nb_hard_sources, nb_hard_sinks);
        printf("*** WARNING make sure pixels in seed file are either RED, GREEN, or some level of "
               "GRAY\n");
        printf("\n");
    }
    if (hard_seeds != soft_seeds && nb_soft_errors > 0) {
        printf("\n");
        printf("*** WARNING %d invalid pixels in soft seed file\n", nb_hard_errors);
        printf("*** WARNING %zu soft object pixels and %zu soft background pixels were found\n",
            nb_soft_sources, nb_soft_sinks);
        printf("\n");
    }

    max_capacity += 1;

    if (nb_soft_sources != 0) {
        for (size_t i = 0; i < NB_COLORS; i++) {
            proba_S[i] /= nb_soft_sources;
        }
    }
    if (nb_soft_sinks != 0) {
        for (size_t i = 0; i < NB_COLORS; i++) {
            proba_T[i] /= nb_soft_sinks;
        }
    }

    G->_source_arcs = malloc(sizeof(arc));
    Vec S;
    S.data = G->_source_arcs;
    S.len = 0;
    S.alloc = 1;

    G->_sink_arcs = malloc(sizeof(arc));
    Vec T;
    T.data = G->_sink_arcs;
    T.len = 0;
    T.alloc = 1;
    for (size_t k = 0; k < G->_width * G->_height; k++) {
        if (label[k] == 'R') {
            // red seed: sink node (background)
            if (!((opt >> 1) & 1) || !surrounded(G, label, k, nb_neighboors)) {
                append_arc(&T, k, max_capacity);
            }
        } else if (label[k] == 'G') {
            // green seed: source node (object)
            if (!((opt >> 1) & 1) || !surrounded(G, label, k, nb_neighboors)) {
                append_arc(&S, k, max_capacity);
            }
        } else if (label[k] == 'N') {
            if (lambda > 0.0) {
                double cap = capacity_external(img->P[k], proba_S, max_capacity, lambda, coeff);
                append_arc(&T, k, cap);

                cap = capacity_external(img->P[k], proba_T, max_capacity, lambda, coeff);
                append_arc(&S, k, cap);
            }
        } else {
            fprintf(stderr, "That's impossible, you shouldn't see this!\n");
            assert(0);
        }
    }
    append_arc(&S, NULL_NODE, -1);
    append_arc(&T, NULL_NODE, -1);

    G->_source_arcs = realloc_or_exit(S.data, S.len * sizeof(arc), "image_to_graph");
    G->_nb_source_arcs = S.len - 1; // -1 because of the extra dummy node
    G->_sink_arcs = realloc_or_exit(T.data, T.len * sizeof(arc), "image_to_graph");
    G->_nb_sink_arcs = T.len - 1; // -1 because of the extra dummy node

    G->degree[G->source] = G->_nb_source_arcs;
    G->degree[G->sink] = G->_nb_sink_arcs;

    // add sink as neighboors to each background node
    for (size_t i = 0; i < G->_nb_sink_arcs; i++) {
        size_t k = G->_sink_arcs[i].target;
        int cap = G->_sink_arcs[i].capacity;
        add_inner_arc(G, k, G->sink, cap, nb_neighboors);
    }

    // add source as neighboors to each foreground node
    for (size_t i = 0; i < G->_nb_source_arcs; i++) {
        size_t k = G->_source_arcs[i].target;
        int cap = G->_source_arcs[i].capacity;
        add_inner_arc(G, k, G->source, cap, nb_neighboors);
    }

    // create the neighboor array
    G->Adj = malloc(G->nb_nodes * sizeof(arc*));
    for (size_t k = 0; k < G->nb_nodes - 2; k++) {
        G->Adj[k] = G->_arcs + (nb_neighboors + 3) * k;
    }
    G->Adj[G->source] = G->_source_arcs;
    G->Adj[G->sink] = G->_sink_arcs;

    // initialize opposite arcs
    for (size_t k = 0; k < G->nb_nodes; k++) {
        for (size_t i = 0; i < G->degree[k]; i++) {
            arc* a = G->Adj[k] + i;
            if (a->opposite || a->target == G->source || a->target == G->sink) {
                continue;
            }
            for (arc* aa = G->Adj[a->target]; aa->target != NULL_NODE; aa++) {
                if (aa->target == k) {
                    a->opposite = aa;
                    aa->opposite = a;
                    break;
                }
            }
        }
    }

    // make sure all adjacency lists end with a dummy arc
    for (size_t k = 0; k < G->nb_nodes; k++) {
        G->Adj[k][G->degree[k]].target = NULL_NODE;
        G->Adj[k][G->degree[k]].capacity = -1;
        G->Adj[k][G->degree[k]].flow = -1;
        G->Adj[k][G->degree[k]].opposite = NULL;
    }

    return G;
}

int boundary_pixel(PPM* img, const char mask[], size_t x, size_t y, int w)
{
    int b_bg = 0;
    int b_obj = 0;
    for (int dx = -w; dx <= w; dx++) {
        for (int dy = -w; dy <= w; dy++) {
            // hack: if x+dx is negative, there is an overflow, and the test x+dx >= width becomes
            // true!
            if (x + dx >= img->width || y + dy >= img->height || dx * dx + dy * dy > w * w) {
                continue;
            }
            size_t idx = (x + dx) + (y + dy) * img->width;
            if (mask[idx]) {
                b_obj = 1;
            } else {
                b_bg = 1;
            }
            if (b_obj && b_bg) {
                return 1;
            }
        }
    }
    return 0;
}
// NOTE: modifies img!
int apply_mask_PPM(PPM* img, graph* G, const char mask[], const char* output_filename,
    color_mode obj_mode, color_mode bg_mode, color_mode boundary_mode, int boundary_width)
{
    for (size_t k = 0; k < img->width * img->height; k++) {
        if (G->_label[k] == 'N'
            && boundary_pixel(img, mask, k % img->width, k / img->width, boundary_width)) {
            img->P[k] = pixel_color(img->P[k], boundary_mode);
        } else if (mask[k]) {
            img->P[k] = pixel_color(img->P[k], obj_mode);
        } else {
            img->P[k] = pixel_color(img->P[k], bg_mode);
        }
    }
    return save_PPM(img, output_filename);
}

// vim: textwidth=100 foldmethod=indent
