# FlyMovieFormat.py
# KMB 11/06/2008

import sys
import struct
import warnings
import os.path

import numpy as nx
from numpy import nan

import time

import math

# version 1 formats:
VERSION_FMT = '<I'
FORMAT_LEN_FMT = '<I'
BITS_PER_PIXEL_FMT = '<I' 
FRAMESIZE_FMT = '<II'
CHUNKSIZE_FMT = '<Q'
N_FRAME_FMT = '<Q'
TIMESTAMP_FMT = 'd' # XXX struct.pack('<d',nan) dies

# additional version 2 formats:
CHUNK_N_FRAME_FMT = '<Q'
CHUNK_TIMESTAMP_FMT = 'd' # XXX struct.pack('<d',nan) dies
CHUNK_DATASIZE_FMT = '<Q'

class NoMoreFramesException( Exception ):
    pass
    
class InvalidMovieFileException( Exception ):
    pass
    
class FlyMovie:
    
    def __init__(self, filename,check_integrity=False):
        self.filename = filename
        try:
            self.file = open(self.filename,mode="r+b")
        except IOError:
            self.file = open(self.filename,mode="r")
            self.writeable = False
        else:
            self.writeable = True
        
	# get the extension
	tmp,ext = os.path.splitext(self.filename)
	if ext == '.sbfmf':
	    self.init_sbfmf()
	    self.issbfmf = True
	    return
	else:
	    self.issbfmf = False
        
        r=self.file.read # shorthand
        t=self.file.tell # shorthand
        size=struct.calcsize
        unpack=struct.unpack

        version_buf = r(size(VERSION_FMT))
        if len(version_buf)!=size(VERSION_FMT):
            raise InvalidMovieFileException("could not read data file")
            
        version, = unpack(VERSION_FMT,version_buf)
        if version not in (1,3):
            raise NotImplementedError('Can only read version 1 and 3 files')

        if version  == 1:
            self.format = 'MONO8'
            self.bits_per_pixel = 8
        elif version == 3:
            format_len = unpack(FORMAT_LEN_FMT,r(size(FORMAT_LEN_FMT)))[0]
            self.format = r(format_len)
            self.bits_per_pixel = unpack(BITS_PER_PIXEL_FMT,r(size(BITS_PER_PIXEL_FMT)))[0]

        try:
            self.framesize = unpack(FRAMESIZE_FMT,r(size(FRAMESIZE_FMT)))
        except struct.error:
            raise InvalidMovieFileException('file could not be read')

        self.bytes_per_chunk, = unpack(CHUNKSIZE_FMT,r(size(CHUNKSIZE_FMT)))
        self.n_frames, = unpack(N_FRAME_FMT,r(size(N_FRAME_FMT)))
        self.timestamp_len = size(TIMESTAMP_FMT)
        self.chunk_start = self.file.tell()
        self.next_frame = None

        if self.bytes_per_chunk != self.bits_per_pixel/8*self.framesize[0]*self.framesize[1] + self.timestamp_len:
            print "FMF reading will probably end badly:", self.bytes_per_chunk, self.bits_per_pixel, self.framesize, self.timestamp_len, self.bits_per_pixel*self.framesize[0]*self.framesize[1] + self.timestamp_len

	if self.n_frames == 0: # unknown movie length, read to find out
            # seek to end of the movie
            self.file.seek(0,2)
            # get the byte position
            eb = self.file.tell()
            # compute number of frames using bytes_per_chunk
            self.n_frames = int((eb-self.chunk_start)/self.bytes_per_chunk)
            # seek back to the start
            self.file.seek(self.chunk_start,0)
            
        if check_integrity:
            n_frames_ok = False
            while not n_frames_ok:
                try:
                    self.get_frame(-1)
                    n_frames_ok = True
                except NoMoreFramesException:
                    self.n_frames -= 1
	    self.file.seek(self.chunk_start,0) # go back to beginning

        self._all_timestamps = None # cache

    def init_sbfmf(self):
	
	#try:
        # read the version number
        format = '<I'
        nbytesver, = struct.unpack(format,self.file.read(struct.calcsize(format)))
        version = self.file.read(nbytesver)
        
        # read header parameters
        format = '<4IQ'
        nr,nc,self.n_frames,difference_mode,self.indexloc = \
	      struct.unpack(format,self.file.read(struct.calcsize(format)))

        # read the background image
        self.bgcenter = nx.fromstring(self.file.read(struct.calcsize('<d')*nr*nc),'<d')
        # read the std
        self.bgstd = nx.fromstring(self.file.read(struct.calcsize('<d')*nr*nc),'<d')
        
        # read the index
        ff = self.file.tell()
        self.file.seek(self.indexloc,0)
        self.framelocs = nx.fromstring(self.file.read(self.n_frames*8),'<Q')

	#except:
        #    raise InvalidMovieFileException('file could not be read')

	if version == "0.1":
	    self.format = 'MONO8'
	    self.bits_per_pixel = 8

	self.framesize = (nr,nc)
	self.bytes_per_chunk = None
	self.timestamp_len = struct.calcsize('<d')
        self.chunk_start = self.file.tell()
        self.next_frame = None
        self._all_timestamps = None # cache

    def close(self):
        self.file.close()
        self.writeable = False
        self.n_frames = None
        self.next_frame = None
        
    def get_width(self):
        return self.framesize[1]

    def get_height(self):
        return self.framesize[0]

    def get_n_frames(self):
        return self.n_frames

    def get_format(self):
        return self.format

    def get_bits_per_pixel(self):
        return self.bits_per_pixel

    def read_some_bytes(self,nbytes):
        return self.file.read(nbytes)

    def _read_next_frame(self):
        if self.issbfmf:
            format = '<Id'
            try:
                npixels,timestamp = struct.unpack(format,self.file.read(struct.calcsize(format)))
                x = self.file.read(npixels*4)
                idx = nx.fromstring(x,'<I')
                v = nx.fromstring(self.file.read(npixels*1),'<B')
                frame = self.bgcenter.copy()
                frame[idx] = v
            except:
                print "sbfmf indexing error", self.file.tell()
                print len( x )
                print idx.shape, nx.max( idx )
                print frame.shape
                raise
            frame.shape = self.framesize
        else:
            data = self.file.read( self.bytes_per_chunk )
            if data == '':
                raise NoMoreFramesException('EOF')
            if len(data)<self.bytes_per_chunk:
                raise NoMoreFramesException('short frame')
            timestamp_buf = data[:self.timestamp_len]
            timestamp, = struct.unpack(TIMESTAMP_FMT,timestamp_buf)

            frame = nx.fromstring(data[self.timestamp_len:],'<B')
            frame.shape = self.framesize
        
##        if self.format == 'MONO8':
##            frame = nx.fromstring(data[self.timestamp_len:],nx.uint8)
##            frame.shape = self.framesize
##        elif self.format in ('YUV411','YUV422'):
##            frame = nx.fromstring(data[self.timestamp_len:],nx.uint16)
##            frame.shape = self.framesize
##        elif self.format in ('MONO16',):
##            print 'self.framesize',self.framesize
##            frame = nx.fromstring(data[self.timestamp_len:],nx.uint8)
##            frame.shape = self.framesize
##        else:
##            raise NotImplementedError("Reading not implemented for %s format"%(self.format,))
        return frame, timestamp
        
    def _read_next_timestamp(self):
        if self.issbfmf:
            format = '<Id'
            self.npixelscurr,timestamp = struct.unpack(format,self.file.read(struct.calcsize(format)))
            return timestamp
        read_len = struct.calcsize(TIMESTAMP_FMT)
        timestamp_buf = self.file.read( read_len )
        self.file.seek( self.bytes_per_chunk-read_len, 1) # seek to next frame
        if timestamp_buf == '':
            raise NoMoreFramesException('EOF')
        timestamp, = struct.unpack(TIMESTAMP_FMT,timestamp_buf)
        return timestamp
        
    def is_another_frame_available(self):
        try:
            if self.next_frame is None:
                self.next_frame = self._read_next_frame()
        except NoMoreFramesException:
            return False
        return True
        
    def get_next_frame(self):
        if self.next_frame is not None:
            frame, timestamp = self.next_frame
            self.next_frame = None
            return frame, timestamp
        else:
            frame, timestamp = self._read_next_frame()
            return frame, timestamp

    def get_frame(self,frame_number):
        if frame_number < 0:
            frame_number = self.n_frames + frame_number
        if frame_number < 0:
            raise IndexError("index out of range (n_frames = %d)"%self.n_frames)
        if self.issbfmf:
            seek_to = self.framelocs[frame_number]
        else:
            seek_to = self.chunk_start+self.bytes_per_chunk*frame_number
        self.file.seek(seek_to)
        self.next_frame = None
        try:
            x = self.get_next_frame()
        except:
            print "error after seeking to", seek_to, "for frame", frame_number
            raise
        else:
            return x
    
    def get_all_timestamps(self):
        if self._all_timestamps is None:

            self._all_timestamps = []

            if self.issbfmf:
                self.seek(0)
                format = '<Id'
                l = struct.calcsize(format)
                for i in range(self.n_frames):
                    self.seek(i)
                    npixels,timestamp = struct.unpack(format,self.file.read(l))
                    self._all_timestamps.append(timestamp)
            else:
                self.seek(0)
                read_len = struct.calcsize(TIMESTAMP_FMT)
                while 1:
                    timestamp_buf = self.file.read( read_len )
                    self.file.seek( self.bytes_per_chunk-read_len, 1) # seek to next frame
                    if timestamp_buf == '':
                        break
                    timestamp, = struct.unpack(TIMESTAMP_FMT,timestamp_buf)
                    self._all_timestamps.append( timestamp )
            self.next_frame = None
            self._all_timestamps = nx.asarray(self._all_timestamps)
        return self._all_timestamps
        
    def seek(self,frame_number):
        if frame_number < 0:
            frame_number = self.n_frames + frame_number
        if self.issbfmf:
            seek_to = self.framelocs[frame_number]
        else:
            seek_to = self.chunk_start+self.bytes_per_chunk*frame_number
        self.file.seek(seek_to)
        self.next_frame = None

    def get_next_timestamp(self):
        if self.next_frame is not None:
            frame, timestamp = self.next_frame
            self.next_frame = None
            return timestamp
        else:
            timestamp = self._read_next_timestamp()
            return timestamp
        
    def get_frame_at_or_before_timestamp(self, timestamp):
        tss = self.get_all_timestamps()
        at_or_before_timestamp_cond = tss <= timestamp
        nz = nx.nonzero(at_or_before_timestamp_cond)
        if len(nz)==0:
            raise ValueError("no frames at or before timestamp given")
        fno = nz[-1]
        return self.get_frame(fno)
        
class FlyMovieSaver:
    def __init__(self,
                 filename,
                 version=1,
                 seek_ok=True,
                 
                 compressor=None,
                 comp_level=1,
                 
                 format=None,
                 bits_per_pixel=None,
                 ):
        """create a FlyMovieSaver instance

        arguments:

          filename
          version    -- 1, 2, or 3
          seek_ok    -- is seek OK on this filename?

          For version 2:
          --------------
          compressor -- None or 'lzo' (only used if version == 2)
          comp_level -- compression level (only used if compressed)
          
          For version 3:
          --------------
          format     -- string representing format (e.g. 'MONO8' or 'YUV422')
          bits_per_pixel -- number of bytes per pixel (MONO8 = 8, YUV422 = 16)
          

        """

        # filename
        path, ext = os.path.splitext(filename)
        if ext == '':
            ext = '.fmf'
        self.filename = path+ext

        # seek_ok
        if seek_ok:
            mode = "w+b"
        else:
            mode = "wb"
        self.seek_ok = seek_ok
            
        self.file = open(self.filename,mode=mode)

        if version == 1:
            self.add_frame = self._add_frame_v1
            self.add_frames = self._add_frames_v1
        elif version == 2:
            self.add_frame = self._add_frame_v2
            self.add_frames = self._add_frames_v2
        elif version == 3:
            self.add_frame = self._add_frame_v1
            self.add_frames = self._add_frames_v1
        else:
            raise ValueError('only versions 1, 2, and 3 exist')

        self.file.write(struct.pack(VERSION_FMT,version))
        
        if version == 2:
            self.compressor = compressor
            if self.compressor is None:
                self.compressor = 'non'
                
            if self.compressor == 'non':
    	        self.compress_func = lambda x: x
    	    elif self.compressor == 'lzo':
    	        import lzo
    	        self.compress_func = lzo.compress
    	    else:
    	        raise ValueError("unknown compressor '%s'"%(self.compressor,))
    	    assert type(self.compressor) == str and len(self.compressor)<=4
    	    self.file.write(self.compressor)
            
        if version == 3:
            if type(format) != str:
                raise ValueError("format must be string (e.g. 'MONO8', 'YUV422')")
            if type(bits_per_pixel) != int:
                raise ValueError("bits_per_pixel must be integer")
            format_len = len(format)
            self.file.write(struct.pack(FORMAT_LEN_FMT,format_len))
            self.file.write(format)
            self.file.write(struct.pack(BITS_PER_PIXEL_FMT,bits_per_pixel))
            
            self.format = format
            self.bits_per_pixel = bits_per_pixel
        else:
            self.format = 'MONO8'
            self.bits_per_pixel = 8

        self.framesize = None
        
        self.n_frames = 0
        self.n_frame_pos = None

    def _add_frame_v1(self,origframe,timestamp=nan):
        TIMESTAMP_FMT = 'd' # XXX struct.pack('<d',nan) dies
        frame = nx.asarray(origframe)
        if self.framesize is None:
            self._do_v1_header(frame)
        else:
            if self.framesize != frame.shape:
                raise ValueError('frame shape is now %s, but it used to be %s'%(str(frame.shape),str(self.framesize)))

        b1 = struct.pack(TIMESTAMP_FMT,timestamp)
        self.file.write(b1)
        if hasattr(origframe,'dump_to_file'):
            nbytes = origframe.dump_to_file( self.file )
            assert nbytes == self._bytes_per_image
        else:
            if not hasattr(self,'gave_dump_fd_warning'):
                warnings.warn('could save faster if %s implemented dump_to_file()'%(str(type(origframe)),))
                self.gave_dump_fd_warning = True
            b2 = frame.tostring()
            if len(b2) != self._bytes_per_image:
                raise ValueError("expected buffer of length %d, got length %d (shape %s)"%(self._bytes_per_image,len(b2),str(frame.shape)))
            self.file.write(b2)
        self.n_frames += 1

    def _add_frames_v1(self, frames, timestamps=None):
        if 0:
            for frame, timestamp in zip(frames,timestamps):
                self._add_frame_v1(frame,timestamp)
        else:
            if timestamps is None:
                timestamps = [nan]*len(frames)
            TIMESTAMP_FMT = 'd' # XXX struct.pack('<d',nan) dies
            if self.framesize is None:
                self._do_v1_header(frames[0])
            else:
                assert self.framesize == frames[0].shape
            mega_buffer = ''
            for frame, timestamp in zip(frames,timestamps):
                b1 = struct.pack(TIMESTAMP_FMT,timestamp)            
                mega_buffer += b1
                b2 = frame.tostring()
                assert len(b2) == self._bytes_per_image
                mega_buffer += b2
            self.file.write(mega_buffer)
            self.n_frames += len(frames)
        
    def _do_v1_header(self,frame):
        # first frame

        # frame data are always type uint8, so frame shape (width) varies if data format not MONO8
        self.framesize = frame.shape

        buf = struct.pack(FRAMESIZE_FMT,frame.shape[0],frame.shape[1])
        self.file.write(buf)

        #bits_per_image = frame.shape[0] * frame.shape[1] * self.bits_per_pixel
        bits_per_image = frame.shape[0] * frame.shape[1] * 8
        if bits_per_image % 8 != 0:
            raise ValueError('combination of frame size and bits_per_pixel make non-byte aligned image')
        self._bytes_per_image = bits_per_image / 8
        bytes_per_chunk = self._bytes_per_image + struct.calcsize(TIMESTAMP_FMT)

        buf = struct.pack(CHUNKSIZE_FMT,bytes_per_chunk)
        self.file.write(buf)

        self.n_frame_pos = self.file.tell()

        buf = struct.pack(N_FRAME_FMT,self.n_frames) # will fill in later
        self.file.write(buf)

        ####### end of header ###########################
        
    def _add_frames_v2(self, frames, timestamps=None):
        if self.framesize is None:
            # header stuff dependent on first frame
            
            frame = frames[0]
            assert len(frame.shape) == 2 # must be MxN array
            self.framesize = frame.shape
            
            buf = struct.pack(FRAMESIZE_FMT,frame.shape[0],frame.shape[1])
            self.file.write(buf)
            
            self.n_frame_pos = self.file.tell() # may fill value later
            buf = struct.pack(N_FRAME_FMT,self.n_frames) 
            self.file.write(buf)
            
            ####### end of header ###########################

        # begin chunk
        chunk_n_frames = len(frames)
        if timestamps is None:
            timestamps = [nan]*chunk_n_frames
        else:
            assert len(timestamps) == chunk_n_frames

        buf = struct.pack(CHUNK_N_FRAME_FMT,chunk_n_frames)
        self.file.write(buf)

        for timestamp in timestamps:
            self.file.write(struct.pack(CHUNK_TIMESTAMP_FMT,timestamp))

        # generate mega string
        bufs = [ frame.tostring() for frame in frames ]
        buf = ''.join(bufs)
        del bufs
        compressed_data = self.compress_func(buf)
        
        chunk_datasize = len( compressed_data)
        self.file.write(struct.pack(CHUNK_DATASIZE_FMT,chunk_datasize))

        self.file.write(compressed_data)
        
        self.n_frames += chunk_n_frames

    def _add_frame_v2(self,frame,timestamp=nan):
        self._add_frames_v2([frame],[timestamp])

    def close(self):
        if self.n_frames == 0:
            warnings.warn('no frames in FlyMovie')
            # no frames added
            self.file.close()
            del self.file
            return

        if self.seek_ok:
            self.file.seek( self.n_frame_pos )
            buf = struct.pack(N_FRAME_FMT,self.n_frames) # will fill in later
            self.file.write(buf)
        self.file.close()
        del self.file # make sure we can't access self.file again

    def __del__(self):
        if hasattr(self,'file'):
            self.close()

