#encoding:latin-1
from numpy import *
from pylab import *

from py_nnma import *

from glob import glob
from pgm_utils import *
import os

###############################################################################
## 
##  reads several frames of movie below 'movie_dir' and tries to separate 
##  single objects with NNMA.
## 
###############################################################################


## read all frames ############################################################
movie_dir="movie/"
if not os.path.exists(movie_dir):
    print "please run make_movie.py before you use this program"
    exit()

collected_images = []
for path in glob(os.path.join(movie_dir,"frame*.pgm")):
    data = read_pgm(path)
    collected_images.append(data.flatten())

## build matrix as input for FNAMI  ###########################################
Y = array(collected_images, dtype=double).T # one frame per col of A #


# take averaged image as starting point for B  ################################
# Astart = Y.mean(axis=1)[newaxis].repeat(4, axis=0).T

# start with random matrix
Astart = None

# run FNMAI  ##################################################################
Y += 1e-3 # avoid zeros

print
print "RUN FNMAI"

A,X, res, count, convergence =FNMAI(Y, k=4, A=Astart, tau=2, alpha=0.2, 
                                    maxcount=1000, verbose=10)

# plot seperated objects ######################################################
figure()
for i, row in enumerate(A.T):
    subplot(221+i)
    title("FNMAI %d" % (i+1))
    imshow(row.reshape(70,70), cmap=cm.gray)
    xticks([])
    yticks([])


# run FNMAI_SPARSE   ##################################################################
print
print "RUN GDCLS_L1"
A, X, res, count, convergence =GDCLS_L1(Y, k=4, A=Astart, regul=5e-1, tau=2, alpha=.2,
                                 maxcount=1000, verbose=10)

# plot seperated objects ######################################################
figure()
title("sparse")
for i, row in enumerate(A.T):
    subplot(221+i)
    title("GDCLS_L1 %d" % (i+1))
    imshow(row.reshape(70,70), cmap=cm.gray)
    xticks([])
    yticks([])
show()