# FIXME: sizeof(V) could be different when writing/reading a code file. The
# size in bits already is stored.. but it's currently ignored and the code
# assumes size on disk and at runtime is the same.

import config, getopt, tiff, canny, haar, daubechies
import detect_eye, pupil, hough, iris, iristexture, lids
import util
import algo1/process, algo2/process, algo3/process
import wavelets/cdf_9x7
import code
import tests/cdf97, tests/psize, tests/treeplot, tests/wp, tests/floatvector
import tests/noisewp, tests/randwp, tests/key, tests/zou, tests/fft
import tests/loggabor, tests/bitcom
import algo1/verify
import algo2/wed, algo2/training

static def usage(int argc, char **argv):
    printf("Usage: %s [options]\n", argv[0])
    printf("Options:\n")
    printf("-a ALGO Choose algorithm:\n")
    printf("   no\n")
    printf("   qsw bi-orthogonal quadratic spline wavelet by Mallat [default]\n")
    printf("   cdf97\n")
    printf("   d4\n")
    printf("   haar\n")
    printf("   pollen\n")
    printf("   d42 (pattern recognition)\n")
    printf("   lg log-gabor filter by Masek (implementing Daugman)\n")
    printf("-c CODE -c CODE1 {-c CODE*} compare iris codes/features.\n")
    printf("-f FILE uses the given feature sets for training and testing.\n")
    printf("-h Print this help.\n")
    printf("-i INPUT Specify input image.\n")
    printf("-k KEY Apply key for code generation.\n")
    printf("-l Use lid removal.\n")
    printf("-m INPUT Specify texture mask.\n")
    printf("-o OUTPUT Specify output image (or signal).\n")
    printf("-b BASE Specify the basename of output images. [default: derived from -i]\n")
    printf("-p PREFIX Specify a path to prefix to all output images.\n")
    printf("-r X Specify the number of rotational offsets to try.\n")
    printf("   This is given in bits to rotate left/right.\n")
    printf("-R Do not use bit rotation.\n")
    printf("-t INPUT Specify input texture.\n")
    printf("-T WHAT Run a test.\n")
    printf("-T psize Print pupil size (-i needs to come first).\n")
    printf("-T randwp Do random WP test (-i needs to come first).\n")
    printf("-w WHAT Specify what to do:\n")
    printf("   circles - draw found circles.\n")
    printf("   blackout - blackout the pupil.\n")
    printf("   lids, lidscanny, lidsmask - masked out lids.\n")
    printf("   mask - masked out parts.\n")
    printf("   textureback\n")
    printf("   iristexture\n")
    printf("   finaltexture\n")
    printf("   [pupil|iris]prep - prepare for pupil/iris detection.\n")
    printf("   [pupil|iris]canny - pupil/iris canny edges.\n")
    printf("   [pupil|iris]streaks - canny streaks postprocessing.\n")
    printf("   [pupil|iris]accums - write out hough accumulator visualizations.\n")
    printf("   algo1signals\n")
    printf("   algo1wavelet\n")
    printf("   algo1extrema\n")
    printf("   algocode\n")
    printf("   algo1decode\n")
    printf("   code - write binary iris code\n")
    printf("-x LxTxRxB ignore left/top/right/bottom pixels [default: 0x0x0x0]\n")

def display_code(Vector *code, char const *what, *algo, *key, Image *image):
    if hasword(what, "algocode"):
        char name[MAX_PATH]
        get_image_name_algo_key(image, "algocode", algo, key, name)
        Image *rgb
        if hasword(algo, "d42"):
            rgb = image_new_undef_rgb(13 * 10, 2 * 10)
            for int i = 0; i < 13; i++:
                for int j = 0; j < 2; j++:
                    V val = *(V *)vector_get(code, i * 2 + j)
                    val = fabs(val)
                    if j == 0:
                        val += 127.5
                    else:
                        val *= 4

                    if val < 0: val = 0
                    if val > 255: val = 255
                    Color color = {val, val, val}
                    draw_blank(rgb, i * 10, j * 10, 10, 10, color)
        else:
            int code_bytes = (code->count * code->element_size) / 2
            int w, h
            # FIXME
            if code_bytes * 8 == 512 * 20:
                w = 512
                h = 20
            elif code_bytes * 8 == 256 * 64:
                w = 256
                h = 64
            else:
                w = 8
                h = code_bytes
            rgb = image_new_undef_rgb(w, h)
            Color white = {255, 255, 255}
            Color black = {0, 0, 0}
            Color blue1 = {0, 0, 255}
            Color blue2 = {0, 255, 255}
            int codepos = 0
            int maskpos = code_bytes / 4
            int bitpos = 0
            uint32_t *cp = code->data

            for int i = 0; i < h; i++:
                for int j = 0; j < w; j++:
                    int x = j
                    # With wavelet algorithms, two rows are two scales.
                    # With "lg", two rows are real/imag part
                    int y = i / 2
                    if i & 1: y += h / 2
                    int masked = cp[maskpos] & (1 << bitpos)
                    if cp[codepos] & (1 << bitpos):
                        draw_vline(rgb, x, y, 1, masked ? blue1 : black)
                    else:
                        draw_vline(rgb, x, y, 1, masked ? blue2 : white)
                    bitpos++
                    if bitpos == 32:
                        bitpos = 0
                        codepos++
                        maskpos++

        save_tiff_rgb(name, rgb)

char *basename = None
char *prefix = None

def get_code_name_algo_key(
    char const *algo, char const *key, char *out):
    char const *prefix2 = "."
    if prefix:
        prefix2 = prefix
    else:
        if basename[0] == '/':
            prefix2 = ""

    if key and key[0]:
        sprintf(out, "%s/%s_%s:%s.code", prefix2, basename, algo, key)
    else:
        sprintf(out, "%s/%s_%s.code", prefix2, basename, algo)

static def apply_keys(Image *texture, *mask,
    char const *key, *what, *algo):

    char keydata[256]
    char *real_key

    char *s1, *s2
    # Are multiple keys given?
    s1 = (char *)key
    while 1:
        if s1:
            s2 = strchr(s1, ' ')
            if s2:
                strncpy(keydata, s1, s2 - s1)
                keydata[s2 - s1] = '\0'
                s1 = s2 + 1
            else:
                strcpy(keydata, s1)
                s1 = NULL
            real_key = keydata
        else:
            real_key = None

        printf("key: %s\n", real_key)

        BitCode *bitcode = ALLOCATE0(sizeof *bitcode)
        if hasword(algo, "d42"):
            Vector *code = algo2(texture, real_key, what, algo)
            display_code(code, what, algo, real_key, texture)

            bitcode->algorithm = strdup("d42")
            bitcode->bits = code->data
            bitcode->w = code->element_size * 8
            bitcode->h = code->count
            
        elif hasword(algo, "d42b"):
            assert(0) # FIXME

        elif hasword(algo, "lg"):
            Vector *code = algo3(texture, mask, what)
            display_code(code, what, algo, real_key, texture)

            bitcode->algorithm = strdup("lg")
            bitcode->bits = code->data
            bitcode->w = 512
            bitcode->h = 20
            bitcode->has_mask = True
        else:
            Vector *code = algo1(texture, 5, 10, real_key, what, algo)

            uint32_t *mc = generate_mask_code(mask, 5, 10)
            vector_cat(code, mc, 512 * 20 / 8)

            display_code(code, what, algo, real_key, texture)

            bitcode->algorithm = strdup("qsw")
            bitcode->bits = code->data
            bitcode->w = 512
            bitcode->h = 20
            bitcode->has_mask = True

        if hasword(what, "code"):
            char name[MAX_PATH]
            get_code_name_algo_key(algo, real_key, name)
            bitcode_to_file(bitcode, name)

        if not s1: break

int def main(int argc, char **argv):
    srand(time(NULL))
    char *infilename = None
    char *outfilename = None
    char const *what = ""
    char const *algo = ""
    char *codenames[1024]
    char *key = None
    char *maskname = None
    char *testsetname = None
    int codenames_count = 0
    int got_texture = False
    int norotate = False
    int rotation_offsets = 16
    int remove_lids = 0
    bool half_texture = False
    int cut_l = 0, cut_t = 0, cut_r = 0, cut_b = 0
    while 1:
        int c = getopt(argc, argv, "hi:lo:w:p:c:f:t:Rr:a:k:T:m:b:x:H")
        if c < 0: break
        switch c:
            case 'h':
                usage(argc, argv)
                exit(0)
                break
            case 'b':
                basename = optarg
                break
            case 'f':
                testsetname = optarg
                break
            case 'H':
                half_texture = true;
                break
            case 'i':
                infilename = optarg
                break
            case 'k':
                key = optarg
                break
            case 'l':
                remove_lids = 1
                break
            case 't':
                infilename = optarg
                got_texture = True
                break;
            case 'm':
                maskname = optarg
                break;
            case 'c':
                codenames[codenames_count++] = optarg
                break;
            case 'o':
                outfilename = optarg
                break
            case 'w':
                what = optarg
                break
            case 'a':
                algo = optarg
                break
            case 'p':
                prefix = optarg
                break
            case 'r':
                rotation_offsets = strtol(optarg, NULL, 0)
                break
            case 'R':
                norotate = True
                rotation_offsets = 0
                break
            case 'T':
                if hasword(optarg, "cdf97"):
                    exit(test_cdf97())
                if hasword(optarg, "psize"):
                    exit(test_psize(infilename, what))
                if hasword(optarg, "treeplot"):
                    exit(test_treeplot())
                if hasword(optarg, "wp"):
                    exit(test_wp(infilename, what))
                if hasword(optarg, "floatvector"):
                    exit(test_floatvector(infilename, what))
                if hasword(optarg, "noisewp"):
                    exit(test_noisewp())
                if hasword(optarg, "randwp"):
                    exit(test_randwp(infilename, what))
                if hasword(optarg, "key"):
                    exit(test_key(infilename, what))
                if hasword(optarg, "zou"):
                    exit(test_zou())
                if hasword(optarg, "fft"):
                    exit(test_fft())
                if hasword(optarg, "loggabor"):
                    exit(test_loggabor())
                if hasword(optarg, "bitcom"):
                    exit(test_bitcom())
                break;
            case 'x':
                char *p = optarg
                cut_l = strtol(p, &p, 10)
                cut_t = strtol(p + 1, &p, 10)
                cut_r = strtol(p + 1, &p, 10)
                cut_b = strtol(p + 1, &p, 10)
                break;
    if not infilename and not testsetname and not (codenames_count >= 2):
        usage(argc, argv)
        exit(-1)

    if prefix:
        set_output_prefix(prefix)

    if basename:
        set_output_basename(basename)
    elif infilename:
        char *autobase = strdup(infilename)
        char *bar = strchr(autobase, '_')
        if bar:
            *bar = '\0'
        basename = autobase

    if testsetname:
        int use_algo2 = 0
        Algo2Set *set = None
        set = data_load_sets(testsetname)
        Algo2Sample *fs = data_get_first_sample(set)
        if !strcmp(fs->bitcode->algorithm, "d42"):
            algo2_train(set)
            use_algo2 = 1
        else:
            set->rotation_offsets = rotation_offsets
            
        bool ok = False
        if hasword(what, "identify"):
            if use_algo2:
                algo2_test(set)
                ok = True
        if hasword(what, "verify"):
            if use_algo2:
                algo2_verify(set)
            else:
                algo1_verify(set)
            ok = True
            printf("intra:")
            for int i = 0; i < set->bins_count; i++:
                printf(" %d", set->intra_bins[i])
            printf("\n")
            printf("inter:")
            for int i = 0; i < set->bins_count; i++:
                printf(" %d", set->inter_bins[i])
            printf("\n")

        if not ok:
            fprintf(stderr, "You need to either specify 'identify' or 'verify'\n")
    elif got_texture:
        Image *texture = image_load_b(infilename)

        Image *mask
        if maskname and remove_lids: mask = image_load_b(maskname)
        else:
            mask = generate_empty_mask(texture->w, texture->h)
            if remove_lids:
                printf("No mask specified - using empty mask!\n")

        if half_texture:
            int p = texture->w / 8
            Image *halfed = image_copy_crop(texture, 0, 0,
                texture->w / 2, texture->h)
            image_transfer(texture, halfed, p * 3, 0, p, 0, p * 2, texture->h)
            image_transfer(texture, halfed, p * 7, 0, p * 3, 0, p, texture->h)

            printf(" using only half image.\n");
            if hasword(what, "half"):
                char name[MAX_PATH]
                get_image_name(texture, "half", name)
                save_tiff_luminance(name, halfed)

            image_del(texture)
            texture = halfed

        apply_keys(texture, mask, key, what, algo)

    elif infilename:
        Image *image = image_load_b(infilename)
        if not image:
            printf("Cannot load image %s. Aborting.\n", infilename)
            abort()

        printf("Image %s size is %d x %d.\n", image->name, image->w, image->h)
        if cut_l or cut_t or cut_r or cut_b:
            Image *cropped = image_copy_crop(image, cut_l, cut_t,
                image->w - cut_l - cut_r, image->h - cut_t - cut_b)
            image_del(image)
            image = cropped
            printf(" cropped to %d x %d.\n", image->w, image->h)
            

        int px, py, pr
        Image *result = find_pupil(image, what, &px, &py, &pr)
        Image *result2 = result
        int ix = px, iy = py, ir = pr
        if pr:
            result2 = find_iris(image, what, px, py, pr, &ix, &iy, &ir)

        Image *lidmask = None
        if remove_lids:
            lidmask = mask_lids(image, what, px, py, pr, ix, iy, ir)

        if hasword(what, "circles"):
            char name[MAX_PATH]
            get_image_name(image, "circles", name)
            save_tiff_rgb(name, result2)

        Image *mask
        if lidmask:
            int ox = ix - ir
            int oy = iy - ir
            mask = generate_mask(lidmask,
                px - ox, py - oy, pr, ix - ox, iy - oy, ir, 512, 64)
        else:
            mask = generate_empty_mask(512, 64)

        if hasword(what, "mask"):
            char name[MAX_PATH]
            get_image_name(image, "mask", name)
            save_tiff_luminance(name, mask)

        Image *texture = extract_iris_texture(image, px, py, pr, ix, iy, ir,
            512, 64, what)

        if hasword(what, "finaltexture"):
            char name[MAX_PATH]
            get_image_name(image, "finaltexture", name)
            save_tiff_luminance(name, texture)

        if key:
            apply_keys(texture, mask, key, what, algo)

    else:
        # if not strcmp(algo, "d42") or not strcmp(algo, "d42b"):
        #    Vector *c1 = readfile_datasize(codenames[0], sizeof(V))
        #    for int i = 1; i < codenames_count; i++:
        #        Vector *c2 = readfile_datasize(codenames[i], sizeof(V))
        #        if c1 and c2:
        #            V distance = WED(c1, c2, None)
        #            printf("%d: %f\n", i, distance)
        # else:
            #
            #int code_w = 512
            #int code_h = 20
            # if not strcmp(algo, "lg"):
            #    code_w = 256
            #    code_h = 64
            #int code_bytes = code_w * code_h / 8

            #Vector *c1 = readfile(codenames[0])
            #for int i = 1; i < codenames_count; i++:
            #    Vector *c2 = readfile(codenames[i])
            #    if c1 and c2:
            #        int d, m = 0
            #        if norotate:
            #            d = code_compare_bitwise_masked(
            #                c1->data, (char *)c1->data + code_bytes,
            #                c2->data, (char *)c2->data + code_bytes,
            #                code_bytes, &m)
            #        else:
            #            d = code_compare_masked(code_w, code_h,
            #                c1->data, (char *)c1->data + code_bytes,
            #                c2->data, (char *)c2->data + code_bytes,
            #                rotation_offsets, &m)
            #        printf("%d: %d/%d\n", i, d, m)
            #    else:
            #        printf("%d: failed\n", i)

        BitCode *c1 = bitcode_from_file(codenames[0])
        for int i = 1; i < codenames_count; i++:
            BitCode *c2 = bitcode_from_file(codenames[i])
            if c1 and c2:
                int d, m = 0
                if not strcmp(c1->algorithm, "d42"):
                    Vector *v1 = vector_new(c1->w / 8)
                    v1->count = c1->h;
                    v1->data = c1->bits;
                    Vector *v2 = vector_new(c2->w / 8)
                    v2->count = c2->h;
                    v2->data = c2->bits;
                    V distance = WED(v1, v2, None)
                    printf("%d: %f\n", i, distance)
                else:
                    d = bitcode_compare(c1, c2, rotation_offsets, &m)
                    printf("%d: %d/%d\n", i, d, m)
            else:
                printf("%d: failed\n", i)

    return 0
