#include "hash_nb.h"
#include "segment.h"

int VERBOSITY;

void help(const char* exe)
{
    printf("usage: %s [OPTIONS] IMG SEEDS OUTPUT\n", exe);
    printf("    or %s [OPTIONS] IMG HARD_SEEDS SOFT_SEEDS OUTPUT\n", exe);
    printf("    or %s [OPTIONS] BASE\n", exe);
    printf("    or %s --test [OPTIONS] IMG SEEDS [ARGS]\n", exe);
    printf("In the second case, IMG, SEEDS and OUTPUT will defined as "
           "'BASE.ppm', 'BASE-seeds.ppm' and 'BASE-output.ppm'.\n");
    printf("\n");
    printf("Note: SOFT_SEEDS can be the empty string. In that case, no soft seeds are used.\n");
    printf("SOFT_SEEDS is only used to get the colour distribution, not for 'hard' edges.\n");
    printf("\n");
    printf("Options:\n");
    printf("  -i  /  --invert               invert source and sink\n");
    printf("  -4  /  -8                     4 or 8 neighbooring system\n");
    printf("  -S FLOAT  /  --sigma=FLOAT    sigma constant (default: 0.15)\n");
    printf("  -L FLOAT  /  --lambda=FLOAT   lambda constant (default: 0.0, or 0.1 if soft seed "
           "file is given)\n");
    printf("  --opt=N                       activate graph optimizations (default: 0, possible "
           "values: 0,1,2,3)\n");
    printf("  -C INT  /  --coeff=INT        coefficient (default: 10000)\n");
    printf("  -O MODE  /  --object=MODE     object (foreground) output mode (see below)\n");
    printf("  -B MODE  /  --background=MODE background output mode (see below)\n");
    printf("  --boundary=MODE               the boundary output mode (see below)\n");
    printf("  --boundary-width=NB           the boundary width (default: 0)\n");

    printf("  --assert-flow=NUMBER          check that the max flow is equal to this number\n");
    printf("  -T  /  --test                 call the main_test function instead of maxflow\n");
    printf("                                arguments are of the form IMG SEEDS [ARGS] and\n");
    printf("                                ARGS are passed to the main_test function\n");

    printf("  -h  /  --help                 display this message\n");
    printf("  -V  /  --version              show version information\n");
    printf("  -v  /  --verbosity            increase verbosity of log messages\n");
    printf("  -q  /  --quiet                set verbosity to 0\n");
    printf("\n");
    printf("MODE specifies by what color should object / background pixels should be replaced:\n");
    printf("  none              as in original image\n");
    printf("  solid:COLOR       solid color\n");
    printf("  light:COLOR       light version of this color based on original intensity\n");
    printf("  dark:COLOR        dark version of this color based on original intensity\n");
    printf("COLOR is alway a 6 characted hex code, as in 'ff0000' for red.\n");
    printf("\n");
    printf("Notes: optimization levels for the graph correspond to the following:\n");
    printf("    0: no optimization\n");
    printf("    1: do not add inner edges between similarly annotated pixels\n");
    printf("    2: do not add outer edges to/from annotated pixels surrounded by similarly "
           "annotated pixels\n");
    printf("    3: both optimizations\n");
}

Pixel hexcode2pixel(const char* s)
{
    Pixel p;
    uint32_t c = strtoul(s, NULL, 16);
    p.R = c >> 16 & 0xff;
    p.G = c >> 8 & 0xff;
    p.B = c & 0xff;
    return p;
}

int main(int argc, char** argv)
{
    char* exe = argv[0];
    if (argc == 1) {
        help(exe);
        return 1;
    }

    char short_options[] = "hVvS:L:C:48O:B:qiT";
    static struct option long_options[] = { { "help", no_argument, 0, 'h' },
        { "verbose", no_argument, 0, 'h' }, { "sigma", required_argument, 0, 'S' },
        { "lambda", required_argument, 0, 'L' }, { "coeff", required_argument, 0, 'C' },
        { "object", required_argument, 0, 'O' }, { "background", required_argument, 0, 'B' },
        { "quiet", required_argument, 0, 'q' }, { "invert", required_argument, 0, 'i' },
        { "test", no_argument, 0, 'T' }, { "boundary", required_argument, 0, 'b' },
        { "assert-flow", required_argument, 0, 'A' }, { "opt", required_argument, 0, 'G' },
        { "boundary-width", required_argument, 0, 'w' }, { 0, 0, 0, 0 } };

    VERBOSITY = 1;
    double sigma = 0.15;
    double lambda = -1.0; // set to default value after argument have been parsed
    int coeff = 10000;
    int nb_neighboors = 4;
    int boundary_width = 0;
    color_mode bg_mode, obj_mode, boundary_mode;
    obj_mode.mode = NONE;
    obj_mode.color.R = 0x00;
    obj_mode.color.G = 0x00;
    obj_mode.color.B = 0x00;
    bg_mode.mode = SOLID;
    bg_mode.color.R = 0xff;
    bg_mode.color.G = 0x00;
    bg_mode.color.B = 0x00;
    boundary_mode.mode = DARK;
    boundary_mode.color.R = 0xff;
    boundary_mode.color.G = 0xff;
    boundary_mode.color.B = 0xff;
    int assert_flow = -1;
    int graph_opt = 0;

    int invert = 0;
    int test = 0;

    int opt;
    int long_index;
    while ((opt = getopt_long(argc, argv, short_options, long_options, &long_index)) != -1) {
        switch (opt) {
        case 'S':
            // TODO check for errors
            sigma = strtod(optarg, NULL);
            break;
        case 'L':
            // TODO check for errors
            lambda = strtod(optarg, NULL);
            break;
        case 'C':
            // TODO check for errors
            coeff = strtol(optarg, NULL, 10);
            break;
        case 'O':
            if (0 == strncmp(optarg, "none", 4)) {
                obj_mode.mode = NONE;
            } else if (0 == strncmp(optarg, "solid:", 6)) {
                obj_mode.mode = SOLID;
                obj_mode.color = hexcode2pixel(optarg + 6);
            } else if (0 == strncmp(optarg, "light:", 6)) {
                obj_mode.mode = LIGHT;
                obj_mode.color = hexcode2pixel(optarg + 6);
            } else if (0 == strncmp(optarg, "dark:", 5)) {
                obj_mode.mode = DARK;
                obj_mode.color = hexcode2pixel(optarg + 5);
            } else {
                fprintf(stderr, "unknown mode: %s\n", optarg);
                exit(1);
            }
            break;
        case 'B':
            if (0 == strncmp(optarg, "none", 4)) {
                bg_mode.mode = NONE;
            } else if (0 == strncmp(optarg, "solid:", 6)) {
                bg_mode.mode = SOLID;
                bg_mode.color = hexcode2pixel(optarg + 6);
            } else if (0 == strncmp(optarg, "light:", 6)) {
                bg_mode.mode = LIGHT;
                bg_mode.color = hexcode2pixel(optarg + 6);
            } else if (0 == strncmp(optarg, "dark:", 5)) {
                bg_mode.mode = DARK;
                bg_mode.color = hexcode2pixel(optarg + 5);
            } else {
                fprintf(stderr, "unknown mode: %s\n", optarg);
                exit(1);
            }
            break;
        case 'b':
            if (0 == strncmp(optarg, "none", 4)) {
                boundary_mode.mode = NONE;
            } else if (0 == strncmp(optarg, "solid:", 6)) {
                boundary_mode.mode = SOLID;
                boundary_mode.color = hexcode2pixel(optarg + 6);
            } else if (0 == strncmp(optarg, "light:", 6)) {
                boundary_mode.mode = LIGHT;
                boundary_mode.color = hexcode2pixel(optarg + 6);
            } else if (0 == strncmp(optarg, "dark:", 5)) {
                boundary_mode.mode = DARK;
                boundary_mode.color = hexcode2pixel(optarg + 5);
            } else {
                fprintf(stderr, "unknown mode: %s\n", optarg);
                exit(1);
            }
            break;
        case 'w':
            // TODO check for errors
            boundary_width = strtod(optarg, NULL);
            break;
        case '4':
            nb_neighboors = 4;
            break;
        case '8':
            nb_neighboors = 8;
            break;
        case 'i':
            invert = 1;
            break;
        case 'A':
            assert_flow = atoi(optarg);
            break;
        case 'G':
            graph_opt = atoi(optarg);
            if (graph_opt < 0 || graph_opt > 3) {
                fprintf(stderr, "possible optimization levels for graphs are 0, 1, 2 or 3\n");
                exit(2);
            }
            break;
        case 'T':
            test = 1;
            break;
        case 'v':
            VERBOSITY += 1;
            break;
        case 'q':
            VERBOSITY = 0;
            break;
        case 'h':
            help(exe);
            return 0;
        case 'V':
            printf("maxflow segmenter: %s\n", HASH_NB);
            printf("%s\n", VERSION);
            return 0;
        default: /* '?' */
            // TODO
            return 1;
        }
    }

    argv += optind;
    argc -= optind;

    if ((test == 0 && argc != 1 && argc != 3 && argc != 4) || (test == 1 && argc < 2)) {
        help(exe);
        return 1;
    }

    char input_filename[256];
    char hard_seeds_filename[256];
    char soft_seeds_filename[256];
    char output_filename[256];
    if (test == 0) {
        if (argc == 3) {
            strncpy(input_filename, argv[0], sizeof(input_filename) - 1);
            strncpy(hard_seeds_filename, argv[1], sizeof(hard_seeds_filename) - 1);
            strncpy(soft_seeds_filename, argv[1], sizeof(hard_seeds_filename) - 1);
            lambda = lambda < 0 ? 0.0 : lambda;
            strncpy(output_filename, argv[2], sizeof(output_filename) - 1);
        } else if (argc == 4) {
            strncpy(input_filename, argv[0], sizeof(input_filename) - 1);
            strncpy(hard_seeds_filename, argv[1], sizeof(hard_seeds_filename) - 1);
            strncpy(soft_seeds_filename, argv[2], sizeof(hard_seeds_filename) - 1);
            lambda = lambda < 0 ? 0.1 : lambda;
            strncpy(output_filename, argv[3], sizeof(output_filename) - 1);
        } else if (argc == 1) {
            snprintf(input_filename, sizeof(input_filename) - 1, "%s.ppm", argv[0]);
            snprintf(hard_seeds_filename, sizeof(hard_seeds_filename) - 1, "%s-seeds.ppm", argv[0]);
            snprintf(soft_seeds_filename, sizeof(hard_seeds_filename) - 1, "%s-seeds.ppm", argv[0]);
            lambda = lambda < 0 ? 0.0 : lambda;
            snprintf(output_filename, sizeof(output_filename) - 1, "%s-output.ppm", argv[0]);
        }
    } else {
        strncpy(input_filename, argv[0], sizeof(input_filename) - 1);
        strncpy(hard_seeds_filename, argv[1], sizeof(hard_seeds_filename) - 1);
        strncpy(soft_seeds_filename, argv[1], sizeof(hard_seeds_filename) - 1);
        strncpy(output_filename, "none", sizeof(output_filename) - 1);
        argv += 2;
        argc -= 2;
    }

    if (test == 0) {
        LOG(2, ">> max-flow based image segmentation\n");
        LOG(2, ">> %s\n", VERSION);
    }

    LOG(2, ">> config:\n");
    LOG(2, ">>   original image: %s\n", input_filename);
    LOG(2, ">>   seeds:          %s\n", hard_seeds_filename);
    LOG(2, ">>   output image:   %s\n", output_filename);
    LOG(2, ">>   sigma  = %f\n", sigma);
    LOG(2, ">>   lambda = %f\n", lambda);
    LOG(2, ">>   coeff  = %i\n", coeff);

    LOG(2, ">> opening %s ... ", input_filename);
    PPM* img = open_PPM(input_filename);
    LOG(2, "done\n");

    LOG(2, ">> opening %s ... ", hard_seeds_filename);
    PPM* hard_seeds = open_PPM(hard_seeds_filename);
    LOG(2, "done\n");

    PPM* soft_seeds;
    if (strlen(soft_seeds_filename) == 0) {
        soft_seeds = NULL;
    } else if (0 == strcmp(soft_seeds_filename, hard_seeds_filename)) {
        soft_seeds = hard_seeds;
    } else {
        LOG(2, ">> opening %s ... ", soft_seeds_filename);
        soft_seeds = open_PPM(soft_seeds_filename);
        LOG(2, "done\n");
    }

    // image and seeds must have the same size!
    if (img->width != hard_seeds->width && img->height == hard_seeds->height) {
        fprintf(stderr, "image file and seed file have different sizes (%zux%zu vs %zux%zu)\n",
            img->width, img->height, hard_seeds->width, hard_seeds->height);
        exit(2);
    }

    if (soft_seeds != NULL
        && (hard_seeds->width != soft_seeds->width && hard_seeds->height == soft_seeds->height)) {
        fprintf(stderr,
            "hard seeds file and soft seed file have different sizes (%zux%zu vs %zux%zu)\n",
            hard_seeds->width, hard_seeds->height, hard_seeds->width, hard_seeds->height);
        exit(2);
    }

    LOG(2, ">> creating graph ... ");
    graph* G = image_to_graph(
        img, hard_seeds, soft_seeds, sigma, lambda, coeff, nb_neighboors, graph_opt);
    if (invert) {
        size_t tmp = G->source;
        G->source = G->sink;
        G->sink = tmp;
    }
    LOG(2, "done\n");
    if (VERBOSITY > 1) {
        int nb_edges = 0;
        for (size_t k = 0; k < G->nb_nodes; k++) {
            nb_edges += G->degree[k];
        }
        LOG(2, ">> The graph contains %zu nodes; its source is %zu and its sink is %zu\n",
            G->nb_nodes, G->source, G->sink);
        LOG(2,
            ">> It countains %i inner edges, "
            "%i edges from the source "
            "and %i edges to the sink\n",
            nb_edges / 2 - G->degree[G->source] - G->degree[G->sink], G->degree[G->source],
            G->degree[G->sink]);
    }

    if (test == 1) {
        return (main_test(G, argc, argv));
    } else {
        char* Seen = calloc(G->nb_nodes, sizeof(char));
        LOG(2, ">> calling the maxflow function\n");
        int flow = maxflow(G, Seen);
        LOG(1, "> final flow value is %d\n", flow);
        if (assert_flow > 0 && flow != assert_flow) {
            LOG(1, "*** This is not the expected value (%d)\n", assert_flow);
        }

        // make sure all annotated nodes are in the correct part
        for (size_t k = 0; k < G->nb_nodes; k++) {
            // background annotated pixels cannot be Seen
            // assert(G->_label[k] != 'R' || !Seen[k]);
            if (G->_label[k] == 'R' && Seen[k]) {
                fprintf(stderr, "!!! background node %zu is wrongly labeled 'object'\n", k);
            }

            // NOTE: with both graph optimizations are activated, surrounded source pixels become
            // inaccessible from the source and are thus labeled wrongly as "background". To deal
            // with that, annotated source pixels must be forced as object pixels after the maxflow
            // has finished. This is done by recording them (via the "label" array) into an
            // additional field of the graph structure.
            if (G->_label[k] == 'G') {
                Seen[k] = 1;
            }
        }

        LOG(2, ">> creating and saving output image ...");
        apply_mask_PPM(
            img, G, Seen, output_filename, obj_mode, bg_mode, boundary_mode, boundary_width);
        LOG(2, "done\n");

        LOG(1, "> resulting image saved to file '%s'\n", output_filename);

        // free memory
        free(Seen);
        free_graph(G);
        if (soft_seeds != NULL && soft_seeds != hard_seeds) {
            free_PPM(soft_seeds);
        }
        free_PPM(hard_seeds);
        free_PPM(img);

        if (assert_flow > 0 && flow != assert_flow) {
            return 1;
        } else {
            return 0;
        }
    }
}

// vim: textwidth=100 foldmethod=indent
