import sys, os        # some basic Python services
import Image          # loading, manipulating, and displaying images
import Numeric        # fast arrays and matrices
import LinearAlgebra  # LAPACK services (in particular, eigenvectors)
import cPickle        # object serialization (write arbitrary objects to files)
from glob import glob # wildcard matching
from math import sqrt

def norm(v): return sqrt(Numeric.dot(v, v))
def normalize(v): return v / norm(v)

def mesg(s):
    sys.stdout.write(s)
    sys.stdout.flush()

class FaceSpace:
    def __init__(self, dir=''):
        if dir != '':
            self.load_eigenfaces(dir)
        
    # takes a filename and returns a vector with the data from that image
    def image_to_vector(self, filename):
        try:
            im = Image.open(filename)
        except IOError:
            print 'couldn\'t load ' + filename + ', aborting'
            sys.exit(1)
        self.im_size = im.size
        a = Numeric.array(im.getdata(), Numeric.Float)
        # a.shape = im.size   # this would make it a matrix, not a vector
        return a - 128

    # takes a vector and a size tuple and saves it to filename (type inferred)
    def vector_to_image(self, v, filename):
        v.shape = (-1,)
        a, b = min(v), max(v)
        span = max(abs(b), abs(a))
        im = Image.new('L', self.im_size)
        im.putdata((v * 127. / span) + 128)
        im.save(filename)

    # saves a vector to a temporary file and displays it asynchronously
    def vector_to_screen(self, v, name):
        self.vector_to_image(v, name + '.gif')
        os.system('display ' + name + '.gif &')

    # takes a directory (which it empties), a list of files to analyze, and a
    # number of eigenfaces to construct (M_prime)
    def construct_eigenfaces(self, dir, lof, M_prime):
        # empty the directory (create it if necessary)
        try:
            for f in os.listdir(dir):
                os.unlink(dir + '/' + f)
        except os.error:    # doh! directory didn't exist
            os.mkdir(dir)
        # load all the images (gamma)
        mesg('Loading images...')
        Gamma = map(self.image_to_vector, lof)
        print 'done.'
        self.M = len(Gamma)
        # compute psi and phi
        self.Psi = reduce(lambda a, b: a+b, Gamma) / len(Gamma)
        self.Phi = []
        for g in Gamma: self.Phi.append(g - self.Psi)
        # throw out gamma (to save memory)
        Gamma = None
        # build L, the matrix of image dot products
        L = Numeric.array([0.0] * self.M * self.M)
        L.shape = (self.M, self.M)
        for i in range(self.M):
            for j in range(self.M):
                L[i][j] = Numeric.dot(self.Phi[i], self.Phi[j])
        # save psi into the directory
        self.vector_to_image(self.Psi, dir + '/psi.gif')
        # compute the eigenvectors
        evalues, evectors = LinearAlgebra.eigenvectors(L)
        # sort them by eigenvalue and keep the top M_prime
        evs = map(None, evalues, evectors)
        evs.sort()
        evs.reverse()
        evs = evs[0:M_prime]
        # write those into the directory
        v = map(lambda x: x[1], evs)
        self.u = []
        for k in range(M_prime):
            mesg(' ' + str(k+1))
            self.u.append(Numeric.matrixmultiply(v[k], self.Phi))
            self.vector_to_image(self.u[-1], '%s/eig%03d.gif' % (dir, k))
        print
        self.M_prime = M_prime

    # load a set of eigenfaces from dir, overwriting any previous set
    def load_eigenfaces(self, dir):
        mesg('Loading eigenfaces... ')
        eigs = glob(dir + '/eig*')
        eigs.sort()
        self.M_prime = len(eigs)
        self.u = map(self.image_to_vector, eigs)
        self.Psi = self.image_to_vector(dir + '/psi.gif')
        print 'done.'

    # takes an image vector, projects it into face space, returns that vector
    def project_to_face_space(self, im_face):
        # have to futz with the shape so I can use matrixmultiply
        for v in self.u: v.shape = (v.shape[0], 1)
        fs_face = Numeric.matrixmultiply(im_face - self.Psi, self.u)
        # put it back like I found it
        for v in self.u: v.shape = (-1, )
        fs_face.shape = (-1, )
        return normalize(fs_face)

    # takes a face_space vector and multiplies it by the eigenfaces
    # if n_eigs is nonzero, use only that many eigenfaces
    def lift_to_image_space(self, fs_face, n_eigs=0):
        if n_eigs == 0: n_eigs = self.M_prime
        scaled = []
        for i in range(n_eigs):
            scaled.append(self.u[i] * fs_face[i])
        im_face = reduce(lambda a, b: a+b, scaled)
        return im_face + self.Psi

    # take a file name for a learned face and return the vector it contains
    def load_fs_face(self, filename):
        file = open(filename, 'r')
        fs_face = cPickle.load(file)
        file.close()
        fs_face.shape = (-1, )  # flatten the vector
        return normalize(fs_face)

    # take a face space vector and write it to the file named by filename
    def write_fs_face(self, fs_face, filename):
        file = open(filename, 'w')
        fs_face = cPickle.dump(fs_face, file)
        file.close()