import numpy as np
import sys
import os
import subprocess as sp
from PIL import Image, ImageDraw
from scipy.misc import imsave
from scipy.signal import medfilt2d
import ast
import json

from fit_ellipse import fit_ellipse, FitEllipse
from itracker_utils import generate_rays, initial_pupil_point, initial_cr_point, sobel_grad
import logging

import matplotlib.pyplot as plt

# import cv2

color_list = ['b','g','r','c','m','y','k']

class iTracker (object):
    def __init__(self, output_folder, 
                 im_shape, num_frames, 
                 input_stream, 
                 threshold_factor=1.3, auto=True,
                 cutoff_pixels=10,
                 bbox_pupil=None,
                 bbox_cr=None):

        self.im_shape = im_shape
        self.num_frames = num_frames
        self.movie_shape = (num_frames, im_shape[0], im_shape[1])
        self.input_stream = input_stream

        self.threshold_factor = threshold_factor
        self.auto = auto
        self.folder = output_folder
        self.cutoff_pixels = cutoff_pixels
        self.bbox_pupil = bbox_pupil
        self.bbox_cr = bbox_cr

        self._mean_frame = None

        if not os.path.exists(self.folder):
            os.mkdir(self.folder)

        self.run_params_file = os.path.join(self.folder, 'run_params.json')
        self.run_params = { 'threshold_factor': threshold_factor,
                            'auto':  auto,
                            'cutoff_pixels': cutoff_pixels,
                            'movie_shape':  self.movie_shape,
                            'im_shape':  im_shape,
                            'bbox_pupil':  bbox_pupil,
                            'bbox_cr':  bbox_cr }
        with open(self.run_params_file, 'w') as f:
            f.write(json.dumps(self.run_params))

        self.movie_path_storage_file = os.path.join(self.folder, 'movie_path.txt')
        if os.path.exists(self.movie_path_storage_file):
            with open(self.movie_path_storage_file, 'r') as f:
                self.movie_path = f.read()
        else:
            self.movie_path = None

        self.input_image_folder = os.path.join(self.folder, 'input_images')
        if not os.path.exists(self.input_image_folder):
            os.mkdir(self.input_image_folder)

        self.results_folder = os.path.join(self.folder,'results')

        if not os.path.exists(self.results_folder):
            os.mkdir(self.results_folder)

        # self.rays_folder = os.path.join(self.results_folder,'rays')
        # if not os.path.exists(self.rays_folder):
        #     os.mkdir(self.rays_folder)

        self.qc_folder = os.path.join(self.results_folder,'qc')
        if not os.path.exists(self.qc_folder):
            os.mkdir(self.qc_folder)

        self.frames_folder = os.path.join(self.results_folder,'output_frames')
        if not os.path.exists(self.frames_folder):
            os.mkdir(self.frames_folder)

        self.pupil_file = os.path.join(self.results_folder, 'pupil_params.npy')
        self.cr_file = os.path.join(self.results_folder, 'cr_params.npy')
        self.mean_frame_file = os.path.join(self.results_folder, 'mean_frame.npy')
        self.annotated_movie_file = os.path.join(self.results_folder, 'annotated_movie.mp4')

        # add variables to determine whether to provide diagnostic, QC and other output
        # method to regnerate image frames, with or without results?


    # fix this so it sets an absolute path
    def set_movie(self, file_path):
        logging.debug("Setting movie_path to: %s", file_path)
        self.movie_path = file_path
        with open(self.movie_path_storage_file, 'w') as f:
            f.write(self.movie_path)

    def set_bbox_pupil(self, bbox):
        self.bbox_pupil = bbox
        self.run_params['bbox_pupil']=self.bbox_pupil
        with open(self.run_params_file, 'w') as f:
            f.write(json.dumps(self.run_params))

    def set_bbox_cr(self, bbox):
        self.bbox_cr = bbox
        self.run_params['bbox_cr']=self.bbox_cr
        with open(self.run_params_file, 'w') as f:
            f.write(json.dumps(self.run_params))

    def create_input_images(self, image_type='png'):
        self.input_stream.create_images(self.input_image_folder, image_type)

    @property
    def mean_frame(self):
        if self._mean_frame is None:
            self._mean_frame = self.compute_mean_frame()
        return self._mean_frame

    @mean_frame.setter
    def mean_frame(self, mean_frame):
        self._mean_frame = mean_frame

    def estimate_bbox_from_mean_frame(self, margin=75, image_type='png'):
        try:
            import keras
        except ImportError:
            logging.debug("keras failed to import.  Returning None for bbox_pupil and bbox_cr")
            return None, None

        from keras.applications import InceptionV3

        logging.debug("Estimating bbox parameters from 'mean_frame'")
        # compute the representation for mean_frame
        model = InceptionV3(include_top=False, weights='imagenet')
        # print self.mean_frame.dtype, self.mean_frame.shape
        mp_temp = self.mean_frame.astype(np.float32)
        mp_temp -= 128
        mp_temp /= 128
        rep = model.predict(mp_temp.reshape((1,)+mp_temp.shape))  # shape (1,13,18,2048)
        rep[rep<0]=0
        rep = np.mean(rep, axis=(0,1,2))  # shape (2048,)

        # load regression weights
        module_folder = os.path.dirname(os.path.abspath(__file__))
        W_pupil = np.load(os.path.join(module_folder,'resources','pupil_weights.npy'))  # shape (2048, 5)
        W_cr = np.load(os.path.join(module_folder,'resources','cr_weights.npy'))    # shape (2048, 5)

        estimated_pupil_point = np.dot(rep, W_pupil)  # shape (5,)
        estimated_cr_point = np.dot(rep, W_cr)   # shape (5,)

        x_pupil, y_pupil = estimated_pupil_point[:2]*np.array([640,480])
        x_pupil = int(x_pupil)
        y_pupil = int(y_pupil)
        logging.debug("estimated pupil point is ({0},{1})".format(x_pupil,y_pupil))
        print "estimated pupil point is ({0},{1})".format(x_pupil,y_pupil)
        # bbox is xmin, xmax, ymin, ymax
        # x, y = 320, 240
        bbox_pupil = [x_pupil-margin, x_pupil+margin, y_pupil-margin, y_pupil+margin]

        x_cr, y_cr = estimated_cr_point[:2]
        x_cr = int(x_cr)
        y_cr = int(y_cr)
        logging.debug("estimated cr point is ({0},{1})".format(x_cr,y_cr))
        print "estimated cr point is ({0},{1})".format(x_cr,y_cr)

        # bbox is xmin, xmax, ymin, ymax
        bbox_cr = [x_cr-margin, x_cr+margin, y_cr-margin, y_cr+margin]
        # bbox_cr = None

        # plot bbox on mean_frame for QC check
        mean_frame_annotated = np.dstack([self.mean_frame, self.mean_frame, self.mean_frame])
        mean_frame_annotated = self.annotate_frame_with_bbox(mean_frame_annotated,pupil_bbox=bbox_pupil,cr_bbox=bbox_cr)
        mean_frame_annotated = self.annotate_frame_with_point(mean_frame_annotated,pupil=(x_pupil, y_pupil),cr=(x_cr, y_cr))

        dpi = 100.0
        fig, ax = plt.subplots(figsize=(mean_frame.shape[1]/dpi, mean_frame.shape[0]/dpi))
        fig.subplots_adjust(left=0,right=1,bottom=0,top=1)

        ax.imshow(mean_frame_annotated, aspect='normal')
        ax.axis('off')
        fig.savefig(os.path.join(self.qc_folder, 'mean_frame_annotated.'+image_type), dpi=dpi)

        self.bbox_pupil = bbox_pupil
        self.bbox_cr = bbox_cr

        return bbox_pupil, bbox_cr

    def compute_mean_frame(self, image_file_type='png'):
        logging.debug("computing mean frame")

        mean_frame = np.zeros(self.im_shape)

        frames_read = 0
        for input_frame in self.input_stream:
            mean_frame += input_frame
            frames_read += 1

        mean_frame /= frames_read
        mean_frame = mean_frame.astype(np.uint8)

        np.save(self.mean_frame_file, mean_frame)

        dpi = 100.0
        fig, ax = plt.subplots(figsize=(mean_frame.shape[1]/dpi, mean_frame.shape[0]/dpi))
        fig.subplots_adjust(left=0,right=1,bottom=0,top=1)
        ax.imshow(mean_frame, aspect='normal', cmap='gray')
        ax.axis('off')
        fig.savefig(os.path.join(self.qc_folder, 'mean_frame.' + image_file_type), dpi=dpi)
        plt.close()

        if self.bbox_cr and self.bbox_pupil:
            mean_frame_annotated = np.dstack([mean_frame,mean_frame,mean_frame])
            mean_frame_annotated = self.annotate_frame_with_bbox(mean_frame_annotated,pupil_bbox=self.bbox_pupil,cr_bbox=self.bbox_cr)        

            fig, ax = plt.subplots(figsize=(mean_frame.shape[1]/dpi, mean_frame.shape[0]/dpi))
            fig.subplots_adjust(left=0,right=1,bottom=0,top=1)
            ax.imshow(mean_frame_annotated, aspect='normal')
            ax.axis('off')
            fig.savefig(os.path.join(self.qc_folder, 'mean_frame_bbox.' + image_file_type), dpi=dpi)
            plt.close()

        return mean_frame



    def detect_eye_closed(self):

        try:
            import keras
        except ImportError:
            logging.debug("keras failed to import.  Can't detect eye closure")
            return None

        from keras.applications import InceptionV3

        logging.debug("Detecting eye closed frames")
        # compute the representation for mean_frame
        model = InceptionV3(include_top=False, weights='imagenet')
        # print self.mean_frame.dtype, self.mean_frame.shape

        # get pre-trained svm
        from sklearn.externals import joblib
        module_folder = os.path.dirname(os.path.abspath(__file__))
        svm = joblib.load(os.path.join(module_folder, 'resources','svm_trained.pkl'))

        def compute_rep(frame):
            mp_temp = frame.astype(np.float32)
            mp_temp -= 128
            mp_temp /= 128
            rep = model.predict(mp_temp.reshape((1,)+mp_temp.shape))  # shape (1,13,18,2048)
            rep[rep<0]=0
            rep = np.mean(rep, axis=(0,1,2))  # shape (2048,)

            return rep

        is_closed = np.zeros(self.num_frames)

        for input_frame in self.input_stream:
            rep = compute_rep(input_frame)
            is_closed[i] = svm.predict(rep.reshape(-1,len(rep)))[0]

        self.is_closed = is_closed
        save_path = os.path.join(self.results_folder, 'is_closed.npy')

        logging.debug("Saving is_closed to:")
        logging.debug("\t%s", save_path)
        #
        np.save(save_path, self.is_closed)

        return is_closed

    def process_movie(self, movie_output_stream=None,
                      output_frames=False,
                      output_annotation_frames=False,
                      image_file_type = 'jpg' ):

        # these aren't really used yet.
        # self.pupil_loc = (0,0)
        # self.cr_loc = (0,0)

        self.pupil_params = np.zeros([self.num_frames, 5])
        self.cr_params = np.zeros([self.num_frames, 5])

        if movie_output_stream:
            movie_output_stream.open(self.annotated_movie_file)
        else:
            movie_output_stream = None

        if output_frames:
            frame_output_stream = ImageOutputStream()
            frame_output_stream.open(os.path.join(self.input_image_folder, 'input_frame-%06d.'+image_file_type))
        else:
            frame_output_stream = None

        if output_annotation_frames:
            annotation_frame_output_stream = ImageOutputStream()
            annotation_frame_output_stream.open(os.path.join(self.frames_folder, 'output_frame-%06d.'+image_file_type))
        else:
            annotation_frame_output_stream = None

        for i, input_frame in enumerate(self.input_stream):
            # get pupil and corneal reflection parameters, this line is the actual eye tracking algorithm
            pupil, cr = self.process_image(input_frame, bbox_pupil=self.bbox_pupil, bbox_cr=self.bbox_cr)

            pupil_params = (pupil[0][0],pupil[0][1],pupil[1],pupil[2][0],pupil[2][1])
            cr_params = (cr[0][0],cr[0][1],cr[1],cr[2][0],cr[2][1])

            if frame_output_stream:
                frame_output_stream.write(input_frame)

            if movie_output_stream or annotation_frame_output_stream:
                annotated_frame = self.annotate_frame(np.dstack([input_frame,input_frame,input_frame]), 
                                                      pupil_params, 
                                                      cr_params)

                if movie_output_stream:
                    movie_output_stream.write( annotated_frame )

                if annotation_frame_output_stream:
                    annotation_frame_output_stream.write( annotated_frame )

            # save results in arrays
            self.pupil_params[i] = (pupil[0][0],pupil[0][1],pupil[1],pupil[2][0],pupil[2][1])
            self.cr_params[i] = (cr[0][0],cr[0][1],cr[1],cr[2][0],cr[2][1])

            if i % 100 == 0:
                logging.debug("tracked frame %d", i)

        logging.debug("Saving pupil and cr parameters to:")
        logging.debug("\t%s", self.pupil_file)
        logging.debug("\t%s", self.cr_file)
        #
        np.save(self.pupil_file, self.pupil_params)
        np.save(self.cr_file, self.cr_params)

        if movie_output_stream:
            movie_output_stream.close()

        if frame_output_stream:
            frame_output_stream.close()

        if annotation_frame_output_stream:
            annotation_frame_output_stream.close()

        # return mean_frame

    def clear_input_images(self):
        logging.debug("Deleting input image folder")
        shutil.rmtree(os.path.join(self.folder, 'input_images'))

    def process_image(self, im, bbox_pupil=None, bbox_cr=None):
        # let's try median filtering the image first
        im = medfilt2d(im, kernel_size=3)

        # find pupil and corneal reflection if auto==True
        if self.auto:
            self.pupil_loc = initial_pupil_point(im, bbox=bbox_pupil)
            self.cr_loc = initial_cr_point(im, bbox=bbox_cr)


        # find rays projecting from seed point
        pupil_rays, pupil_ray_values = generate_rays(im,self.pupil_loc)

        # save values for analysis
        self.pupil_rays = pupil_rays
        self.pupil_ray_values = pupil_ray_values

        # code for finding pupil ellipse, start with candidate points from rays
        pupil_candidate_points = self.get_candidate_points(self.pupil_rays,self.pupil_ray_values,self.threshold_factor,above_threshold=True)

        # fit pupil ellipse with all candidate points
        #pupil_params = fit_ellipse(pupil_candidate_points)

        # fit pupil ellipse with ransac algorithm
        fe=FitEllipse(10,10,0.0001,4)
        result = fe.ransac_fit(pupil_candidate_points)

        # if np.any(np.isnan(result)):    #should use np.any(np.isnan(result))
        #     pupil_params = result  #fe.ransac_fit(pupil_candidate_points)
        # else:
        #     pupil_params = ((np.nan,np.nan),np.nan,(np.nan,np.nan))  #  np.nan*np.ones(5)


        if result!=None:    #should use np.any(np.isnan(result))
            pupil_params = result  #fe.ransac_fit(pupil_candidate_points)
        else:
            logging.debug("No good fit found")
            pupil_params = ((np.nan,np.nan),np.nan,(np.nan,np.nan))  #  np.nan*np.ones(5)


        # code for finding corneal reflection, start with finding rays from center of cr
        cr_rays, cr_ray_values = generate_rays(im,self.cr_loc)

        self.cr_rays = cr_rays
        self.cr_ray_values = cr_ray_values

        cr_candidate_points = self.get_candidate_points(self.cr_rays,self.cr_ray_values,0.75,above_threshold=False)

        try:
            #cr_params = fit_ellipse(cr_candidate_points)
            fe=FitEllipse(10,10,0.0001,4)
            result = fe.ransac_fit(cr_candidate_points)

            if result!=None:
                cr_params = result #fe.ransac_fit(cr_candidate_points)
            else:
                logging.debug("No good fit found")
                cr_params = ((np.nan,np.nan),np.nan,(np.nan,np.nan))
        except Exception as e:
            logging.error("Error during fit: %s", e.message)
            cr_params = ((np.nan,np.nan),np.nan,(np.nan,np.nan))

        # update instance variables
        if not np.isnan(pupil_params[0][0]):   #should use np.any(np.isnan(result))
            self.pupil_loc = (int(pupil_params[0][1]),int(pupil_params[0][0]))

        self.pupil_candidate_points = pupil_candidate_points
        self.cr_candidate_points = cr_candidate_points

        return pupil_params, cr_params

    def get_candidate_points(self,rays,ray_values,threshold_f,above_threshold=True):

        candidate_points = []
        # find candidate points for ellipse from threshold crossing of the image over the rays
        for i, ray in enumerate(rays):

            sample_ray = ray_values[i][:self.cutoff_pixels]
            threshold = threshold_f*np.mean(sample_ray)

            for t,g in enumerate(ray_values[i][self.cutoff_pixels:]):
                if above_threshold:
                    if g > threshold:
                        new_point = ray.T[t+self.cutoff_pixels]
                        candidate_points += [new_point]
                        break
                else:
                    if g < threshold:
                        new_point = ray.T[t+self.cutoff_pixels]
                        candidate_points += [new_point]
                        break

        return candidate_points

    def set_seed_points(self, initial_pupil_x, initial_pupil_y, initial_cr_x, initial_cr_y):
        self.initial_pupil_x = initial_pupil_x
        self.initial_pupil_y = initial_pupil_y
        self.initial_cr_x = initial_cr_x
        self.initial_cr_y = initial_cr_y

    def process_all_images(self):
        """ deprecated """

        # these aren't really used yet.
        self.pupil_loc = (0,0)
        self.cr_loc = (0,0)

        frame_list = os.listdir(self.input_image_folder)
        num_frames = len(frame_list)

        self.pupil_params = np.zeros([num_frames, 5])
        self.cr_params = np.zeros([num_frames, 5])

        for i,frame in enumerate(frame_list):
            logging.debug("Processing frame %d", i)
            if frame[-4:]!='.jpg' and frame[-4:]!='.png':  continue  # just in case some OS specific files snuck in (like in OS X)
            frame_path = os.path.join(self.input_image_folder,frame)

            # open Image, convert to gray scale and then to numpy array
            im = Image.open(frame_path)
            im = im.convert('L')
            im = np.array(im)


            # get pupil and corneal reflection parameters, this line is the actual eye tracking algorithm
            pupil, cr = self.process_image(im)

            # save results in arrays
            self.pupil_params[i] = (pupil[0][0],pupil[0][1],pupil[1],pupil[2][0],pupil[2][1])
            self.cr_params[i] = (cr[0][0],cr[0][1],cr[1],cr[2][0],cr[2][1])

        logging.debug("Saving pupil and cr parameters to:")
        logging.debug("\t%s", self.pupil_file)
        logging.debug("\t%s", self.cr_file)

        np.save(self.pupil_file, self.pupil_params)
        np.save(self.cr_file, self.cr_params)


    @staticmethod
    def rotate(X, Y, center_x, center_y, theta):

        Xp = (X-center_x)*np.cos(theta) - (Y-center_y)*np.sin(theta) + center_x
        Yp = (X-center_x)*np.sin(theta) + (Y-center_y)*np.cos(theta) + center_y

        return Xp, Yp

    @staticmethod
    def get_ellipse_mask(X, Y, params):

        center_x, center_y, theta, axis1, axis2 = params

        dX = X - center_x
        dY = Y - center_y

        theta = theta*np.pi/180.

        Xp = dX*np.cos(theta) + dY*np.sin(theta)
        Yp = -dX*np.sin(theta) + dY*np.cos(theta)

        mask1 = (Xp/axis1)**2 + (Yp/axis2)**2 < 1 + 0.1
        mask2 = (Xp/axis1)**2 + (Yp/axis2)**2 > 1 - 0.1

        mask = np.logical_and(mask1, mask2)

        return mask

    def annotate_frame_old(self, im, pupil=None, cr=None):

        y, x, c = im.shape

        X, Y = np.meshgrid(np.arange(x), np.arange(y))

        # pupil in red
        if pupil is not None:
            pupil_mask = self.get_ellipse_mask(X, Y, pupil)
            im.T[0].T[pupil_mask] = 255
            im.T[1].T[pupil_mask] = 0
            im.T[2].T[pupil_mask] = 0

        # cr in blue
        if cr is not None:
            cr_mask = self.get_ellipse_mask(X, Y, cr)
            im.T[0].T[cr_mask] = 0
            im.T[1].T[cr_mask] = 0
            im.T[2].T[cr_mask] = 255

        return im

    @classmethod
    def ellipse_points_from_params(cls, params):
        center_x, center_y, theta, axis1, axis2 = params
        theta = theta*np.pi/180. # convert to radians

        points_x = np.array([ axis1*np.cos(phi) + center_x for phi in np.linspace(0,2*np.pi, 1000)])
        points_y = np.array([ axis2*np.sin(phi) + center_y for phi in np.linspace(0,2*np.pi, 1000)])

        points_x, points_y = cls.rotate(points_x, points_y, center_x, center_y, theta)

        return points_x, points_y

    def annotate_frame(self, im, pupil=None, cr=None):
        
        y, x, c = im.shape

        im_pil = Image.fromarray(im)

        draw = ImageDraw.Draw(im_pil)

        if pupil is not None:
            points_x, points_y = self.ellipse_points_from_params(pupil)
            draw.point(zip(points_x, points_y), fill=(255,0,0))
            # for i, px in enumerate(points_x):
            #     im[int(points_y[i]), int(px), 0] = 255

        if cr is not None:
            points_x, points_y = self.ellipse_points_from_params(cr)
            draw.point(zip(points_x, points_y), fill=(0,0,255))
            # for i, px in enumerate(points_x):
            #     im[int(points_y[i]), int(px), 0] = 255

        # return im
        return np.array(im_pil)

    def annotate_frame_with_bbox(self, im, pupil_bbox=None, cr_bbox=None):

        # y, x, c = im.shape

        im_pil = Image.fromarray(im)

        draw = ImageDraw.Draw(im_pil)

        if pupil_bbox is not None:
            xmin, xmax, ymin, ymax = pupil_bbox
            # print pupil_bbox
            draw.rectangle([xmin,ymin,xmax,ymax],outline=(255,0,0))

            # for i, px in enumerate(points_x):
            #     im[int(points_y[i]), int(px), 0] = 255

        if cr_bbox is not None:
            # print cr_bbox
            xmin, xmax, ymin, ymax = cr_bbox
            draw.rectangle([xmin,ymin,xmax,ymax],outline=(0,0,255))
            # for i, px in enumerate(points_x):
            #     im[int(points_y[i]), int(px), 0] = 255

        # return im
        return np.array(im_pil)

    def annotate_frame_with_point(self, im, pupil=None, cr=None):

        # y, x, c = im.shape

        # print im.shape, im.dtype

        im_pil = Image.fromarray(im)

        draw = ImageDraw.Draw(im_pil)

        if pupil is not None:
            # points_x, points_y = ellipse_points_from_params(pupil)
            # draw.point(pupil, fill=(255,0,0))
            draw.ellipse([pupil[0]-5,pupil[1]-5,pupil[0]+5,pupil[1]+5],fill=(255,0,0))
            # for i, px in enumerate(points_x):
            #     im[int(points_y[i]), int(px), 0] = 255

        if cr is not None:
            # points_x, points_y = ellipse_points_from_params(cr)
            # draw.point(cr, fill=(0,0,255))
            draw.ellipse([cr[0]-5,cr[1]-5,cr[0]+5,cr[1]+5],fill=(0,0,255))
            # for i, px in enumerate(points_x):
            #     im[int(points_y[i]), int(px), 0] = 255

        # return im
        return np.array(im_pil)

    @staticmethod
    def get_frame_index(frame_name):

        return int(frame_name[12:-4])-1  # change 7 to 12

    # def annotate_frame(self, frame, im, pupil, cr):
    #     # this function is not done yet
    #     frame_index = self.get_frame_index(frame)
    #
    #     new_im = im.copy()
    #     pupil = self.pupil_params[frame_index]
    #     cr = self.cr_params[frame_index]
    #     new_im = self.annotate_frame(new_im, pupil, cr)
    #
    #     im_fig.set_data(new_im)
    #
    #     fig.savefig(os.path.join(self.frames_folder, input_frame), dpi=100)

    def output_annotation(self, frames_to_output=None):
        """generate a the series of images with eyetracking results superimposed"""

        fig, ax = plt.subplots(figsize=(4,3))  #, frameon=False)
        fig.subplots_adjust(left=0,right=1,bottom=0,top=1)

        ax.axis('off')
        # ax.axis('tight')

        self.pupil_params = np.load(self.pupil_file)
        self.cr_params = np.load(self.cr_file)

        if frames_to_output is None:
            frames_to_output = os.listdir(self.input_image_folder)

        first_frame = frames_to_output[0]

        frame_index = self.get_frame_index(first_frame)
        im = Image.open(os.path.join(self.input_image_folder, first_frame))
        im = np.array(im)

        new_im = np.dstack([im,im,im])
        pupil = self.pupil_params[frame_index]
        cr = self.cr_params[frame_index]
        new_im = self.annotate_frame(new_im, pupil, cr)

        im_fig = ax.imshow(new_im, aspect='normal')  #extent=(0,1,1,0)

        fig.savefig(os.path.join(self.frames_folder, first_frame), dpi=100)

        for input_frame in frames_to_output[1:]:

            frame_index = self.get_frame_index(input_frame)
            im = Image.open(os.path.join(self.input_image_folder, input_frame))
            im = np.array(im)

            new_im = np.dstack([im,im,im])
            pupil = self.pupil_params[frame_index]
            cr = self.cr_params[frame_index]
            new_im = self.annotate_frame(new_im, pupil, cr)

            im_fig.set_data(new_im)

            fig.savefig(os.path.join(self.frames_folder, input_frame), dpi=100)

    def output_QC(self, image_type='png'):
        """generate a set of summary statistics and plots for QC purposes"""

        logging.debug("saving QC images")
        self.pupil_params = np.load(self.pupil_file)
        self.cr_params = np.load(self.cr_file)

        logging.debug("saving pupil position")
        fig, ax = plt.subplots(1)
        ax.plot(self.pupil_params.T[0], label='pupil x')
        ax.plot(self.pupil_params.T[1], label='pupil y')
        ax.set_xlabel('frame index')
        ax.set_title('pupil position')
        ax.legend()
        fig.savefig(os.path.join(self.qc_folder, 'pupil_position.'+image_type))

        logging.debug("saving cr position")
        fig, ax = plt.subplots(1)
        ax.plot(self.cr_params.T[0], label='cr x')
        ax.plot(self.cr_params.T[1], label='cr y')
        ax.set_xlabel('frame index')
        ax.set_title('CR position')
        ax.legend()
        fig.savefig(os.path.join(self.qc_folder, 'cr_position.'+image_type))

        logging.debug("saving pupil axes")
        fig, ax = plt.subplots(1)
        ax.plot(self.pupil_params.T[3], label='pupil axis 1')
        ax.plot(self.pupil_params.T[4], label='pupil axis 2')
        ax.set_xlabel('frame index')
        ax.set_title('Pupil major and minor axis size')
        ax.legend()
        fig.savefig(os.path.join(self.qc_folder, 'pupil_axes.'+image_type))

        logging.debug("saving cr major/minor axis")
        fig, ax = plt.subplots(1)
        ax.plot(self.cr_params.T[3], label='cr axis 1')
        ax.plot(self.cr_params.T[4], label='cr axis 2')
        ax.set_xlabel('frame index')
        ax.set_title('CR major and minor axis size')
        ax.legend()
        fig.savefig(os.path.join(self.qc_folder, 'cr_axes.'+image_type))

        logging.debug("saving pupil angle")
        fig, ax = plt.subplots(1)
        ax.plot(self.pupil_params.T[2], label='pupil angle')
        ax.set_xlabel('frame index')
        ax.set_title('pupil major axis angle')
        ax.legend()
        fig.savefig(os.path.join(self.qc_folder, 'pupil_angle.'+image_type))

        logging.debug("saving cr angle")
        fig, ax = plt.subplots(1)
        ax.plot(self.cr_params.T[2], label='cr angle')
        ax.set_xlabel('frame index')
        ax.set_title('corneal reflection major axis angle')
        ax.legend()
        fig.savefig(os.path.join(self.qc_folder, 'cr_angle.'+image_type))


        logging.debug("computing density")
        # the remainder of these take a *very* long time
        T = self.pupil_params.shape[0]
        y, x = self.im_shape

        mean_frame = np.dstack([self.mean_frame,self.mean_frame,self.mean_frame])
        pupil_density = np.zeros((y, x, 3))
        # pupil_all = 255*np.ones(pupil_density.shape, np.uint8)
        # pupil_all = np.stack([mean_frame, mean_frame, mean_frame], axis=2)
        pupil_all = mean_frame.copy()

        cr_density = np.zeros((y, x, 3))
        # cr_all = 255*np.ones(cr_density.shape, np.uint8)
        # cr_all = np.stack([mean_frame, mean_frame, mean_frame], axis=2)
        cr_all = mean_frame.copy()
        temp = np.zeros((y,x,3), dtype=np.uint8)
        for t in range(T):
            if t % 100 == 0:
                logging.debug("finished %d frames", t)
            ptemp = self.annotate_frame(temp.copy(), self.pupil_params[t])
            pupil_density += ptemp #, self.cr_params[t])
            pupil_all = self.annotate_frame(pupil_all, self.pupil_params[t])

            crtemp = self.annotate_frame(temp.copy(), cr=self.cr_params[t])
            cr_density += crtemp #, self.cr_params[t])
            cr_all = self.annotate_frame(cr_all, cr=self.cr_params[t])

        logging.debug("plotting pupil density")
        fig, ax = plt.subplots(1)
        ax.imshow(np.log(1+pupil_density[:,:,0]), cmap='Greys', interpolation='nearest')
        # ax.axis('off')
        ax.set_title('Pupil ellipse density')
        fig.savefig(os.path.join(self.qc_folder, 'pupil_density.'+image_type))


        logging.debug("plotting pupil all")
        fig, ax = plt.subplots(1)
        ax.imshow(pupil_all, cmap='Greys', interpolation='nearest')
        # ax.axis('off')
        ax.set_title('All pupil ellipses combined')
        fig.savefig(os.path.join(self.qc_folder, 'pupil_all_plot.'+image_type))

        logging.debug("plotting cr density")
        fig, ax = plt.subplots(1)
        ax.imshow(np.log(1+cr_density[:,:,2]), cmap='Greys', interpolation='nearest')
        # ax.axis('off')
        ax.set_title('CR ellipse density')
        fig.savefig(os.path.join(self.qc_folder, 'cr_density.'+image_type))

        logging.debug("plotting cr all")
        fig, ax = plt.subplots(1)
        ax.imshow(cr_all, cmap='Greys', interpolation='nearest')
        ax.set_title('All CR ellipses combined')
        # ax.axis('off')
        fig.savefig(os.path.join(self.qc_folder, 'cr_all_plot.'+image_type))

