import config, image
import random

class Wavelet:
    """Interface for a wavelet transform."""
    char *name

    # Method to transform a one-dimensional signal.
    void (*transform)(V const *indata, int n, int instep,
        V *outlow, int lowstep, V *outhigh, int highstep, char const *key)

    char *key # For parameterized transform

class WPNode:
    """
    One node in a wavelet packets tree. Each node can use a different wavelet
    transform.
    """
    WPNode *parent
    # The order for the children is:
    #  _____________
    # | 0 LL | 1 HL |
    # |______|______|
    # | 2 LH | 3 HH |
    # |______|______|
    #
    # A leaf has no children.
    bool is_leaf # in this case, child is ignored
    bool is_1d # in this case, only child[0, 1] are valid
    WPNode *child[4]
    Wavelet *hwavelet, *vwavelet
    # For leafs, level means:
    # 0 - the leaf is the original, untransformed image
    # 1 - the leaf is the result of 1 decomposition (4 possible subbands)
    # and so on..
    int level
    int x, y, w, h

static def _recurse_pyramidal(WPNode *parent, int max_level,
    Wavelet *(*callback)(int level, bool vertical, void *data), void *data):
    for int i = 0; i < 4; i++:
        if parent->is_1d and i == 2: break
        WPNode *node = parent->child[i] = ALLOCATE0(sizeof *node)
        node->level = parent->level + 1
        node->parent = parent
        node->w = parent->w / 2
        node->h = parent->h / 2
        node->x = parent->x
        node->y = parent->y
        node->is_1d = parent->is_1d
        if i == 1 or i == 3: node->x += node->w
        if i == 2 or i == 3: node->y += node->h
        node->hwavelet = callback(node->level, False, data)
        if not node->is_1d:
            node->vwavelet = callback(node->level, True, data)
        if i != 0 or node->level == max_level:
            node->is_leaf = True
        else:
            _recurse_pyramidal(node, max_level, callback, data)

WPNode *def wptree_create_pyramidal_tree_with_callback(int max_level, w, h,
    bool is_1d,
    Wavelet *(*callback)(int level, bool vertical, void *data), void *data):
    """
    Like [wptree_create_pyramidal_tree], but with a callback function querying
    the wavelet to use instead of the same wavelet for each subband. This can
    be used to use different wavelets or wavelet parameterizations for each
    subband.
    """
    WPNode *root = ALLOCATE0(sizeof *root)
    root->hwavelet = callback(0, False, data)
    root->vwavelet = callback(0, True, data)
    root->w = w
    root->h = h
    root->level = 0
    if max_level == 0:
        root->is_leaf = True
        return root

    _recurse_pyramidal(root, max_level, callback, data)
    return root

static Wavelet *def _static_callback(int level, bool vertical, void *data):
    return data

WPNode *def wptree_create_pyramidal_tree(Wavelet *wavelet, int max_level, w, h,
    bool is_1d):
    """
    Return a pyramidal tree, with the same wavelet for each subband.
    If max_level is 0, a leaf node is returned (no wavelet decomposition).
    If max_level is 1, only a single transform into 4 subbands is done. And so
    on.
    """
    return wptree_create_pyramidal_tree_with_callback(max_level, w, h,
        is_1d, _static_callback, wavelet)

static bool def _random_decision(double base_value, change_factor, int cur_depth):
    double x = mersenne_rnd_open(0, 2)
    double weight = base_value + change_factor * cur_depth
    return x / weight >= 1

static def _recurse_random(WPNode *parent, int min_ll_depth, max_ll_depth,
    max_depth, double base_value, change_factor,
    Wavelet *(*callback)(int level, bool vertical, void *data), void *data):
    """
    Decomposes the given parent node into 4 child nodes.
    The depth parameters always refer to the leaf level.
    """
    for int i = 0; i < 4; i++:
        if parent->is_1d and i == 2: break
        WPNode *node = parent->child[i] = ALLOCATE0(sizeof(WPNode))
        node->parent = parent
        node->level = parent->level + 1
        node->w = parent->w / 2
        node->h = parent->h / 2
        node->x = parent->x
        node->y = parent->y
        node->is_1d = parent->is_1d
        if i == 1 or i == 3: node->x += node->w
        if i == 2 or i == 3: node->y += node->h
        node->hwavelet = callback(node->level, False, data)
        if not node->is_1d:
            node->vwavelet = callback(node->level, True, data)

        bool decompose = False
        if i == 0 and parent->x == 0 and parent->y == 0:
            if node->level < min_ll_depth:
                decompose = True
            elif node->level >= max_ll_depth:
                pass
            else:
                decompose = _random_decision(base_value, change_factor, node->level)
        else:
            if node->level >= max_depth:
                pass
            else:
                decompose = _random_decision(base_value, change_factor, node->level)       

        if decompose:
            #for int j = 0; j < cur_depth; j++: printf(" ")
            #printf("decompose: %d\n", i)            
            _recurse_random(node, min_ll_depth, max_ll_depth,
                max_depth, base_value, change_factor, callback, data)
        else:
            node->is_leaf = True

WPNode *def wptree_create_random_tree_with_callback(int w, h,
    int min_ll_depth, max_ll_depth, max_depth,
    double base_value, change_factor, bool is_1d,
    Wavelet *(*callback)(int level, bool vertical, void *data), void *data):
    """
    What we do here is, we use a PRNG to take a decision at each
    decomposition level if and how to subdivide (following Pommer/Uhl 2003).

    If max_depth is 0, a leaf node is returned (no decomposition).
    if max_depth is 1, the only possible decomposition is the first level. And
    so on.
    The max_ll_depth parameter is the same as max_depth, but for the LL band.

    The min_ll_depth parameter gives the minimum depth for the approximation
    subband. A value of 1 means, the first decomposition
    will always be done. A value of 2 means, the LL band from the first
    subdivision will be divided again, and so on.
    """
    WPNode *root = ALLOCATE0(sizeof *root)
    root->is_1d = is_1d
    root->hwavelet = callback(0, False, data)
    root->vwavelet = callback(0, True, data)
    root->w = w
    root->h = h
    root->parent = None

    bool decompose = False
    if root->level < min_ll_depth:
        decompose = True
    elif root->level >= max_ll_depth:
        pass
    else:
        decompose = _random_decision(base_value, change_factor, root->level)
    if decompose:
        _recurse_random(root, min_ll_depth, max_ll_depth, max_depth, base_value,
            change_factor, callback, data)
    else:
        root->is_leaf = True
    return root

WPNode *def wptree_create_random_tree(Wavelet *wavelet, int w, h,
    int min_ll_depth, max_ll_depth, max_depth,
    double base_value, change_factor, bool is_1d):
    return wptree_create_random_tree_with_callback(w, h,
        min_ll_depth, max_ll_depth, max_depth, base_value, change_factor, is_1d,
        _static_callback, wavelet)

static def _recurse_plot(WPNode *node, FILE *f, int pw, ph, float gray):
    if node->is_leaf: return

    int x, y, w, h, w2, h2
    # Recurse to subbands.
    gray -= 0.05
    for int i = 0; i < 4; i++:
        WPNode *c = node->child[i]
        if c->is_leaf: continue
        x = c->x
        y = c->y
        w = c->w
        h = c->h
        w2 = w / 2
        h2 = h / 2

        fprintf(f, "%f setgray\n", gray)
        fprintf(f, "newpath\n")
        fprintf(f, "%d %d moveto %d %d lineto\n", x, ph - y, x + w, ph - y)
        fprintf(f, "%d %d lineto %d %d lineto\n", x + w, ph - y - h, x, ph - y - h)
        fprintf(f, "closepath fill\n")

        _recurse_plot(c, f, pw, ph, gray)

    # Draw dividing cross.
    x = node->x
    y = node->y
    w = node->w
    h = node->h
    w2 = w / 2
    h2 = h / 2

    fprintf(f, "newpath\n")
    fprintf(f, "%d %d moveto %d %d lineto\n",
        x, ph - y - h2, x + w, ph - y - h2)
    fprintf(f, "%d %d moveto %d %d lineto\n",
        x + w2, ph - y, x + w2, ph - y - h)
    fprintf(f, "0 setgray stroke\n")

def wptree_plot(WPNode *root, char const *filename):
    FILE *f = fopen(filename, "w")
    int w = root->w
    int h = root->h
    float gray = 0.9
    fprintf(f, "%%!PS-Adobe-2.0 EPSF-2.0\n%%%%BoundingBox: 0 0 %d %d\n", w, h)
    fprintf(f, "%%%%EndComments\n")
    fprintf(f, "%f setgray\n", gray);
    fprintf(f, "0 setlinewidth\n");
    fprintf(f, "newpath 0 0 moveto %d 0 lineto %d %d lineto 0 %d lineto",
        w, w, h, h)
    fprintf(f, " closepath gsave fill\n")

    _recurse_plot(root, f, w, h, gray)

    fprintf(f, "grestore 0 setgray stroke\n")

    fclose(f)

static def _recurse_transform(WPNode *node, Image *v, *temp, *result):
    """
    Transform the area corresponding to the given node into its four sub-bands,
    and then recurse. The result image can be the same as the v image, the
    temp image must be a different one for the intermediate step of just one
    transform.
    """
    int w = node->w
    int h = node->h
    int x = node->x
    int y = node->y
    int w2 = w / 2
    int h2 = h / 2

    # Horizontal transform: Each row is split into low/highpass
    for int i = y; i < y + h; i++:
        node->hwavelet->transform(v->data.v + i * v->w + x, w, 1,
            temp->data.v + i * v->w + x, 1,
            temp->data.v + i * v->w + x + w2, 1, node->hwavelet->key)

    if not node->is_1d:
        # Vertical transform: Each columns is split into low/highpass
        for int i = x; i < x + w; i++:
            node->vwavelet->transform(temp->data.v + y * v->w + i, h, v->w,
                result->data.v + y * v->w + i, v->w,
                result->data.v + (y + h2) * v->w + i, v->w, node->vwavelet->key)

    for int i = 0; i < 4; i++:
        WPNode *c = node->child[i]
        if c->is_leaf: continue
        if c->is_1d and i == 2: break
        _recurse_transform(c, result, temp, result)

Image *def wptree_transform(WPNode *root, Image *v):
    """
    Given an image, performs a wavelet packet decomposition and returns the
    result in a new image of the same size.
    """

    if root->is_leaf: return None

    Image *temp = image_new(None, v->w, v->h, v->type)
    Image *result = image_new(None, v->w, v->h, v->type)

    _recurse_transform(root, v, temp, result)

    image_del(temp)
    return result

int def wptree_count_leafs(WPNode *root):
    # The whole tree is just one leaf.
    if root->is_leaf: return 1
    int n = 0
    for int i = 0; i < 4; i++:
        n += wptree_count_leafs(root->child[i])
    return n

WPNode *def _select_leaf(WPNode *root, int n, int *c):
    if root->is_leaf:
        if *c == n:
            return root
        (*c)++
        return None
    for int i = 0; i < 4; i++:
        WPNode *found = _select_leaf(root->child[i], n, c)
        if found: return found
    return None

WPNode *def wptree_select_leaf(WPNode *root, int n):
    int c = 0
    return _select_leaf(root, n, &c)
