#!/usr/bin/env python
# -*- coding: utf-8 -*-

import sys
import argparse
import fileinput
import numpy

class Image:
	def __init__(self, matrix=[[]], width=0, height=0, depth=0):
		self.matrix = matrix
		self.width = width
		self.height = height
		self.depth = depth

	def set_width_and_height(self, width, height):
		self.width = width
		self.height = height
		self.matrix = [[0 for j in xrange(height)] for i in xrange(width)]


def write_matrices_to_file(matrixU, matrixS, matrixVt, kmin, kmax, file_handle, width, height, depth, rescale=False):
	"""
	Write a decomposed matrix to file uncompressed as it would show compressed

	Keyword Arguments:
	matrixU -- the U portion of the SVD
	matrixS -- the S (sigma) portion of the SVD
	matrixVt -- the V transpose portion of the SVD
	kmin -- the minimum k value to use for compresion (ignored if kmax = 0)
	kmax -- the maximum kvalue to use for compresion (find optimal if zero)
	filename -- the file to write to (stdout if blank)
	width -- the image width
	height -- the image height
	depth -- the maximum grey scale value (normally 255)
	rescale -- True to shift resulting image into 0 < n < depth bounds
	"""
	matrixScopy = matrixS.copy()
	# when kmax is not 0 use the provided kmax
	if kmax > 0:
		i = 0
		for t in numpy.nditer(matrixScopy, op_flags=["readwrite"]):
			if i < kmin or i >= kmax:
				t[...] = 0
			i += 1
	# when kmax is 0 then drop eigen values less than 1.0E-14
	else:
		for t in numpy.nditer(matrixScopy, op_flags=["readwrite"]):
			if round(t, 14) <= 0:
				t[...] = 0
	
	# recompose the trimmed SVD matrices back into matrix A
	A = numpy.dot(numpy.dot(matrixU, numpy.diag(matrixScopy)), matrixVt)

	# attempt the handle out of range values (TODO: pull out to own function)
	if rescale:
		curMin = 0
		curMax = 0
		# find min and max values
		for n in numpy.nditer(A):
			if int(round(n)) < curMin:
				curMin = int(round(n))
			if int(round(n)) > curMax:
				curMax = int(round(n))
		# shift values up
		if curMax < depth and curMin < 0:
			shiftVal = depth - curMax
			for t in numpy.nditer(A, op_flags=["readwrite"]):
				t[...] = t + shiftVal
				if t > depth:
					t[...] = depth
				elif t < 0:
					t[...] = 0
		# shift values down
		elif curMax > depth and curMin > 0:
			shiftVal = curMin
			for t in numpy.nditer(A, op_flags=["readwrite"]):
				t[...] = t - shiftVal
				if t > depth:
					t[...] = depth
				elif t < 0:
					t[...] = 0
		# no chance to shift, just chop (TODO: perform some sort of scaling)
		else:
			for t in numpy.nditer(A, op_flags=["readwrite"]):
				if t > depth:
					t[...] = depth
				elif t < 0:
					t[...] = 0

	file_handle.write("P2\n")
	file_handle.write("# Generated by Stoll \n")
	file_handle.write(str(width))
	file_handle.write(" ")
	file_handle.write(str(height))
	file_handle.write("\n")
	file_handle.write(str(depth))
	file_handle.write("\n")
	for n in numpy.nditer(A):
		file_handle.write(str(int(round(n))))
		file_handle.write(" ")
	file_handle.write("\n")


def read_matrix_from_file(file_handle):
	"""
	Read an ASCII PGM file and create an Image object from it
	"""
	row = 0
	col = 0
	image = Image()
	for line in file_handle:
		if line[0] == '#':
			pass
		elif line[0] == 'P' and line[1] == '2':
			pass
		elif image.width == 0 and image.height == 0:
			x = 0
			y = 0
			x, y = [int(n) for n in line.split()]
			image.set_width_and_height(x, y)
		elif image.depth == 0:
			image.depth = int(line)
		else:
			for value in line.split():
				if col >= len(image.matrix[row]):
					row += 1
					col = 0
				image.matrix[row][col] = value
				col += 1
	return image
	

def process_svd(source_file, destination_file, kmin, kmax):
	"""
	Read from file provided on the command line or from stdin
	then save uncompressed representations of the SVD compressed version
	"""
	image = read_matrix_from_file(source_file)

	M = numpy.asmatrix(image.matrix)
	U, s, Vt = numpy.linalg.svd(M, full_matrices=True)
	
	write_matrices_to_file(U, s, Vt, kmin, kmax, destination_file, image.width, image.height, image.depth)



if __name__ == "__main__":
	parser = argparse.ArgumentParser()
	parser.add_argument("infile", nargs='?', help="The source ASCII PGM file", type=argparse.FileType('r'), default=sys.stdin)
	parser.add_argument("outfile", nargs='?', help="The destination ASCII PGM file", type=argparse.FileType('w'), default=sys.stdout)
	parser.add_argument("-j", "--kmin", help="The number of high k values to exlude", type=int, default=0)
	parser.add_argument("-k", "--kmax", help="The number k values to use", type=int, default=0)
	parser.add_argument("-s", "--scale", help="Fit resulting image depth into '0 < n < depth' bounds", action="store_true")
	args = parser.parse_args()
	process_svd(args.infile, args.outfile, args.kmin, args.kmax)

@@ -3643,16 +3643,25 @@
 in, kmax
+, rescale
 ):%0A%09%22%22%22%0A
@@ -4698,16 +4698,23 @@
 _args()%0A
+%09try:%0A%09
 %09process
@@ -4765,10 +4765,59 @@
 rgs.kmax
+, args.scale)%0A%09except KeyboardInterrupt:%0A%09%09exit(0
 )%0A

Include scale and trap ctrl-c