#!/usr/local/bin/python

import argparse
import os
import dpy
import shutil

# TODO: might not work in some cases
def crop_resize_convert_to_match_target_image(image, target_image):
    from scipy.misc import imresize
    
    target_width  = target_image.shape[1]
    target_height = target_image.shape[0]

    image_width = image.shape[1]
    image_height = image.shape[0]
    image_aspect = float(image_width) / float(image_height)

    resized_width = target_width
    resized_height = int(target_width / image_aspect)
    if resized_height < target_height:
      scale = float(target_height) / float(resized_height)
      resized_width = int(scale * resized_width)
      resized_height = target_height
    new_image = imresize(image, (resized_height, resized_width))
    
    left = new_image.shape[1] / 2 - target_width / 2
    top = new_image.shape[0] / 2 - target_height / 2

    #print image.shape
    #print top, left

    new_image = new_image[top:top+target_height, left:left+target_width]
    new_image = new_image.astype(target_image.dtype)

    return new_image

def pca_image_dir(args):
   import sys
   import cv2
   import numpy as np
   from scipy.misc import imresize

   in_path = args.images_directory
   stride = args.stride
   target_mean_path = args.target_mean
   target_mean = None
   if target_mean_path != None:
       target_mean = np.load(target_mean_path)
       print "TARGET_MEAN shape: " + str(target_mean.shape) + ", type: " + str(target_mean.dtype)

   files = os.listdir(in_path)
   out_dir = os.path.abspath(in_path) + "_every-" + str(stride) + "_scale-" + str(args.scale) + "_PCA"

   f = 0
   matrix_test = None
   imgraw = None
   print " "
   for image in files:
       f = f + 1
       if f % stride != 0:
	   continue
       imgraw = cv2.imread(os.path.join(in_path, image), -1)
       imgraw = imresize(imgraw, (int(args.scale * imgraw.shape[0]), int(args.scale * imgraw.shape[1])))
       if target_mean != None:
           imgraw = crop_resize_convert_to_match_target_image(imgraw, target_mean)
       imgvector = imgraw.reshape(imgraw.size)
       #print matrix_test
       try:
	   matrix_test = np.vstack((matrix_test, imgvector))
       except:
	   matrix_test = imgvector
       sys.stderr.write( str(f) + " of " + str(len(files)) + "\r" )
   print ""

   # PCA
   print "Running PCA ..."
   #print matrix_test
   mean, eigenvectors = cv2.PCACompute(matrix_test, np.mean(matrix_test, axis=0).reshape(1,-1))

   mean = mean.reshape(imgraw.shape)
   #print "mean: " + str(mean)
   print " "

   os.mkdir(out_dir)

   mean_filepath = os.path.join(out_dir, "mean.npy")
   np.save(mean_filepath, mean)

   meanimg = mean
   #meanimg *= 255.0/mean.max()
   cv2.imwrite(os.path.join(out_dir, "mean.png"), meanimg)

   i = 0
   for eigenvector in eigenvectors:
       filename = "%05d" % i
       path_no_extension = os.path.join(out_dir, filename)
       npy_path = path_no_extension + ".npy"
       cvxml_path = path_no_extension + ".xml"
       eigenvector = eigenvector.reshape(imgraw.shape)
       np.save(npy_path, eigenvector)
       #cv.Save(cvxml_path, eigenvector)
   #    print "Saved " + path
   #    print eigenvector
   #    print "--------------------------"
   #    print "min: " + str(evimg.min()) + ", max: " + str(evimg.max())
   #    evimg += evimg.min()
       evimg = eigenvector
       evimg *= 255.0/eigenvector.max()
       cv2.imwrite(os.path.join(out_dir, filename + ".png"), evimg)
       i = i + 1

   print "Saved " + str(len(eigenvectors)) + " eigenvectors to " + out_dir

def _images_as_basis(args):
  import numpy as np
  from scipy.misc import imread, imsave, imresize, imshow, fromimage, toimage
  from PIL import Image
  #from skimage.io import imread
  
  target_width  = None
  target_height = None

  mean = None
  if args.mean != None:
    mean = np.load(args.mean)
    target_width = int(args.scale * mean.shape[1])
    target_height = int(args.scale * mean.shape[0])
  else:
    import enzyme
    video_meta = enzyme.parse(args.target_video)
    target_width = int(args.scale * video_meta.video[0].width)
    target_height = int(args.scale * video_meta.video[0].height)
    if args.grayscale:
      mean = np.ones((target_height, target_width), dtype=np.dtype('float32'))
      mean = 128.0 * mean
    else:
      mean = np.ones((target_height, target_width, 3), dtype=np.dtype('float32'))
      mean = 128.0 * mean
  print "Mean: "
  print mean
    
  array_out_directory = args.input_directory + "_as_basis"
  
  if args.grayscale:
     print "Will convert images to grayscale"
     array_out_directory = array_out_directory + "_gray"

     if len(mean.shape) > 2:
       print " "
       print args.mean + "\nhas more than 1 channel. Need a grayscale mean.\n"
       exit()

  print str(target_width) + " x " + str(target_height) + " each..."

  dpy.ensure_dir(array_out_directory)

  #shutil.copy(args.mean, os.path.join(array_out_directory, "mean.npy"))
  mean_out_path = os.path.join(array_out_directory, "mean.npy")
  np.save(mean_out_path, mean)
  imsave(mean_out_path + ".jpg", mean.astype(np.dtype('uint8')))

  index = 0
  image_names = sorted(os.listdir(args.input_directory))
  for image_name in image_names:
    image_path = os.path.join(args.input_directory, image_name)
    #image = imread(image_path)
    image = Image.open(image_path)
    
    if args.grayscale:
      image = image.convert("L")
    
    image = fromimage(image)
    
    image = crop_resize_convert_to_match_target_image(image, mean)    

    image = image / np.linalg.norm(image)

    out_name = "%04i.npy" % (index)
    out_path = os.path.join(array_out_directory, out_name)
    np.save(out_path, image)
    imsave(out_path + ".jpg", image)
    print image_name + " ----constrain( "+str(target_width)+" x "+str(target_height)+" )---> " + out_path

    index = index + 1

parser = argparse.ArgumentParser()
subparsers = parser.add_subparsers(dest='command')

# use images as PCA basis
images_as_basis = subparsers.add_parser('images_as_basis')
images_as_basis.add_argument('-g', '--grayscale', dest='grayscale', action='store_const', const=True, default=False, 
                             help='convert input images to grayscale')
size_specifiers = images_as_basis.add_mutually_exclusive_group(required=True)
size_specifiers.add_argument('-m', '--mean', type=str,
                             help='set size and mean from this NumPy array')
size_specifiers.add_argument('-t', '--target-video', type=str,
                             help='set size from this video, generate gray mean')
images_as_basis.add_argument('scale', type=float)
images_as_basis.add_argument('input_directory', type=str)
images_as_basis.set_defaults(func=_images_as_basis)

# do pca of images
pca_command = subparsers.add_parser('pca')
pca_command.set_defaults(func=pca_image_dir)
pca_command.add_argument('images_directory', type=str,
                         help='Directory containing input images')
pca_command.add_argument('--target-mean', type=str,
                         help='crop and resize images to fit (NumPy array) TARGET_MEAN\'s shape')
pca_command.add_argument('--stride', type=int, default=1,
                         help='use 1 out of every STRIDE images')
pca_command.add_argument('--scale', type=float, default=1.0,
                         help='Scale images by SCALE')

args = parser.parse_args()
args.func(args)
