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)

import dct, numpy

def fdctNxN(b):
    '''
    Forward 2D DCT on NxN array.
    '''
    c = __init_dctN(b.shape[0])
    return numpy.dot(c, numpy.dot(b, numpy.transpose(c)))

def idctNxN(b):
    '''
    Inverse 2D DCT on NxN array.
    '''
    c = __init_dctN(b.shape[0])
    return numpy.dot(numpy.transpose(c), numpy.dot(b, c))

def zigzag(x, start=0, stop=None):
    """Generates zig-zag scan sequence for two dimensional array."""
    i = j = count = 0
    if not stop: stop = x.size
    while count < stop:
        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
            if start <= count < stop: yield i, j
            i -= 1
            j += 1
            count += 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
            if start <= count < stop: yield i, j
            i += 1
            j -= 1
            count += 1

        if i > x.shape[0]: break