import numm, numpy

# watch images disintegrate into patterns.
# - rmo@NUMM.ORG

import sys
images = sys.argv[1:]
images = [numm.image2np(X)[:240,:320] for X in images]

#img = numm.image2np('a.jpg')[:240,:320]
img_idx = 0
img = images[img_idx]

fft = numpy.fft.fft2(img)
ifft = None
x_idx = 0

rweights = numpy.ones(320, dtype=numpy.int)
iweights = numpy.ones(320, dtype=numpy.int)

modes = ['ncols', 'nop', 'rset', 'nop', 'iset', 'nop', 'reset', 'nop']
mode_idx = 0

a_idx = 0

ncols = 0

def video_out(a):
    global ifft
    # fft[:,x_idx] *= 1.1
    # fft[:,(319-x_idx)] /= 1.1

    fft2 = fft.copy()
    for i in range(3):
        fft2.real[:,:,i] *= pow(rweights, 0.5)
        fft2.imag[:,:,i] *= pow(iweights, 0.5)
    
    cutoff = 500

    ifft = numpy.fft.ifft2(fft2)
    a[:] = ifft.clip(0,255)

    a[-5:, numpy.logical_and(rweights==0, iweights == 0)] = (255, 0, 0)
    a[-5:, numpy.logical_and(rweights==1, iweights == 1)] = (0, 255, 0)
    a[-5:, numpy.logical_or(rweights>1, iweights > 1)] = (0, 0, 255)

    a[rweights,range(320)] = (255,255,0)
    a[iweights,range(320)] = (0,255,255)

    # overflow == psychodelic?
    # a[:] = numpy.fft.ifft2(fft).real

def mouse_in(type, px, py, button):
    global x_idx, rmove, ncols, mode_idx

    if type == 'mouse-button-press':
        if button == 1:
            mode_idx = (mode_idx + 1) % (len(modes))
        else:
            mode_idx = (mode_idx - 1) % (len(modes))
        print modes[mode_idx]

    x_idx = int(px*320)

    if modes[mode_idx] == 'ncols':
        p_ncols = ncols
        ncols = x_idx

        fmean = fft.mean(axis=0).mean(axis=1)

        if ncols < p_ncols:
            rweights[abs(fmean).argsort()[ncols:p_ncols]] = 1
            iweights[abs(fmean).argsort()[ncols:p_ncols]] = 1

        rweights[abs(fmean).argsort()[:ncols]] = 0
        iweights[abs(fmean).argsort()[:ncols]] = 0
    elif modes[mode_idx] == 'rset':
        rweights[x_idx] = int(py*240)
    elif modes[mode_idx] == 'iset':
        iweights[x_idx] = int(py*240)
    elif modes[mode_idx] == 'reset':
        rweights[x_idx] = 1
        iweights[x_idx] = 1

def keyboard_in(type, key):
    # reset & move to next image
    global fft, img, img_idx, rweights, iweights
    
    if type == 'key-press':

        img_idx = (img_idx + 1) % len(images)
        img = images[img_idx]

        fft = numpy.fft.fft2(img)
        rweights = numpy.ones(320, dtype=numpy.int)
        iweights = numpy.ones(320, dtype=numpy.int)


def audio_out(a):
    global a_idx
    nframes = numpy.prod(a.shape)
    if nframes > (numpy.prod(ifft.shape)-a_idx):
        a_idx = 0

    a[:,0].flat = 2**15*(ifft.real.flat[a_idx:a_idx+nframes]/ifft.real.max())
    a[:,1].flat = 2**15*(ifft.imag.flat[a_idx:a_idx+nframes]/ifft.imag.max())
    a_idx += nframes