import numpy, math

__dctNxM = {}

def __init_dctNxM(b):
    global __dctNxM
    if not __dctNxM.has_key(b.shape): 
        c = numpy.sqrt([2.0 / b.shape[0], 2.0 / b.shape[1]])
        __cosf0 = lambda j, i: numpy.cos((math.pi * (2*i + 1) * j) / (2.0 * b.shape[0]))
        y = c[0] * numpy.fromfunction(__cosf0, (b.shape[0], b.shape[0]))
        __cosf1 = lambda j, i: numpy.cos((math.pi * (2*i + 1) * j) / (2.0 * b.shape[1]))
        x = c[1] * numpy.fromfunction(__cosf1, (b.shape[1], b.shape[1]))
        __dctNxM[b.shape] = (y, x), c / math.sqrt(2.0)
    return __dctNxM[b.shape]
    
def fdct2d(b):
    '''
    Compute forward 2D DCT on the array of arbitrary size.
    Note: slow.
    '''
    global __dctNxM
    ct, c = __init_dctNxM(b)
    
    d = numpy.empty(b.shape, numpy.float)
    d[0][0] = c[0] * c[1] * numpy.sum(numpy.sum(b))

    for i in xrange(1, b.shape[1]):
        d[0][i] = c[0] * numpy.sum(ct[1][i] * numpy.sum(b, 0))

    for j in xrange(1, b.shape[0]):
        d[j][0] = c[1] * numpy.sum(ct[0][j] * numpy.sum(b, 1))

    bt = numpy.transpose(b)
    for j in xrange(1, b.shape[0]):
        d[j,1:] = numpy.sum(ct[1][1:] * numpy.sum(bt * ct[0][j], 1), 1)
    
    return d

def idct2d(b):
    '''
    Compute inverse 2D DCT on array of arbitrary size.
    '''
    global __dctNxM
    ct, c = __init_dctNxM(b)

    p0 = c[0] * c[1] * b[0][0]
    bt = numpy.transpose(b[1:,1:])
    ctt = ct[1][1:,:].transpose()
    d = numpy.empty(b.shape, numpy.float)
    for y in xrange(b.shape[0]):
        d[y,:] = p0 + \
            c[0] * numpy.sum(b[0,1:] * ctt, 1) + \
            c[1] * numpy.sum(b[1:,0] * ct[0][1:,y]) + \
            numpy.sum(numpy.sum(bt * ct[0][1:,y],1) * ctt,1)
    return d