#
# figuring out the correct forward/reverse procedure
#

# demonstrate FFT and inverse FFT on a grayscale image.
# 
# need:
#   the fftw library installed as a shared library (on OS X, use "configure --enable-shared")
#   numpy
#
# I wrote this while trying to remove patterned noise from some scans of old photos.
# I never quite figured out exactly *what* to do to the complex FFT image in the star-shaped
#   areas... I if make the slightest change at all to the array it will damage the image.
#
# Having since discovered ImageJ (which thankfully comes with source), I'm fairly certain that
#   those areas are just being set to zero (masked out), but I get different results.
#
# If you can explain it to me, send me email!
# http://www.nightmare.com/~rushing/
#
# Regardless, this shows how to use ctypes with the fftw library, and how cool numpy is.
#

import ctypes
import numpy
import time
import pickle

fftw3 = ctypes.cdll.LoadLibrary ("/usr/local/lib/libfftw3.3.dylib")

def read_ch (f):
    while 1:
        ch = f.read (1)
        if ch == '#':
            f.readline()
        else:
            return ch

def read_pgm_num (f):
    r = []
    found = False
    while 1:
        ch = read_ch (f)
        if not found:
            if ch in ' \r\n\t':
                pass
            elif ch in '0123456789':
                found = True
                r.append (ch)
            else:
                raise ValueError
        else:
            if ch in '0123456789':
                r.append (ch)
            else:
                break
    return int (''.join (r))

def read_pgm_header (f):
    h0 = f.read (1)
    h1 = f.read (1)
    m = h0+h1
    assert (m == 'P5')
    w = read_pgm_num (f)
    h = read_pgm_num (f)
    maxval = read_pgm_num (f)
    return w, h, maxval

def image_test (pgm_file):
    global h, w, scale
    f = open (pgm_file, 'rb')
    w, h, maxval = read_pgm_header (f)
    pixels = w * h
    a = numpy.fromfile (f, dtype=numpy.int8, count=pixels)
    a = a.astype (numpy.double)
    a2 = numpy.zeros ((w,h/2+1), dtype=numpy.complex)
    print 'planning to plan...'
    plan = fftw3.fftw_plan_dft_r2c_2d (w, h, a.ctypes.data, a2.ctypes.data, 64)
    # print 'printing the plan...'
    # fftw3.fftw_print_plan (plan)
    print 'executing the plan...', plan
    fftw3.fftw_execute (plan)
    print 'done!'
    # now go backward...
    print 'reversing...'
    a3 = numpy.zeros ((w,h), dtype=numpy.double)
    plan = fftw3.fftw_plan_dft_c2r_2d (w, h, a2.ctypes.data, a3.ctypes.data, 64)
    fftw3.fftw_execute (plan)
    a4 = a3 / (w*h)
    print 'done!'
    a4 = a4.astype (numpy.int8)
    f = open ('fft_inv.pgm', 'wb')
    f.write ('P5\n%d\n%d\n255\n' % (w, h))
    a4.tofile (f)
    f.close()
    
if __name__ == '__main__':
    import sys
    image_test (sys.argv[1])