import config, image, draw, vector, util
import wavelets/cdf_9x7, wavelets/daubechies_d4, wavelets/mallat_qsw
import wavelets/haar, wavelets/cdf_9x7, wavelets/pollen, wavelets/zou
import wavelets/schneid
"""

M = 3
N = 4
K * 78% = M * N = 12
L = 2

        I1[2]
S1[2]   I2[2]
        I3[2]

        I4[2]
S2[2]   I5[2]
        I6[2]

        I7[2]
S3[2]   I8[2]
        I9[2]

        I10[2]
S4[2]   I11[2]
        I12[2]

Si = (1/M) sum(j = 1..M) Ij+(i-1)*M

S1 = (1/M) sum(j = 1..3) Ij =
   = (1/M) I1 + I2 + I3

used: M = 5, N = 10, K * 78% = 50, K = 64, L = 512

"""

# used for drawing graphs
#macro DWTSCALE 32.0
macro DWTSCALE 64.0

# Select subbands
macro QSW_SCALE1 0x011
macro QSW_SCALE2 0x009

macro D4_SCALE1 3
macro D4_SCALE2 5

macro H_SCALE1 1
macro H_SCALE2 3

macro CDF97_SCALE1 3
macro CDF97_SCALE2 7

macro POLLEN_SCALE1 3
macro POLLEN_SCALE2 7

macro ZOU_SCALE1 3
macro ZOU_SCALE2 5

int def extrema_symmetric(V *s, int n, int *out, int debug):
    """
    List all local extremum points in the given signal of n samples.
    A threshold symmetric to the origin is used.
    """
    # TODO: Should we do some preprocessing like smoothing?

    int first = 1
    int first_mm = 0 # 0 = min, 1 = max
    int pos[n]
    int minmax = 0
    int found = 0
    V threshold1 = 0.05 # 0.01
    # TODO: consider the signal to be cyclic
    for int i = 0; i < n; i++:
        V d1 = s[i] - s[(i - 1) % n]
        V d2 = s[i] - s[(i + 1) % n]
        # a flat point
        if d1 == 0 and d2 == 0: continue
        int mm
        # a maximum
        if d1 >= 0 and d2 >= 0:
            if s[i] < threshold1: continue
            mm = 1
        # a minimum
        elif d1 <= 0 and d2 <= 0:
            if s[i] > -threshold1: continue
            mm = 0
        else:
            continue
        if not first and mm == minmax:
            # double extremum - we ignore it
            pass
        else:
            if first: first_mm = mm
            first = 0
            minmax = mm
            pos[found] = i
            found++

    int skip = 0
    if minmax == first_mm:
        first_mm = !first_mm
        skip += 1

    # TODO: Threshold.
    #V threshold = 0.05
    int j = 1
    for int i = skip; i < found; i++:
        out[j++] = pos[i]
        #int p1 = pos[i]
        #int p2 = pos[(i + 1) % found]
        #V d = fabs(s[p1] - s[p2])
        # if d > threshold and (fabs(s[p1]) > threshold or fabs(s[p2]) > threshold):
        #    out[j++] = pos[i]
        # else:
        #    i++

    # element 0 is 1 for maximum, 0 for minimum
    out[0] = first_mm
    if debug:
        printf("%d (%d) extrema found (%d).\n", j - 1, found, out[0])

    return j

int def extrema_relative(V *s, int n, int *out, int debug):
    """
    List all local extremum points in the given signal of n samples. A
    threshold using the relative distance is used.
    """
    # TODO: Should we do some preprocessing like smoothing?

    int first = 1
    int first_mm = 0
    int pos[n]
    int minmax = 0
    int found = 0

    # TODO: consider the signal to be cyclic
    for int i = 0; i < n; i++:
        V d1 = s[i] - s[(i - 1) % n]
        V d2 = s[i] - s[(i + 1) % n]
        # a flat point
        if d1 == 0 and d2 == 0: continue
        int mm
        # a maximum
        if d1 >= 0 and d2 >= 0:
            mm = 1
        # a minimum
        elif d1 <= 0 and d2 <= 0:
            mm = 0
        else:
            continue
        if not first and mm == minmax:
            # double extremum - we ignore it
            pass
        else:
            if first: first_mm = mm
            first = 0
            minmax = mm
            pos[found] = i
            found++

    # TODO: Threshold.
    int mm = !first_mm
    V threshold = 0.1
    int j = 1
    int prev = found - 1
    for int i = 0; i < found; i++:
        mm = !mm
        int next = (i + 1) % found
        int next2 = (next + 1) % found
        int p1 = pos[prev]
        int p2 = pos[i]
        int p3 = pos[next]
        int p4 = pos[next2]
        prev = i
        V dx = p3 - p2
        if dx < 0:
            dx += 512
        V dy = s[p2] - s[p3]
        if mm:
            if dy * dx < threshold:
               if s[p1] < s[p3] and s[p2] < s[p4]:
                i++
                mm = !mm
                continue
        else:
            if dy * dx > -threshold:
                if s[p1] > s[p3] and s[p2] > s[p4]:
                    i++
                    mm = !mm
                    continue

        out[j++] = pos[i]

    # element 0 is 1 for maximum, 0 for minimum
    out[0] = first_mm
    if debug:
        printf("%d (%d) extrema found (%d).\n", j, found, out[0])

    return j

def add_positions(Vector *result, int *pos, n):
    uint16_t c = n
    vector_add(result, &c)
    c = pos[0] # 0 = min, 1 = max
    vector_add(result, &c)
    int prev = 0
    for int j = 1; j < n; j++:
        int d = pos[j] - prev
        assert(d >= 0)
        c = d
        vector_add(result, &c)
        prev = pos[j]

static int def draw_extrema(Image *rgb, Vector *result, int i, vp, x, y, w,
    V ***dwts, int s,
    Color red):
    uint16_t c = *(uint16_t *)vector_get(result, vp++)
    int en = c
    c = *(uint16_t *)vector_get(result, vp++)
    int minmax = c
    int off = 0
    for int j = 1; j < en; j++:
        off += *(uint16_t *)vector_get(result, vp++)
        # int v = s[i * w + x] / 20
        int v = dwts[i][s][off] * DWTSCALE
        if minmax:
            draw_vline(rgb, x + off, y + 32 - v - 9, 9, red)
        else:
            draw_vline(rgb, x + off, y + 32 - v + 1, 9, red)
        minmax = !minmax
    return vp

static macro addbit {
    if (minmax) code[codepos] |= 1 << bitpos;
    bitpos++;
    if (bitpos == 32) {
        codepos++;
        bitpos = 0;
    }
}

uint32_t *def generate_code(Vector *result, int *w):
    """
    Bits are added to the code starting to the left and with the lowest bit in
    a byte.
    """
    int vp = 0
    unsigned int c;
    c = *(uint16_t *)vector_get(result, vp++)
    int n = c
    uint32_t *code = ALLOCATE0(2 * n * 512 / 8)
    int codepos = 0
    int bitpos = 0
    for int i = 0; i < n; i++:
        for int j = 0; j < 2; j++:
            c = *(uint16_t *)vector_get(result, vp++)
            int sl = c # length of minmax string
            c = *(uint16_t *)vector_get(result, vp++)
            int minmax = c # is the first extremum a minimum (0) or maximum (1)?
            # If first is a maximum, then we start inside a minimum. In case the
            # offset of the first extremum is 0, it will flip immediately in the
            # first loop iteration.
            minmax = !c
            int x = 0, ex = 0
            for int k = 1; k < sl; k++:
                int offset = 0
                c = *(uint16_t *)vector_get(result, vp++)
                offset = c
                ex += offset * 512 / w[j]
                while x < ex:
                    addbit
                    x++
                minmax = !minmax
            while x < 512:
                addbit
                x++
            assert(x == 512)
    assert(codepos == 2 * n * 512 / 32)
    return code

Vector *def algo1(Image *texture, int M, N, char const *key, *what, *algo):
    """
    We usually expect the texture to be sized 512x64, M = 5 and N = 10.
    This means we should end up with 10 1-D vector of 512 samples.
    """
    int w = texture->w
    # Downscale by M in vertical direction, so we get N 1-D signals s.
    V s[N * w]
    for int x = 0; x < w; x++:
        for int i = 0; i < N; i++:
            V sum = 0
            for int j = 0; j < M; j++:
                sum += texture->data.b[(j + i * M) * w + x]
            s[i * w + x] = sum / M

    # normalize the signals from 0..255 to -1..1 range, probably not necessary
    for int i = 0; i < N; i++:
        for int j = 0; j < w; j++:
            s[i * w + j] /= 255
            s[i * w + j] -= 0.5

    Color white = {255, 255, 255}
    Color white2 = {255, 255, 200}
    Color black = {0, 0, 0}
    Color red = {255, 100, 100}

    if hasword(what, "algo1signals"):
        char name[MAX_PATH]
        get_image_name(texture, "algo1signals", name)
        Image *rgb = image_new_undef_rgb(512, 64 * 10)
        for int i = 0; i < 10; i++:
            int x = 0, y = i * 64
            draw_blank(rgb, x, y, w, 64, white)
            draw_hline(rgb, x, y + 32, w, white2)
            for int j = 0; j < w; j++:
                int v = s[i * w + j] * 32
                draw_vline(rgb, j, y + 32 - v, 1, black)
            
        save_tiff_rgb(name, rgb)
        char *dot = strrchr(name, '.')
        strcpy(dot, ".txt")
        FILE *f = fopen(name, "wb")
        for int i = 0; i < w; i++:
            fprintf(f, "%d %f\n", i, s[i])
        fclose(f)

    # Apply wavelet transform to each 1-D signal
    V **dwts[N] # For each signal, two result vectors
    if hasword(algo, "no"):
        for int i = 0; i < N; i++:
            dwts[i] = ALLOCATE(2 * sizeof(V *))
            dwts[i][0] = ALLOCATE(w * sizeof(V))
            memcpy(dwts[i][0], s + i * w, w * sizeof(V))
            dwts[i][1] = ALLOCATE(w * sizeof(V))
            memcpy(dwts[i][1], s + i * w, w * sizeof(V))
    elif hasword(algo, "d4"):
        for int i = 0; i < N; i++:
            dwts[i] = d4(s + i * w, w)
            #FIXME
            dwts[i][0] = dwts[i][D4_SCALE1]
            dwts[i][1] = dwts[i][D4_SCALE2]
    elif hasword(algo, "haar"):
        for int i = 0; i < N; i++:
            dwts[i] = haar(s + i * w, w)
            #FIXME
            dwts[i][0] = dwts[i][H_SCALE1]
            dwts[i][1] = dwts[i][H_SCALE2]
    elif hasword(algo, "cdf97"):
        for int i = 0; i < N; i++:
            dwts[i] = cdf97(s + i * w, w, key)
            #FIXME
            dwts[i][0] = dwts[i][CDF97_SCALE1]
            dwts[i][1] = dwts[i][CDF97_SCALE2]
    elif hasword(algo, "pollen"):
        for int i = 0; i < N; i++:
            dwts[i] = pollen(s + i * w, w, key)
            #FIXME
            dwts[i][0] = dwts[i][POLLEN_SCALE1]
            dwts[i][1] = dwts[i][POLLEN_SCALE2]
    elif hasword(algo, "zou"):
        #for int i = 0; i < N; i++:
            #dwts[i] = zou(s + i * w, w, key)
            #FIXME
            #dwts[i][0] = dwts[i][ZOU_SCALE1]
            #dwts[i][1] = dwts[i][ZOU_SCALE2]
        #    zou_twice_from_key(s + i * w, w, dwts + i, key)
        zou_ten_twice(s, N, w, dwts, key)
    elif hasword(algo, "zoub"):
        zou_ten_twice(s, N, w, dwts, key)
    elif hasword(algo, "schneid"):
        for int i = 0; i < N; i++:
            dwts[i] = schneid(s + i * w, w, key)
            #FIXME
            dwts[i][0] = dwts[i][ZOU_SCALE1]
            dwts[i][1] = dwts[i][ZOU_SCALE2]
    else:
        for int i = 0; i < N; i++:
            dwts[i] = ALLOCATE(2 * sizeof(V *))
            int which
            int hex = 0
            if key: hex = strtol(key, NULL, 16)
            if key: which = hex & 4095
            else: which = QSW_SCALE1
            dwts[i][0] = qsw(s + i * w, w, which)
            if key: which = (hex >> 12) & 4095
            else: which = QSW_SCALE2
            dwts[i][1] = qsw(s + i * w, w, which)

    if hasword(what, "algo1wavelet"):
        char name[MAX_PATH]
        get_image_name_algo_key(texture, "algo1wavelet", algo, key, name)
        Image *rgb = image_new_undef_rgb(512 * 2, 64 * 10)
        for int i = 0; i < 10; i++:
            int x = 0, y = i * 64
            draw_blank(rgb, x, y, w, 64, white)
            draw_blank(rgb, x + 512, y, w, 64, white2)
            draw_hline(rgb, x, y + 32, w, white2)
            draw_hline(rgb, x + 512, y + 32, w, white)
            for int j = 0; j < w; j++:
                #int v = s[i * w + j] / 20
                int v = dwts[i][0][j] * DWTSCALE
                draw_vline(rgb, j, y + 32 - v, 1, black)
                int v2 = dwts[i][1][j] * DWTSCALE
                draw_vline(rgb, 512 + j, y + 32 - v2, 1, black)
            
        save_tiff_rgb(name, rgb)

        char *dot = strrchr(name, '.')
        strcpy(dot, ".txt")
        FILE *f = fopen(name, "wb")
        for int i = 0; i < w; i++:
            fprintf(f, "%d %f\n", i, dwts[0][0][i])
        fclose(f)

    # Format for result vector:
    # [n] number of signals
    # n times
    #     [n] number of extrema
    #     [s] sign of first extremum
    #     n times
    #         [o] offset to next extremum (first one offset to 0)
    Vector *result = vector_new(2)
    uint16_t c = N
    vector_add(result, &c)
    int pos[w]
    int first_count = 0
    for int i = 0; i < N; i++:
        #int n = extrema(s + i * w + w / 2, w / 2, pos)
        int debug = 0
        int n = extrema_relative(dwts[i][0], w, pos, debug)
        int d = result->count
        add_positions(result, pos, n)
        first_count += result->count - d

        #n = extrema(s + i * w + w / 4, w / 4, pos)
        n = extrema_relative(dwts[i][1], w, pos, debug)
        add_positions(result, pos, n)

    printf("Feature vector: %d (%d+%d)\n", result->count, first_count,
        result->count - first_count)

    if hasword(what, "algo1extrema"):
        char name[MAX_PATH]
        get_image_name_algo_key(texture, "algo1extrema", algo, key, name)

        Image *rgb = image_new_undef_rgb(512 * 2, 64 * 11)

        draw_blank(rgb, 0, 0, w, 64 * 11, white)
        draw_blank(rgb, 512, 0, w, 64 * 11, white2)

        for int i = 0; i < N; i++:
            draw_hline(rgb, 0, i * 64 + 32, w, white2)
            draw_hline(rgb, 0 + 512, i * 64 + 32, w, white)

        int vp = 0
        c = *(uint16_t *)vector_get(result, vp++)
        int sn = c
        for int i = 0; i < sn; i++:
            int y = i * 64
            #vp = draw_extrema(rgb, result, i, vp, w / 2, y, w, s, red)
            #vp = draw_extrema(rgb, result, i, vp, w / 4, y, w, s, yellow)
            vp = draw_extrema(rgb, result, i, vp, 0, y, w, dwts, 0, red)
            vp = draw_extrema(rgb, result, i, vp, 512, y, w, dwts, 1, red)

        for int i = 0; i < N; i++:
            int y = i * 64
            for int j = 0; j < w; j++:
                #int v = s[i * w + j] / 20
                int v = dwts[i][0][j] * DWTSCALE
                draw_vline(rgb, j, y + 32 - v, 1, black)
                int v2 = dwts[i][1][j] * DWTSCALE
                draw_vline(rgb, 512 + j, y + 32 - v2, 1, black)

        save_tiff_rgb(name, rgb)

        # output extrema positions to text file
        char *dot = strrchr(name, '.')
        strcpy(dot, ".txt")
        FILE *f = fopen(name, "wb")
        vp = 0
        vector_get(result, vp++) # number of signals
        uint16_t en = *(uint16_t *)vector_get(result, vp++)
        vector_get(result, vp++) # minmax
        int off = 0
        for int j = 1; j < en; j++:
            off += *(uint16_t *)vector_get(result, vp++)
            fprintf(f, "%d %f\n", off, dwts[0][0][off])
        fclose(f)

    int ws[] = {512, 512}
    uint32_t *code = generate_code(result, ws)

    if hasword(what, "algo1decode"):
        char name[MAX_PATH]
        get_image_name_algo_key(texture, "algo1decode", algo, key, name)
        Image *rgb = image_new_undef_rgb(512 * 2, 11 * 64)

        draw_blank(rgb, 0, 0, 512, 11 * 64, white)
        draw_blank(rgb, 512, 0, 512, 11 * 64, white2)
        int codepos = 0
        int bitpos = 0
        for int i = 0; i < 10; i++:
            for int k = 0; k < 2; k++:
                for int j = 0; j < 512; j++:
                    int y = i * 64 + 32

                    int v = s[i * w + j] * 16
                    draw_vline(rgb, 512 * k + j, y - v, 1, black)
                    
                    if code[codepos] & (1 << bitpos):
                        draw_vline(rgb, 512 * k + j, y - 16, 1, red)
                    else:
                        draw_vline(rgb, 512 * k + j, y + 16, 1, red)

                    bitpos++
                    if bitpos == 32:
                        bitpos = 0
                        codepos++

        save_tiff_rgb(name, rgb)

    Vector *cv = vector_new(1)
    vector_cat(cv, code, 512 * 20 / 8)
    return cv
