import numpy, math

__dctN = {}

def __init_dctN(n):
    global __dctN
    if not __dctN.has_key(n): 
        c = numpy.empty((n, n), numpy.float)
        c[0,:] = 1.0 / math.sqrt(n)
        __cosf = lambda i, j: numpy.cos(math.pi * (2*j + 1) * (i + 1) / (2.0 * n))
        c[1:,:] = math.sqrt(2.0 / n) * numpy.fromfunction(__cosf, (n-1, n))
        __dctN[n] = c
    return __dctN[n]

def fdct(v):
    '''
    Forward DCT on vector.
    '''
    c = __init_dctN(len(v))
    return numpy.dot(c, v)

def idct(v):
    '''
    Inverse DCT on vector.
    '''
    c = __init_dctN(len(v))
    return numpy.dot(numpy.transpose(c), v)

def zigzag(x):
    """Generates zig-zag scan sequence for two dimensional array."""
    j = 0
    i = 0
    while True:

        if i == x.shape[0]: j += 2; i = x.shape[0]-1
        else: j = 0
        while i >= 0 and j < x.shape[1]:
            # run up
            yield i, j
            i -= 1
            j += 1

        if j > x.shape[1]: break

        if j == x.shape[1]: j = x.shape[1]-1; i += 2
        else: i = 0
        while j >= 0 and i < x.shape[0]:
            # run down
            yield i, j
            i += 1
            j -= 1

        if i > x.shape[0]: break