#! /usr/bin/python

# This work is in the public domain.
#
# It was completed by Arun I in 2014.
#
# Though you are not legally obliged to do so,
# I would appreciate it if you credit me for
# this work by not removing this notice, and
# hopefully linking to my web page.

from numpy import *
from scipy import misc, sparse
from PIL import Image
import argparse, itertools
import os, math

def readimg(filename):
    img = Image.open(filename).convert('L') # Read image file using pillow 
    img = misc.fromimage(img)               # Convert to numpy ndarray 
    img = img / 255                         # Normalize image to range [0, 1] 
    return img

def sparsefft(mat, threshold):
    print('Suppressing components with amplitude less than %f of maximum' % threshold)
    matf = fft.fft2(mat)                                # Compute the 2D FFT of the matrix
    matf[abs(matf) < threshold*abs(matf).max()] = 0     # Suppress components with amplitude less than 'threshold' fraction of maximum
    spmatf = sparse.coo_matrix(matf)                    # Create sparse matrix in coordinates format
    print('Shortlisted %d components' % spmatf.nnz)
    return sparse.coo_matrix(matf)                      # Return result

def sortsparsefft(freq):
    permute = list(reversed(abs(freq.data).argsort()))  # Find permutation based on magnitude of frequency component
    return sparse.coo_matrix((freq.data[permute], (freq.row[permute], freq.col[permute])))  # Return permuted sparse FFT

def partition(imgf, nparts):
    power = real(imgf.data * conj(imgf.data))               # Get power signal 
    accumulated = map(log10, itertools.accumulate(power))   # Accumulate power signal in dB
    totalpower = log10(sum(power))                          # Calculate total power in signal

    parts = []
    startindex = 0
    endindex = 0
    currentpower = 0
    for i in range(nparts - 1):
        # Calculate partition's share of energy
        threshold = (totalpower - currentpower)/(nparts - i) + currentpower
        # Keep adding components to current partition till power becomes greater than the partition's share
        while currentpower < threshold:
            currentpower = next(accumulated)
            endindex += 1
        parts.append(sparse.coo_matrix((imgf.data[startindex:endindex], (imgf.row[startindex:endindex], imgf.col[startindex:endindex])), shape=imgf.shape))
        startindex = endindex
    # For the last partition, dump all remaining components
    parts.append(sparse.coo_matrix((imgf.data[startindex:], (imgf.row[startindex:], imgf.col[startindex:])), shape=imgf.shape))
    return parts

def reconstruct(parts, filenameprefix, fileextension):
    imgf = complex(0) * zeros(parts[0].shape)   # Initialize frequency domain image with zeros 
    fileformat = '%s_%%0%dd%s' % (filenameprefix, math.ceil(math.log10(len(parts))), fileextension)
    # For each part in the partitioned frequency spectrum
    for i in range(len(parts)):
        imgf = imgf + parts[i].todense()    # Add components in current partition to frequency domain image
        img = real(fft.ifft2(imgf))         # Compute the inverse FFT
        img = img * 255                     # Denormalize image back to [0, 255] scale
        misc.imsave(fileformat % i, img)    # Save generated image to file
        if not args.disableprogressbar:
            printprogress((i + 1)/len(parts), 'Generating frames')

def printprogress(progress, message):
    barlength = int(os.get_terminal_size().columns / 2)     # Set length of progress bar to half the number of columns available
    print('\r%s: [' % message, end='')                      # Print message
    # Print progress bar
    for i in range(barlength):
        print('#' if i <= round(progress * barlength) else '-', end='')
    print('] %d%%' % round(progress * 100), end='')         # Close progress bar, and print percentage progress

parser = argparse.ArgumentParser()
parser.add_argument('inputfile', help='Path to input image file')
parser.add_argument('nframes', help='Number of output frames to generate', type=int)
parser.add_argument('-p', '--disableprogressbar', help='Disable progress bar display', action='store_true')
parser.add_argument('--threshold', help='Frequency components whose amplitude is below ''threshold'' fraction of the amplitude of the most dominant component will be suppressed (default=0.00001)', default=0.00001)
args = parser.parse_args()

img = readimg(args.inputfile)
imgf = sparsefft(img, args.threshold)
imgf = sortsparsefft(imgf)

parts = partition(imgf, args.nframes)
inputfilepath = os.path.splitext(args.inputfile)
reconstruct(parts, inputfilepath[0], inputfilepath[1])