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/wptree

"""
Base algorithm mostly following Zhu, Tan, Wang.
Our input textures are 512x64.
With pyramidal wavelet subdivision, we can do 6 subdivisions before ending up
with 8x1-pixel strips. (Which could be further subdivided with 3 1D transforms.)

We limit subdivision down to 16x2 areas though, which means 5 subdivision
levels. In the pyramidal case, this means:

   
   128 256     512
_ _.___._______.
X|_|   |(n)    |
_|3|___|       |16
   |2  |       |
___|___|_______|32
(n)    |(n)    |
       |1      |
       |       |
_______|_______|64

(X)
   16  32      64
___.___._______.
   |   |       |
___|___|       | 2
   |5  |       |
___|___|_______| 4
       |4      |
       |       |
       |       |
_______|_______| 8

The 26 features then are mean value and standard deviation from the 13 subbands
not marked with (n), including the smallest 16x2 pixel LL in the top left.

One way to add a key to this with wavelet packets is  to build the complete
wavelet transform, then select 13 sub-bands out of all possibilities. The total
decomposition has: 4^1 + 4^2 + 4^3 + 4^4 + 4^5 = 1364 possible subbands.

[
512 x 64 (no transform): 1 * 1 = 4^0 = 2^0
256 x 32 (n)           : 2 * 2 = 4^1 = 2^2
128 x 16               : 4 * 4 = 4^2 = 2^4
 64 x  8               : 8 * 8 = 4^3 = 2^6
 32 x  4               : 16 * 16 = 4^4 = 2^8
 16 x  2               : 32 * 32 = 4^5 = 2^10
]

Then, we select 13 subbands, so 13 out of 1364 possible keys, or
1364! / (1364 - 13)! = 53416518513355349952044975873073277747200
which is around 2^135. Of course, most of those are rather similar.

We do not directly compute those keys though - instead we use a PRNG to create
a random three out of the above. The PRNG seed is our key.
"""

static V def mean(Image *i, int x, y, w, h):
    return image_mean_area(i, x, y, w, h)

static V def sdev(Image *i, int x, y, w, h):
    float v = image_variance_area(i, x, y, w, h)
    return sqrt(v)

Vector *def get_features_1(Image *wt):
    int w = wt->w
    int h = wt->h
    int x = w
    int y = h
    # build feature vector
    float features[26]
    int fn = 0
    for int i = 0; i < 5; i++:
        w /= 2
        h /= 2
        x -= w
        y -= h
        if i > 0:
            features[fn++] = mean(wt, x, 0, w, h)
            features[fn++] = sdev(wt, x, 0, w, h)
            features[fn++] = mean(wt, x, y, w, h)
            features[fn++] = sdev(wt, x, y, w, h)
            features[fn++] = mean(wt, 0, y, w, h)
            features[fn++] = sdev(wt, 0, y, w, h)
    features[fn++] = mean(wt, 0, 0, w, h)
    features[fn++] = sdev(wt, 0, 0, w, h)
    printf("%d %d\n", w, h)

    Vector *r = vector_new(sizeof(V))

    for int i = 0; i < 26; i++:
        vector_add(r, features + i)
        printf(" %f", features[i])
    printf("\n")

    return r

Vector *def get_features_2(Image *wt):
    int w = wt->w
    int h = wt->h
    int x = 0
    int y = 0
    # build feature vector
    float features[26]
    int fn = 0
    for int i = 0; i < 26; i++:
        features[fn++] = mean(wt, x, y, w, h)
        features[fn++] = sdev(wt, x, y, w, h)
        y += h
  
    Vector *r = vector_new(sizeof(V))

    for int i = 0; i < 26; i++:
        vector_add(r, features + i)

    return r

Vector *def get_random_features(Image *wt, WPNode *tree):
    # Select 13 out of all possible leafs.
    int n = wptree_count_leafs(tree)
    assert(n >= 13)
    bool taken[n]
    CLEAR(taken, n * sizeof(bool))
    int selection[13]
    printf("selection:")
    for int i = 0; i < 13; i++:
        n--
        int x = mersenne_rand(0, n)
        for int j = 0; j < n; j++:
            if taken[j] and j <= x: x++
        selection[i] = x
        taken[x] = True
        printf(" %d", x)
    printf("\n")

    # build feature vector
    float features[26]
    int fn = 0
    for int i = 0; i < 13; i++:
        WPNode *leaf = wptree_select_leaf(tree, selection[i])
        features[fn++] = mean(wt, leaf->y, leaf->y, leaf->w, leaf->h)
        features[fn++] = sdev(wt, leaf->y, leaf->y, leaf->w, leaf->h)

    Vector *r = vector_new(sizeof(V))

    for int i = 0; i < 26; i++:
        vector_add(r, features + i)
        printf(" %f", features[i])
    printf("\n")

    return r

Vector *def algo2_w2d(
    Image *(*wavelet2d)(Image *v, int levels),
    Image *texture, char const *key, *what, *algo):

    Vector *features

    Image *v = BtoV(texture)
    # Center around origin, so mean values of subbands are similar to mean
    # value of picture
    shift_multiply_image(v, -127.5, 1)

    Image *wt = None
    if not key:
        # Following Yong Zhu, Tieniu Tan, Yunhong Wang we use 5 subdivision
        # levels, then use mean and standard deviation of 13 subbands as
        # features, excluding the 3 finest ones.
        wt = wavelet2d(v, 5)
        features = get_features_1(wt)
    else:
        int min_ll_depth = 5
        int max_ll_depth = 5
        int max_depth = 5
        double base_value = 1
        double change_factor = 0

        # We assume our key is an 8-digit hexadecimal number, e.g.
        # "00000000" or "deadbeef"
        uint32_t keynum = strtol(key, None, 16)
        mersenne_seed(keynum)

        Wavelet *wavelet = ALLOCATE0(sizeof *wavelet)
        wavelet->transform = daubechies_d4_transform_no_key
        bool is_1d = False
        WPNode *tree = wptree_create_random_tree(wavelet, 512, 64,
            min_ll_depth, max_ll_depth, max_depth,
            base_value, change_factor, is_1d)
        wt = wptree_transform(tree, v)

        if hasword(what, "wptree"):
            char name[MAX_PATH]
            get_image_name_algo_key_ext(texture, "wptree", algo, key, "eps", name)
            wptree_plot(tree, name)
            printf("Wrote %s.\n", name)

        features = get_random_features(wt, tree)            

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

        Image *wt2 = image_copy(wt)
        # Shift back (just for display, so black <= -127.5, white >= 127.5)
        shift_multiply_image(wt2, 127.5, 1)
        Image *wtb = VtoB(wt2)
        save_tiff_luminance(name, wtb)

    return features

Vector *def algo2(Image *texture, char const *key, *what, *algo):
    return algo2_w2d(d4_2d, texture, key, what, algo)

Vector *def algo2b(Image *texture, char const *key, *what, *algo):
    return algo2_w2d(d4_2d_atrou, texture, key, what, algo)
