import config, filter, random
import wavelets/cdf_9x7

# Code from Peter Meerwald
static def _filter_recurse(Filter f, int n, V const *alpha):
    f.h[0] = 1.0 / sqrt(2.0)
    f.h[1] = 1.0 / sqrt(2.0)
    f.g[0] = -1.0 / sqrt(2.0)
    f.g[1] = 1.0 / sqrt(2.0)

    for int k = 0; k < n; k++: 
        for int i = 2 * (k + 1) + 1; i >= 0; i--:
            V h1, g1, h2, g2

            h1 = (i <= 2 * k + 1) ? f.h[i] : 0.0
            g1 = (i <= 2 * k + 1) ? f.g[i] : 0.0
            h2 = (i - 2 < 0) ? 0.0 : f.h[i - 2]
            g2 = (i - 2 < 0) ? 0.0 : f.g[i - 2]
            # g2 = (2 * (i & 1)-1) * h2;

            f.h[i] = cos(alpha[k]) * (cos(alpha[k]) * h1 + sin(alpha[k]) * g1) +\
                sin(alpha[k]) * (sin(alpha[k]) * h2 - cos(alpha[k]) * g2)
            f.g[2 * (k + 1) - i + 1] = (2 * (i & 1)-1) * f.h[i]

Filter *def zou_get_filter(int fn, V *alpha):
    Filter *f = ALLOCATE0(sizeof *f)
    f->hn = 2 * fn + 2
    f->gn = 2 * fn + 2
    f->h = malloc(f->hn * sizeof(V))
    f->g = malloc(f->gn * sizeof(V))
    _filter_recurse(*f, fn, alpha)
    return f

def zou_nosubsample_param(V const *indata, int n, instep,
    V *outlow, int lowstep, V *outhigh, int highstep, scale, fn, V const *alpha):

    Filter f
    f.hn = f.gn = 2 * fn + 2
    V h[f.hn] # lowpass
    V g[f.gn] # highpass
    f.h = h
    f.g = g

    _filter_recurse(f, fn, alpha)

    filter(indata, n, instep, h, f.hn, outlow, lowstep, 1 << scale)
    filter(indata, n, instep, g, f.gn, outhigh, highstep, 1 << scale)

V **def zou(V const *s, int n, char const *key):
    """
    This expects a signal of n samples, and will return (full-length, not
    downsampled) low and high pass results for each subdivision level.

    0 2 4.../log2(n)\ * 2 - 2 are the low pass results
    1 3 5.../log2(n)\ * 2 - 1 the high pass results

    """
    int l = ceil(log2(n));
    V **r = ALLOCATE(l * 2 * sizeof(V *))

    int kn = strlen(key) / 2
    V *a = malloc(kn * sizeof *a)
    for int i = 0; i < kn; i++:
        a[i] = ((key[i * 2] - '0') * 10 + key[i * 2 + 1] - '0') / 99.0 * 2 * pi - pi

    V const *p = s
    for int i = 0; i < l; i++:
        r[i * 2] = ALLOCATE(n * sizeof(V))
        r[i * 2 + 1] = ALLOCATE(n * sizeof(V))
        zou_nosubsample_param(p, n, 1, r[i * 2], 1, r[i * 2 + 1],
            1, i, kn, a)
        # right now, we always use the lowband for further analysis.
        p = r[i * 2]

    return r

def zou_transform(V const *indata, int n, instep,
    V *outlow, int lowstep, V *outhigh, int highstep, char const *key):
    # Test function.

    if not key:
        key = "00"

    int kn = strlen(key) / 2
    V *a = malloc(kn * sizeof *a)
    for int i = 0; i < kn; i++:
        a[i] = ((key[i * 2] - '0') * 10 + key[i * 2 + 1] - '0') / 99.0 * 2 * pi - pi

    # TODO: We want to manually downscale, so need our own intermediate
    # results, then write the result of that to the output..
    V templow[n], temphigh[n]
    int scale = 1 # TODO: make sure we can do this
    zou_nosubsample_param(indata, n, instep, templow, 1, temphigh, 1,
        scale, kn, a)
    # Now get results to the output (Is averaging the right thing? This is only
    # a test function anyway though..)
    V *lo = outlow
    V *hi = outhigh
    for int i = 0; i < n; i += 2:
        *lo = (templow[i] + templow[i + 1]) * 0.5
        *hi = (temphigh[i] + temphigh[i + 1]) * 0.5
        lo += lowstep
        hi += highstep

def zou_once(V *signal, int n, V **out,
    unsigned int subband, int param_n, V const *params,
    bool randomize):
    uint32_t m = 1

    # 0x1f
    while (1u << m) <= subband:
        m++
    m -= 1
    subband &= (1 << m) - 1

    V temp[n]
    V const *p = signal
    int read = 1 << (m - 1)
    for uint32_t i = 0; i < m; i++:
        V result_lo[n]
        V result_hi[n]
        V const *real_params = params
        if randomize:
            V ps[param_n]
            for int j = 0; j < param_n; j++:
                # FIXME: parameter range...
                #ps[j] = mersenne_rnd(pi * 0.21, pi * 0.29)
                ps[j] = mersenne_rnd(-pi, pi)
            real_params = ps

        if signal:
            #for int j = 0; j < param_n; j++:
            #    if j == 0: print("<")
            #    print("%.1f", real_params[j])
            #    if j == param_n - 1: print(">")
            #    else: print(" ")
                    
            zou_nosubsample_param(p, n, 1, result_lo, 1, result_hi, 1,
                i, param_n, real_params)

        if subband & read:
            #print("H")
            memcpy(temp, result_hi, n * sizeof(V))
        else:
            #print("L")
            memcpy(temp, result_lo, n * sizeof(V))
        p = temp
        read >>= 1
    # if m > 0: print("\n")

    *out = ALLOCATE(n * sizeof(V))
    if p:
        memcpy(*out, p, n * sizeof(V))

int def random_subband():
""" We do not want the following:
    - starts with HH (64 + 32 + 16 + 8 + 4 + 2 + 1 = 127)
    - 7 or 8 levels and starts with LH (32 + 64 = 96)
    - 8 levels and starts with H (128)
    - starts with LHL (32 + 16 + 8 + 4 + 2 + 1 = 63)
"""
    # biggest value is LLHHHHHH as H....... and LH...... are excluded
    int r = mersenne_rand(0x002, 0x13f)
    int highest_bit = 0
    while (1 << highest_bit) <= r:
        highest_bit += 1
    highest_bit -= 1
    # HH
    if highest_bit >= 2:
        int XX = r >> (highest_bit - 2)
        if XX == 0x7: return random_subband()
    # LHL
    if highest_bit >= 3:
        int XXX = r >> (highest_bit - 3)
        if XXX == 0xa: return random_subband()
    if highest_bit == 7:
        if r & (1 << (highest_bit - 1)): return random_subband()
    return r
        

def zou_twice_from_key(V *signal, int n, V ***dwt, char const *key):
    V **two = ALLOCATE(2 * sizeof(V *))
    int subband1 = 0x005
    int subband2 = 0x009
    int param_length = 1
    V dummy[] = {0, 0}
    V *params = dummy
    bool randomize = False

    if key[0] == 'p':
        uint32_t hex = strtol(key + 1, NULL, 16)
        subband1 = hex & 4095
        subband2 = (hex >> 12) & 4095
    elif key[0] == 'P':
        subband1 = random_subband()
        subband2 = random_subband()
    elif key[0] >= '0' and key[0] <= '9':
        int kn = strlen(key) / 2
        V *a = malloc(kn * sizeof *a)
        for int i = 0; i < kn; i++:
            a[i] = ((key[i * 2] - '0') * 10 + key[i * 2 + 1] - '0') / 99.0 * 2 * pi - pi
        param_length = kn
        params = a
    elif key[0] == 'r' or key[0] == 'R' or key[0] == '*':
        if key[1] == '1':
            params[0] = mersenne_rnd(-pi, pi)
        elif key[1] == '2':
            params[0] = mersenne_rnd(-pi, pi)
            params[1] = mersenne_rnd(-pi, pi)
            param_length = 2
        if key[0] == 'R':
            subband1 = random_subband()
            subband2 = random_subband()
        if key[0] == '*':
            subband1 = random_subband()
            subband2 = random_subband()
            randomize = True
    else:
        fprintf(stderr, "Fatal error: Invalid key: \"%s\"\n", key)
        return

    #print(" subband1: %03x, params: %.2f %.2f %s\n", subband1, params[0], params[1],
    #    randomize ? "*" : "-")
    zou_once(signal, n, two + 0, subband1, param_length, params, randomize)
    #print(" subband2: %03x, params: %.2f %.2f %s\n", subband2, params[0], params[1],
    #    randomize ? "*" : "-")
    zou_once(signal, n, two + 1, subband2, param_length, params, randomize)

    *dwt = two

def zou_ten_twice(V *s, int N, w, V ***dwts, char const *key):
    """
    Perform two wavelet transforms on N signals, depending on the key.

    The key is either:
    - If the first character is a digit 0 to 9:
      The key is a sequence of 2-digit decimal numbers. The range 00..99 for
      each number is mapped to -pi..+pi, and the amount determines the filter
      length. All signals use the same filter. Example:
      00 filter length 1, parameter 0
[TODO
    - If the first character is h:
      Like above, but groups of 3 hex digits are used, for a resolution of
      4096 of the -pi..pi range. Example:
      hfff filter length 1, parameter 3.14
]
    - If the first character is r:
      The last 8 hex digits are used as 32-bit seed for a PRNG. The character
      after r is:
      1 for a different random length 1 filter for each signal
      2 for length 2 filters
      Examples:
      r*deadbeef
    - If the first character is p:
      The key consists of 6 hex digits specifying two subbands to use from a
      WP decomposition. They are interpreted as 2 groups of 12 bits.
      The highest bit in those 12 determines the level of the subband to use,
      the lower bits determine the position.
      000: 0000 0000 0000 invalid
      001: 0000 0000 0001 no wavelet transform
      002: 0000 0000 0010 1 decomposition, use low band
      003: 0000 0000 0011 1 decomposition, use high band
      004: 0000 0000 0100 2 decompositions, use LL
      005: 0000 0000 0101 2 dec, use LH
      009: 0000 0000 1001 3 dec, use LLH
      011: 0000 0001 0001 4 dec, use LLLH
      0ff: 0000 1111 1111 7 dec, use HHHHHHH
      100: 0001 0000 0001 8 dec, use LLLLLLLH
      All signals use the same decomposition.
      Example:
      p009011 will use LLH and LLLH
    - If the first character is P, 8 hex digits define a random seed, and a
      random WP selection is done for each signal. Structures like these are
      never generated:
      starts with HH (64 + 32 + 16 + 8 + 4 + 2 + 1 = 127)
      7 or 8 levels and starts with LH (32 + 64 = 96)
      8 levels and starts with H (128)
      starts with LHL (32 + 16 + 8 + 4 + 2 + 1 = 63)
      Example:
      Pdeadbeef
    - If the first character is R, it's as if using r and P. That is, each
      signal uses a random subband, and a random filter. Filter length is
      determined by the second character.
    - If the first character is *, it's as if using r and P, but the filter
      also varies in each WP subdivision (non-stationary decomposition).

    dwts is an array with N elements. We fill in the elements with pointers to
    arrays of length 2, pointing to two transforms, as selected by the key.
    """
    if not key:
        key = "p009011"

    if key[0] == 'r' or key[0] == 'R' or key[0] == '*':
        uint32_t keynum = strtol(key + 2, None, 16)
        mersenne_seed(keynum)
    elif key[0] == 'P':
        uint32_t keynum = strtol(key + 1, None, 16)
        mersenne_seed(keynum)

    for int i = 0; i < N; i++:
        V *p = s
        if p:
            p += i * w
        zou_twice_from_key(p, w, dwts + i, key)