import config, image, vector

Image *def convolution_brute(Image *image, V *kernel, int kw, kh):
    """
    Returns the result of convoluting image with kernel.
    The border is handled by treating all outside pixels as the nearest picture
    pixel.
    """
    int w = image->w
    int h = image->h
    V *data = image->data.v

    int ox = -kw / 2
    int oy = -kh / 2
    Image *result = image_new_undef_f(w, h)
    for int y = oy; y < h + oy; y++:
        for int x = ox; x < w + ox; x++:
            V sum = 0
            for int j = 0; j < kh; j++:
                int yj
                if y + j >= 0 and y + j < h: yj = y + j
                else: yj = y - oy
                for int i = 0; i < kw; i++:
                    int xi
                    if x + i >= 0 and x + i < w: xi = x + i
                    else: xi = x - ox
                    V v = data[w * yj + xi]
                    V k = kernel[j * kw + i]
                    sum += v * k
            result->data.v[(y - oy) * w + x - ox] = sum
    return result

Image *def gaussian_smooth(Image *image, double sigma, int size):
    """
    Returns a new image of size w times h with the result of a gaussian
    smoothing operation.
    TODO: Should do this in frequency domain, especially for big radius.
    """

    double sigma2 = sigma * sigma
    int kw = size
    int kh = size
    V *kernel = ALLOCATE(kw * kh * sizeof *kernel)
    V *pos = kernel
    int o = -size / 2
    for int j = o; j < kh + o; j++:
        for int i = o; i < kw + o; i++:
            # This is 2D (for 1D, sqrt the factor)
            V v = exp(-(i * i + j * j) / (2 * sigma2)) / (2 * pi * sigma2)
            *(pos++) = v

    Image *gaussian = convolution_brute(image, kernel, kw, kh)
    FREE(kernel)
    return gaussian

Image *def horizontal_sobel(Image *image):
    V kernel[] = {
        -1, 0, 1,
        -2, 0, 2,
        -1, 0, 1
    }
    Image *sobelh = convolution_brute(image, kernel, 3, 3)
    return sobelh

Image *def vertical_sobel(Image *image):
    V kernel[] = {
        1, 2, 1,
        0, 0, 0,
        -1, -2, -1
    }
    Image *sobelv = convolution_brute(image, kernel, 3, 3)
    return sobelv

def get_gradient(Image *image, Image **magnitude_out, Image **orientation_out,
    V hfactor, vfactor):
    """
    Use vertical and horizontal sobel operators to estimate the gradient, and
    return it as  magnitude an orientation maps.
    """
    int w = image->w
    int h = image->h

    Image *sobelh = horizontal_sobel(image)
    Image *sobelv = vertical_sobel(image)

    Image *magnitude = image_new_undef_f(w, h)
    Image *orientation = image_new_undef_f(w, h)
    for int j = 0; j < h; j++:
        for int i = 0; i < w; i++:
            V x = sobelh->data.v[i + w * j]
            V y = sobelv->data.v[i + w * j]
            V m = sqrt(x * x * hfactor + y * y * vfactor)
            # if m > 255: m = 255

            m /= 4
            
            magnitude->data.v[i + w * j] = m
            double a = atan2(y, x)

            orientation->data.v[i + w * j] = a
    *magnitude_out = magnitude
    *orientation_out = orientation

static inline def get_gradient_neighbors(V o, int i, j, *i_, *j_, *i__, *j__):
    if o == 0:
        *i_ = i + 1
        *i__ = i - 1
        *j_ = *j__ = j
    elif o == 1:
        *i_ = i + 1
        *i__ = i - 1
        *j_ = j - 1
        *j__ = j + 1
    elif o == 2:
        *i_ = *i__ = i
        *j_ = j - 1
        *j__ = j + 1
    elif o == 3:
        *i_ = i - 1
        *i__ = i + 1
        *j_ = j - 1
        *j__ = j + 1
    else:
        assert(0)

static inline def suppress_gradient_with_border_check(Image *gradient,
    *orientation, int i, j, Image *result):
    int w = gradient->w
    int h = gradient->h
    V g = gradient->data.v[i + w * j]
    V o = orientation->data.v[i + w * j]

    int i_, j_, i__, j__;
    get_gradient_neighbors(o, i, j, &i_, &j_, &i__, &j__)

    V g1 = g, g2 = g
    if i_ >= 0 and i_ < w and j_ >= 0 and j_ < h:
        g1 = gradient->data.v[i_ + w * j_]
    if i__ >= 0 and i__ < w and j__ >= 0 and j__ < h:
        g2 = gradient->data.v[i__ + w * j__]
    if g < g1 or g < g2: result->data.v[i + w * j] = 0
    else: result->data.v[i + w * j] = g

Image *def non_maximum_suppression(Image *gradient, Image *orientation):
    """
    For each pixel, if either the left or right neighbor is brighter, set it to
    0. That way, only thinned areas of maximum brightness are left. Because
    left/right neighbors are used, this is biased to create vertical lines.
    """
    int w = gradient->w
    int h = gradient->h
    Image *result = image_new_undef_f(w, h)

    # handle border pixels in a special way first
    for int j = 0; j < h; j += h - 1:
        for int i = 0; i < w; i++:
            suppress_gradient_with_border_check(gradient, orientation, i, j,
                result)
    for int i = 0; i < w; i += w - 1:
        for int j = 1; j < h - 1; j++:
            suppress_gradient_with_border_check(gradient, orientation, i, j,
                result)

    # now remaining pixels, without check
    for int j = 1; j < h - 1; j++:
        for int i = 1; i < w - 1; i++:
            V g = gradient->data.v[i + w * j]
            V o = orientation->data.v[i + w * j]

            int i_, j_, i__, j__;
            get_gradient_neighbors(o, i, j, &i_, &j_, &i__, &j__)
            
            V g1 = gradient->data.v[i_ + w * j_]
            V g2 = gradient->data.v[i__ + w * j__]
            if g < g1 or g < g2: result->data.v[i + w * j] = 0
            else: result->data.v[i + w * j] = g

    return result

Image *def edge_directions(Image *orientation):
    """
    Quantize the orientations into 4 bins:
    0 right
    1 right up
    2 up
    3 left up
    Those are all we need for pixel-edge-tracking 
    """
    int w = orientation->w
    int h = orientation->h
    Image *e = image_new_undef_f(w, h)
    for int j = 0; j < h; j++:
        for int i = 0; i < w; i++:
            V a = orientation->data.v[i + w * j]
            V c
            if a < 0: a += pi # not 2*pi, left/right both is horizontal
            if a > pi * 0.875: c = 0
            elif a > pi * 0.625: c = 3
            elif a > pi * 0.375: c = 2
            elif a > pi * 0.125: c = 1
            else: c = 0

            e->data.v[i + w * j] = c
    return e

Image *def canny(Image *raster, double gauss_sigma, int gauss_size,
    double hfactor, vfactor,
    char const *what):
    """
    Returns a new image of size w times h with the result of a canny edge
    detection. The implementation here uses the following algorithm:
    - Apply gaussian smoothing to the image.
    - Approximate gradient for each pixel (we use H and V Sobel here right now,
      should compare to e.g. differential gaussian instead)
    - Localize edges.
    - fix edges
    """

    Image *image = BtoV(raster)
    Image *smooth = gaussian_smooth(image, gauss_sigma, gauss_size)
    Image *magnitude, *orientation
    get_gradient(smooth, &magnitude, &orientation, hfactor, vfactor)

    Image *dirs = edge_directions(orientation)
    Image *thinned = non_maximum_suppression(magnitude, dirs)

    Image *binarized = thinned #binarize(thinned, w, h, 10)

    Image *r

    if what:
        if strstr(what, "cannysmooth"): r = VtoB(smooth)
        elif strstr(what, "cannymagnitude"): r = VtoB(magnitude)
        elif strstr(what, "cannyorientation"): r = VtoB(orientation)
        else: r = VtoB(binarized)
    else:
        #multiply_image(dirs, w, h, 255.0 / 3)
        r = VtoB(binarized)
    r->name = strdup(raster->name)
    return r

class Node:
    int x, y

int def follow_streak(Image *image, int x, y, *dirs, int *mark, int tag, kill,
    int *ax, *ay, *bx, *by):
    static Vector *stack
    if not stack: stack = vector_new(sizeof(Node))
    stack->count = 0
    int w = image->w
    int h = image->h
    int dx[] = {+1, +1,  0, -1, -1, -1,  0, +1, +1}
    int dy[] = { 0, -1, -1, -1,  0, +1, +1, +1,  0}
    #
    #  xxxxx
    #  x...x
    #  x.x.x
    #  x...x
    #  xxxxx
    #

    int first = 0
    Node start = {x, y}
    vector_add(stack, &start)
    if kill: image->data.b[x + w * y] = 0
    else: mark[x + w * y] = tag

    *ax = x
    *ay = y
    *bx = x
    *by = y

    while first < stack->count:
        Node *n = vector_get(stack, first++)
        for int i = 0; i < 8; i++:
            int nx = n->x + dx[i]
            int ny = n->y + dy[i]
            if nx >= 0 and ny >= 0 and nx < w and ny < h:
                int ok
                if kill: ok = (mark[nx + w * ny] == tag)
                else: ok = (mark[nx + w * ny] == 0)
                if ok and image->data.b[nx + w * ny]:
                    if kill: image->data.b[nx + w * ny] = 0
                    else: mark[nx + w * ny] = tag
                    dirs[i]++
                    Node next = {nx, ny}
                    vector_add(stack, &next)
                    if nx < *ax: *ax = nx
                    if nx > *bx: *bx = nx
                    if ny < *ay: *ay = ny
                    if ny > *by: *by = ny

    return stack->count

def remove_streaks(Image *image, char const *what,
    int minpixels,
    min_w_pixels, min_h_pixels,
    max_horizontal, max_vertical, max_slash, max_backslash,
    max_bound_w,
    max_bound_h,
    debug):
    """
    We trace some edges and remove them if they certainly aren't part of the
    wanted feature, but could confuse the hough transform later.
    """
    int w = image->w
    int h = image->h
    int mark[w * h]
    memset(mark, 0, w * h * sizeof(int))
    int tag = 1
    for int j = 0; j < h; j++:
        for int i = 0; i < w; i++:
            if not mark[i + j * w] and image->data.b[i + j * w]:
                int dirs[8] = {0, 0, 0, 0, 0, 0, 0, 0}
                int ax, ay, bx, by
                int size = follow_streak(image, i, j, dirs, mark, tag, 0,
                    &ax, &ay, &bx, &by)
                int kill = 0
                # TODO
                # Seems, using just the bounding box is enough (so we can get
                # rid of all the dirs counting). And using the bounding box,
                # instead of just the ratio, we could also already do the test for
                # min/max radius.
                if size >= minpixels and bx - ax >= min_w_pixels and\
                    by - ay >= min_h_pixels:
                    #int ho = dirs[0] + dirs[4]
                    #int ve = dirs[2] + dirs[6]
                    #int d1 = dirs[1] + dirs[5]
                    #int d2 = dirs[3] + dirs[7]

                    #printf("%d/%d: %d %d %d %d\n", i, j, ho, ve, d1, d2)
                    # We don't want to remove the pupil here under no
                    # circumstances, even if it is degenerated to just a small
                    # arc - so must be quite conservative.
                    # if ho > ve * max_horizontal or ve > ho * max_vertical:
                    #    d1 > d2 * max_slash or d2 > d1 * max_backslash or\
                    if (bx - ax) > max_bound_w * (by - ay) or\
                        (by - ay) > max_bound_h * (bx - ax):
                        kill = 1
                        if debug:
                            printf("Failed at %d/%d: %d %d %d %d not within bounds\n",
                                i, j, ax, ay, bx, by)
                else:
                    kill = 1
                    if debug:
                        printf("Failed at %d/%d: %d >= %d and %d - %d >= %d and %d - %d >= %d\n",
                            i, j, size, minpixels, bx, ax, min_w_pixels, by, ay, min_h_pixels)
                if kill:
                    follow_streak(image, i, j, dirs, mark, tag, 1,
                        &ax, &ay, &bx, &by)

                tag++

    if what and strstr(what, "allstreaks"):
        char name[MAX_PATH]
        get_image_name(image, "allstreaks", name)
        Image *out = image_new(None, w, h, IT_RGB)
        Color color[] = {
            {0, 0, 0},
            {255, 0, 0},
            {0, 128, 0},
            {0, 0, 255},
            {128, 128, 0},
            {0, 128, 128},
            {128, 0, 128},
            {255, 255, 255},
            {128, 128, 128},
            {255, 128, 120},
            {0, 255, 0},
            {128, 128, 255},
            {255, 255, 0},
            {0, 255, 255},
            {255, 0, 255},
            {192, 192, 192},
            }
        for int j = 0; j < h; j++:
            for int i = 0; i < w; i++:
                int c = mark[i + j * w]
                out->data.rgb[i + j * w] = color[c ? 1 + c % 15 : 0]
        save_tiff_rgb(name, out)
        image_del(out)
