import config, image, filter

# daubechies d4 constants
double h0, h1, h2, h3
double g0, g1, g2, g3

static bool got_constants

Filter *def d4_get_filter():
    Filter *f = ALLOCATE0(sizeof *f)
    f->hn = 4
    f->gn = 4
    f->h = malloc(f->hn * sizeof(V))
    f->g = malloc(f->gn * sizeof(V))

    daubechies_d4_init()
    f->h[0] = h0
    f->h[1] = h1
    f->h[2] = h2
    f->h[3] = h3
    f->g[0] = g0
    f->g[1] = g1
    f->g[2] = g2
    f->g[3] = g3
    
    return f

def daubechies_d4_transform(V const *indata, int n, instep,
    V *outlow, int lowstep, V *outhigh, int highstep):
    """
    Given a signal of n samples, return low and highband of the wavelet
    transform, downsampling both to half length.
    """
    if not got_constants: daubechies_d4_init()
    V const *v = indata
    V *vl = outlow, *vh = outhigh
    int s1 = instep
    int s2 = instep * 2
    int s3 = instep * 3
    int ii = 0, li = 0, hi = 0
    for int i = 0; i < n - 3; i += 2:
        vl[li] = v[ii] * h0 + v[ii + s1] * h1 + v[ii + s2] * h2 + v[ii + s3] * h3
        vh[hi] = v[ii] * g0 + v[ii + s1] * g1 + v[ii + s2] * g2 + v[ii + s3] * g3
        ii += s2
        li += lowstep
        hi += highstep

    # What to do at the right edge where v[i + 2] and v[i + 3] do not exist?
    # Here we wrap around as the iris actually is circular.
    vl[li] = v[s1 * (n - 2)] * h0 + v[s1 * (n - 1)] * h1 + v[0] * h2 + v[s1] * h3
    vh[hi] = v[s1 * (n - 2)] * g0 + v[s1 * (n - 1)] * g1 + v[0] * g2 + v[s1] * g3

def daubechies_d4_transform_no_key(V const *indata, int n, instep,
    V *outlow, int lowstep, V *outhigh, int highstep, char const *key):
    daubechies_d4_transform(indata, n, instep, outlow, lowstep, outhigh,
        highstep)

def daubechies_d4_transform_atrou(V const *indata, int n, instep,
    V *outlow, int lowstep, V *outhigh, int highstep):
    """
    Given a signal of n samples, return low and highband of the wavelet
    transform, without any downsampling.
    """
    if not got_constants: daubechies_d4_init()
    V const *v = indata
    V *vl = outlow, *vh = outhigh
    int s1 = instep
    int s2 = instep * 2
    int s3 = instep * 3
    int ii = 0, li = 0, hi = 0
    for int i = 0; i < n - 3; i++:
        vl[li] = v[ii] * h0 + v[ii + s1] * h1 + v[ii + s2] * h2 + v[ii + s3] * h3
        vh[hi] = v[ii] * g0 + v[ii + s1] * g1 + v[ii + s2] * g2 + v[ii + s3] * g3
        ii += s1
        li += lowstep
        hi += highstep
    # What to do at the right edge where v[i + 2] and v[i + 3] do not exist?
    # Here we wrap around as the iris actually is circular.
    vl[li] = v[s1 * (n - 3)] * h0 + v[s1 * (n - 2)] * h1 + v[s1 * (n - 1)] * h2 + v[0] * h3
    vh[hi] = v[s1 * (n - 3)] * g0 + v[s1 * (n - 2)] * g1 + v[s1 * (n - 1)] * g2 + v[0] * g3
    li += lowstep
    hi += highstep
    vl[li] = v[s1 * (n - 2)] * h0 + v[s1 * (n - 1)] * h1 + v[0] * h2 + v[s1] * h3
    vh[hi] = v[s1 * (n - 2)] * g0 + v[s1 * (n - 1)] * g1 + v[0] * g2 + v[s1] * g3
    li += lowstep
    hi += highstep
    vl[li] = v[s1 * (n - 1)] * h0 + v[0] * h1 + v[s1] * h2 + v[s2] * h3
    vh[hi] = v[s1 * (n - 1)] * g0 + v[0] * g1 + v[s1] * g2 + v[s2] * g3

def daubechies_d4_init():
    double r3 = sqrt(3)
    # note: sqrt(2) factor included
    double r32 = sqrt(32) * sqrt(2)

    h0 = (1 + r3) / r32
    h1 = (3 + r3) / r32
    h2 = (3 - r3) / r32
    h3 = (1 - r3) / r32

    g0 = h3
    g1 = -h2
    g2 = h1
    g3 = -h0

    got_constants = true

V **def d4(V const *s, int n):
    """
    This expects a signal of n samples, and will return (full-length, not
    downsampled) low and high pass results for each subdivision level.
    """
    int l = ceil(log2(n));
    V **r = ALLOCATE(l * 2 * sizeof(V *))

    V const *p = s
    for int i = 0; i < l; i++:
        V outlow[n]
        V outhigh[n]
        daubechies_d4_transform_atrou(p, n, 1, outlow, 1, outhigh, 1)

        r[i * 2] = ALLOCATE(n * sizeof(V))
        memcpy(r[i * 2], outlow, n * sizeof(V))

        r[i * 2 + 1] = ALLOCATE(n * sizeof(V))
        memcpy(r[i * 2 + 1], outhigh, n * sizeof(V))

        # In each step, re-use the low band. Also, ignore the artificial
        # upsampling..
        p = r[i * 2]

    return r

Image *def d4_2d(Image *v, int levels):
    """
    Given an image, performs a number of wavelet subdivisions and return the
    result in a new image of the same size.
    """
    Image *temp = image_new(None, v->w, v->h, v->type)
    Image *result = image_new(None, v->w, v->h, v->type)

    Image *prev = v

    int h = v->h
    int w = v->w
    for int t = 0; t < levels; t++:
        # Transform each previous low-pass row to lowpass/highpass.
        for int i = 0; i < h; i++:
            daubechies_d4_transform(prev->data.v + i * v->w, w, 1,
                temp->data.v + i * v->w, 1,
                temp->data.v + i * v->w + w / 2, 1)

        # Transform each column of the horizontally transformed image.
        for int i = 0; i < w; i++:
            daubechies_d4_transform(temp->data.v + i, h, v->w,
                result->data.v + i, v->w,
                result->data.v + i + v->w * (h / 2), v->w)

        prev = result

        w /= 2
        h /= 2

    image_del(temp)
    return result

Image *def d4_2d_atrou(Image *v, int levels):
    """
    Given an image, performs a wavelet transformation, but with redundancy.
    That is, no downsampling is performed. The resulting image will have
    dimensions (w, h * (levels * 3 + 1)).
    """
    daubechies_d4_init()

    Image *result = image_new(None, v->w, v->h * (levels * 3 + 1), v->type)
    Image *temp_low = image_new(None, v->w, v->h, v->type)
    Image *temp_high = image_new(None, v->w, v->h, v->type)

    Image *prev = v

    int pos = 0

    for int t = 0; t < levels; t++:
        # Transform each original row to lowpass/highpass.
        for int i = 0; i < v->h; i++:
            daubechies_d4_transform_atrou(prev->data.v + (pos + i) * v->w, v->w, 1,
                temp_low->data.v + i * v->w, 1,
                temp_high->data.v + i * v->w, 1)

        # Transform each column of the row-highpass.
        for int i = 0; i < v->w; i++:
            daubechies_d4_transform_atrou(temp_high->data.v + i, v->h, v->w,
                result->data.v + i + v->w * (pos + v->h), v->w,
                result->data.v + i, v->w * (pos))

        # Transform each column of row-lowpass.
        for int i = 0; i < v->w; i++:
            daubechies_d4_transform_atrou(temp_low->data.v + i, v->h, v->w,
                result->data.v + i + v->w * (pos + v->h * 3), v->w,
                result->data.v + i + v->w * (pos + v->h * 2), v->w)

        prev = temp_low
        pos += v->h * 3
       
    image_del(temp_low)
    image_del(temp_high)
    return result
