''' Classes for read / write of matlab (TM) 5 files

The matfile specification last found here:

http://www.mathworks.com/access/helpdesk/help/pdf_doc/matlab/matfile_format.pdf

(as of December 5 2008)
'''

# Small fragments of current code adapted from matfile.py by Heiko
# Henkelmann

import os
import time
import sys
import zlib
from StringIO import StringIO
from cStringIO import StringIO as cStringIO
from copy import copy as pycopy
import warnings

import numpy as np

import scipy.sparse

import byteordercodes
from miobase import MatFileReader, MatArrayReader, MatMatrixGetter, \
     MatFileWriter, MatStreamWriter, docfiller, matdims, \
     MatReadError

miINT8 = 1
miUINT8 = 2
miINT16 = 3
miUINT16 = 4
miINT32 = 5
miUINT32 = 6
miSINGLE = 7
miDOUBLE = 9
miINT64 = 12
miUINT64 = 13
miMATRIX = 14
miCOMPRESSED = 15
miUTF8 = 16
miUTF16 = 17
miUTF32 = 18

mxCELL_CLASS = 1
mxSTRUCT_CLASS = 2
# The March 2008 edition of "Matlab 7 MAT-File Format" says that
# mxOBJECT_CLASS = 3, whereas matrix.h says that mxLOGICAL = 3.
# Matlab 2008a appears to save logicals as type 9, so we assume that
# the document is correct.  See type 18, below.
mxOBJECT_CLASS = 3
mxCHAR_CLASS = 4
mxSPARSE_CLASS = 5
mxDOUBLE_CLASS = 6
mxSINGLE_CLASS = 7
mxINT8_CLASS = 8
mxUINT8_CLASS = 9
mxINT16_CLASS = 10
mxUINT16_CLASS = 11
mxINT32_CLASS = 12
mxUINT32_CLASS = 13
# The following are not in the March 2008 edition of "Matlab 7
# MAT-File Format," but were guessed from matrix.h.
mxINT64_CLASS = 14
mxUINT64_CLASS = 15
mxFUNCTION_CLASS = 16
# Not doing anything with these at the moment.
mxOPAQUE_CLASS = 17 # This appears to be a function workspace
# https://www-old.cae.wisc.edu/pipermail/octave-maintainers/2007-May/002824.html
mxOBJECT_CLASS_FROM_MATRIX_H = 18

mxmap = { # Sometimes good for debug prints
    mxCELL_CLASS: 'mxCELL_CLASS',
    mxSTRUCT_CLASS: 'mxSTRUCT_CLASS',
    mxOBJECT_CLASS: 'mxOBJECT_CLASS',
    mxCHAR_CLASS: 'mxCHAR_CLASS',
    mxSPARSE_CLASS: 'mxSPARSE_CLASS',
    mxDOUBLE_CLASS: 'mxDOUBLE_CLASS',
    mxSINGLE_CLASS: 'mxSINGLE_CLASS',
    mxINT8_CLASS: 'mxINT8_CLASS',
    mxUINT8_CLASS: 'mxUINT8_CLASS',
    mxINT16_CLASS: 'mxINT16_CLASS',
    mxUINT16_CLASS: 'mxUINT16_CLASS',
    mxINT32_CLASS: 'mxINT32_CLASS',
    mxUINT32_CLASS: 'mxUINT32_CLASS',
    mxINT64_CLASS: 'mxINT64_CLASS',
    mxUINT64_CLASS: 'mxUINT64_CLASS',
    mxFUNCTION_CLASS: 'mxFUNCTION_CLASS',
    mxOPAQUE_CLASS: 'mxOPAQUE_CLASS',
    mxOBJECT_CLASS_FROM_MATRIX_H: 'mxOBJECT_CLASS_FROM_MATRIX_H',
}

mdtypes_template = {
    miINT8: 'i1',
    miUINT8: 'u1',
    miINT16: 'i2',
    miUINT16: 'u2',
    miINT32: 'i4',
    miUINT32: 'u4',
    miSINGLE: 'f4',
    miDOUBLE: 'f8',
    miINT64: 'i8',
    miUINT64: 'u8',
    miUTF8: 'u1',
    miUTF16: 'u2',
    miUTF32: 'u4',
    'file_header': [('description', 'S116'),
                    ('subsystem_offset', 'i8'),
                    ('version', 'u2'),
                    ('endian_test', 'S2')],
    'tag_full': [('mdtype', 'u4'), ('byte_count', 'u4')],
    'tag_smalldata':[('byte_count_mdtype', 'u4'), ('data', 'S4')],
    'array_flags': [('data_type', 'u4'),
                    ('byte_count', 'u4'),
                    ('flags_class','u4'),
                    ('nzmax', 'u4')],
    'U1': 'U1',
    }

mclass_dtypes_template = {
    mxINT8_CLASS: 'i1',
    mxUINT8_CLASS: 'u1',
    mxINT16_CLASS: 'i2',
    mxUINT16_CLASS: 'u2',
    mxINT32_CLASS: 'i4',
    mxUINT32_CLASS: 'u4',
    mxINT64_CLASS: 'i8',
    mxUINT64_CLASS: 'u8',
    mxSINGLE_CLASS: 'f4',
    mxDOUBLE_CLASS: 'f8',
    }


np_to_mtypes = {
    'f8': miDOUBLE,
    'c32': miDOUBLE,
    'c24': miDOUBLE,
    'c16': miDOUBLE,
    'f4': miSINGLE,
    'c8': miSINGLE,
    'i1': miINT8,
    'i2': miINT16,
    'i4': miINT32,
    'i8': miINT64,
    'u1': miUINT8,
    'u2': miUINT16,
    'u4': miUINT32,
    'u8': miUINT64,
    'S1': miUINT8,
    'U1': miUTF16,
    }


np_to_mxtypes = {
    'f8': mxDOUBLE_CLASS,
    'c32': mxDOUBLE_CLASS,
    'c24': mxDOUBLE_CLASS,
    'c16': mxDOUBLE_CLASS,
    'f4': mxSINGLE_CLASS,
    'c8': mxSINGLE_CLASS,
    'i8': mxINT64_CLASS,
    'i4': mxINT32_CLASS,
    'i2': mxINT16_CLASS,
    'u8': mxUINT64_CLASS,
    'u2': mxUINT16_CLASS,
    'u1': mxUINT8_CLASS,
    'S1': mxUINT8_CLASS,
    }



''' Before release v7.1 (release 14) matlab (TM) used the system
default character encoding scheme padded out to 16-bits. Release 14
and later use Unicode. When saving character data, R14 checks if it
can be encoded in 7-bit ascii, and saves in that format if so.'''

codecs_template = {
    miUTF8: {'codec': 'utf_8', 'width': 1},
    miUTF16: {'codec': 'utf_16', 'width': 2},
    miUTF32: {'codec': 'utf_32','width': 4},
    }

miUINT16_codec = sys.getdefaultencoding()

mx_numbers = (
    mxDOUBLE_CLASS,
    mxSINGLE_CLASS,
    mxINT8_CLASS,
    mxUINT8_CLASS,
    mxINT16_CLASS,
    mxUINT16_CLASS,
    mxINT32_CLASS,
    mxUINT32_CLASS,
    mxINT64_CLASS,
    mxUINT64_CLASS,
    )


class mat_struct(object):
    ''' Placeholder for holding read data from structs

    We will deprecate this method of holding struct information in a
    future version of scipy, in favor of the recarray method (see
    loadmat docstring)
    '''
    pass


class MatlabObject(np.ndarray):
    ''' ndarray Subclass to contain matlab object '''
    def __new__(cls, input_array, classname=None):
        # Input array is an already formed ndarray instance
        # We first cast to be our class type
        obj = np.asarray(input_array).view(cls)
        # add the new attribute to the created instance
        obj.classname = classname
        # Finally, we must return the newly created object:
        return obj

    def __array_finalize__(self,obj):
        # reset the attribute from passed original object
        self.classname = getattr(obj, 'classname', None)
        # We do not need to return anything


class MatlabFunction(np.ndarray):
    ''' Subclass to signal this is a matlab function '''
    def __new__(cls, input_array):
        obj = np.asarray(input_array).view(cls)


class MatlabBinaryBlock(object):
    ''' Class to contain matlab unreadable blocks '''
    def __init__(self, binaryblock, endian):
        self.binaryblock = binaryblock
        self.endian = endian


class Mat5ArrayReader(MatArrayReader):
    ''' Class to get Mat5 arrays

    Provides element reader functions, header reader, matrix reader
    factory function
    '''

    def __init__(self,
                 mat_stream,
                 dtypes,
                 processor_func,
                 codecs,
                 class_dtypes,
                 struct_as_record):
        super(Mat5ArrayReader, self).__init__(mat_stream,
                                              dtypes,
                                              processor_func)
        self.codecs = codecs
        self.class_dtypes = class_dtypes
        self.struct_as_record = struct_as_record

    def read_element(self, copy=True):
        raw_tag = self.mat_stream.read(8)
        tag = np.ndarray(shape=(),
                         dtype=self.dtypes['tag_full'],
                         buffer=raw_tag)
        mdtype = tag['mdtype'].item()
        # Byte count if this is small data element
        byte_count = mdtype >> 16
        if byte_count: # small data element format
            if byte_count > 4:
                raise ValueError, 'Too many bytes for sde format'
            mdtype = mdtype & 0xFFFF
            if mdtype == miMATRIX:
                raise TypeError('Cannot have matrix in SDE format')
            raw_str = raw_tag[4:byte_count+4]
        else: # regular element
            byte_count = tag['byte_count'].item()
            # Deal with miMATRIX type (cannot pass byte string)
            if mdtype == miMATRIX:
                return self.current_getter(byte_count).get_array()
            # All other types can be read from string
            raw_str = self.mat_stream.read(byte_count)
            # Seek to next 64-bit boundary
            mod8 = byte_count % 8
            if mod8:
                self.mat_stream.seek(8 - mod8, 1)

        if mdtype in self.codecs: # encoded char data
            codec = self.codecs[mdtype]
            if not codec:
                raise TypeError, 'Do not support encoding %d' % mdtype
            el = raw_str.decode(codec)
        else: # numeric data
            dt = self.dtypes[mdtype]
            el_count = byte_count // dt.itemsize
            el = np.ndarray(shape=(el_count,),
                            dtype=dt,
                            buffer=raw_str)
            if copy:
                el = el.copy()

        return el

    def matrix_getter_factory(self):
        ''' Returns reader for next matrix at top level '''
        tag = self.read_dtype(self.dtypes['tag_full'])
        mdtype = tag['mdtype'].item()
        byte_count = tag['byte_count'].item()
        next_pos = self.mat_stream.tell() + byte_count
        if mdtype == miCOMPRESSED:
            getter = Mat5ZArrayReader(self, byte_count).matrix_getter_factory()
        elif not mdtype == miMATRIX:
            raise TypeError, \
                  'Expecting miMATRIX type here, got %d' %  mdtype
        else:
            getter = self.current_getter(byte_count)
        getter.next_position = next_pos
        return getter

    def current_getter(self, byte_count):
        ''' Return matrix getter for current stream position

        Returns matrix getters at top level and sub levels
        '''
        if not byte_count: # an empty miMATRIX can contain no bytes
            return Mat5EmptyMatrixGetter(self)
        af = self.read_dtype(self.dtypes['array_flags'])
        header = {}
        flags_class = af['flags_class']
        mc = flags_class & 0xFF
        header['mclass'] = mc
        header['is_logical'] = flags_class >> 9 & 1
        header['is_global'] = flags_class >> 10 & 1
        header['is_complex'] = flags_class >> 11 & 1
        header['nzmax'] = af['nzmax']
        ''' Here I am playing with a binary block read of
        untranslatable data. I am not using this at the moment because
        reading it has the side effect of making opposite ending mat
        files unwritable on the round trip.
        
        if mc == mxFUNCTION_CLASS:
            # we can't read these, and want to keep track of the byte
            # count - so we need to avoid the following unpredictable
            # length element reads
            return Mat5BinaryBlockGetter(self,
                                         header,
                                         af,
                                         byte_count)
        '''
        header['dims'] = self.read_element()
        header['name'] = self.read_element().tostring()
        # maybe a dictionary mapping here as a dispatch table
        if mc in mx_numbers:
            return Mat5NumericMatrixGetter(self, header)
        if mc == mxSPARSE_CLASS:
            return Mat5SparseMatrixGetter(self, header)
        if mc == mxCHAR_CLASS:
            return Mat5CharMatrixGetter(self, header)
        if mc == mxCELL_CLASS:
            return Mat5CellMatrixGetter(self, header)
        if mc == mxSTRUCT_CLASS:
            return Mat5StructMatrixGetter(self, header)
        if mc == mxOBJECT_CLASS:
            return Mat5ObjectMatrixGetter(self, header)
        if mc == mxFUNCTION_CLASS:
            return Mat5FunctionGetter(self, header)
        raise TypeError, 'No reader for class code %s' % mc


class Mat5ZArrayReader(Mat5ArrayReader):
    ''' Getter for compressed arrays

    Sets up reader for gzipped stream on init, providing wrapper
    for this new sub-stream.

    '''
    def __init__(self, array_reader, byte_count):
        super(Mat5ZArrayReader, self).__init__(
            cStringIO(zlib.decompress(
                        array_reader.mat_stream.read(byte_count))),
            array_reader.dtypes,
            array_reader.processor_func,
            array_reader.codecs,
            array_reader.class_dtypes,
            array_reader.struct_as_record)


class Mat5MatrixGetter(MatMatrixGetter):
    ''' Base class for getting Mat5 matrices

    Gets current read information from passed array_reader
    '''

    def __init__(self, array_reader, header):
        super(Mat5MatrixGetter, self).__init__(array_reader, header)
        self.class_dtypes = array_reader.class_dtypes
        self.codecs = array_reader.codecs
        self.is_global = header['is_global']
        self.mat_dtype = None

    def read_element(self, *args, **kwargs):
        return self.array_reader.read_element(*args, **kwargs)


class Mat5EmptyMatrixGetter(Mat5MatrixGetter):
    ''' Dummy class to return empty array for empty matrix
    '''
    def __init__(self, array_reader):
        self.array_reader = array_reader
        self.mat_stream = array_reader.mat_stream
        self.header = {}
        self.name = ''
        self.is_global = False
        self.mat_dtype = 'f8'

    def get_raw_array(self):
        return np.array([[]])


class Mat5NumericMatrixGetter(Mat5MatrixGetter):

    def __init__(self, array_reader, header):
        super(Mat5NumericMatrixGetter, self).__init__(array_reader, header)
        if header['is_logical']:
            self.mat_dtype = np.dtype('bool')
        else:
            self.mat_dtype = self.class_dtypes[header['mclass']]

    def get_raw_array(self):
        if self.header['is_complex']:
            # avoid array copy to save memory
            res = self.read_element(copy=False)
            res_j = self.read_element(copy=False)
            res = res + (res_j * 1j)
        else:
            res = self.read_element()
        return np.ndarray(shape=self.header['dims'],
                          dtype=res.dtype,
                          buffer=res,
                          order='F')


class Mat5SparseMatrixGetter(Mat5MatrixGetter):
    def get_raw_array(self):
        rowind = self.read_element()
        indptr = self.read_element()
        if self.header['is_complex']:
            # avoid array copy to save memory
            data   = self.read_element(copy=False)
            data_j = self.read_element(copy=False)
            data = data + (data_j * 1j)
        else:
            data = self.read_element()
        ''' From the matlab (TM) API documentation, last found here:
        http://www.mathworks.com/access/helpdesk/help/techdoc/matlab_external/
        rowind are simply the row indices for all the (nnz) non-zero
        entries in the sparse array.  rowind has nzmax entries, so
        may well have more entries than nnz, the actual number
        of non-zero entries, but rowind[nnz:] can be discarded
        and should be 0. indptr has length (number of columns + 1),
        and is such that, if D = diff(colind), D[j] gives the number
        of non-zero entries in column j. Because rowind values are
        stored in column order, this gives the column corresponding to
        each rowind
        '''
        M,N = self.header['dims']
        indptr = indptr[:N+1]
        nnz = indptr[-1]
        rowind = rowind[:nnz]
        data   = data[:nnz]
        return scipy.sparse.csc_matrix(
            (data,rowind,indptr),
            shape=(M,N))


class Mat5CharMatrixGetter(Mat5MatrixGetter):
    def get_raw_array(self):
        res = self.read_element()
        # Convert non-string types to unicode
        if isinstance(res, np.ndarray):
            if res.dtype.type == np.uint16:
                codec = miUINT16_codec
                if self.codecs['uint16_len'] == 1:
                    res = res.astype(np.uint8)
            elif res.dtype.type in (np.uint8, np.int8):
                codec = 'ascii'
            else:
                raise TypeError, 'Did not expect type %s' % res.dtype
            res = res.tostring().decode(codec)
        return np.ndarray(shape=self.header['dims'],
                          dtype=np.dtype('U1'),
                          buffer=np.array(res),
                          order='F').copy()


class Mat5CellMatrixGetter(Mat5MatrixGetter):
    def get_raw_array(self):
        # Account for fortran indexing of cells
        tupdims = tuple(self.header['dims'][::-1])
        length = np.product(tupdims)
        result = np.empty(length, dtype=object)
        for i in range(length):
            result[i] = self.get_item()
        return result.reshape(tupdims).T

    def get_item(self):
        return self.read_element()


class Mat5StructMatrixGetter(Mat5MatrixGetter):
    def __init__(self, array_reader, header):
        super(Mat5StructMatrixGetter, self).__init__(array_reader, header)
        self.struct_as_record = array_reader.struct_as_record

    def get_raw_array(self):
        namelength = self.read_element()[0]
        names = self.read_element()
        field_names = [names[i:i+namelength].tostring().strip('\x00')
                       for i in xrange(0,len(names),namelength)]
        tupdims = tuple(self.header['dims'][::-1])
        length = np.product(tupdims)
        if self.struct_as_record:
            if not len(field_names):
                # If there are no field names, there is no dtype
                # representation we can use, falling back to empty
                # object
                return np.empty(tupdims, dtype=object).T
            dtype = [(field_name, object) for field_name in field_names]
            result = np.empty(length, dtype=dtype)
            for i in range(length):
                for field_name in field_names:
                    result[i][field_name] = self.read_element()
        else: # Backward compatibility with previous format
            self.obj_template = mat_struct()
            self.obj_template._fieldnames = field_names
            result = np.empty(length, dtype=object)
            for i in range(length):
                item = pycopy(self.obj_template)
                for name in field_names:
                    item.__dict__[name] = self.read_element()
                result[i] = item
        return result.reshape(tupdims).T


class Mat5ObjectMatrixGetter(Mat5StructMatrixGetter):
    def get_raw_array(self):
        '''Matlab objects are like structs, with an extra classname field'''
        classname = self.read_element().tostring()
        result = super(Mat5ObjectMatrixGetter, self).get_raw_array()
        return MatlabObject(result, classname)


class Mat5FunctionGetter(Mat5ObjectMatrixGetter):
    ''' Class to provide warning and message string for unreadable
    matlab function data
    '''
    def get_raw_array(self):
        raise MatReadError('Cannot read matlab functions')


class Mat5BinaryBlockGetter(object):
    ''' Class to read in unreadable binary blocks

    This class could be used to read in matlab functions
    '''

    def __init__(self,
                 array_reader,
                 header,
                 array_flags,
                 byte_count):
        self.array_reader = array_reader
        self.header = header
        self.array_flags = array_flags
        arr_str = array_flags.tostring()
        self.binaryblock = array_reader.mat_stream.read(
            byte_count-len(array_flags.tostring()))
        stream = StringIO(self.binaryblock)
        reader = Mat5ArrayReader(
            stream,
            array_reader.dtypes,
            lambda x : None,
            array_reader.codecs,
            array_reader.class_dtypes,
            False)
        self.header['dims'] = reader.read_element()
        self.header['name'] = reader.read_element().tostring()
        self.name = self.header['name']
        self.is_global = header['is_global']

    def get_array(self):
        dt = self.array_reader.dtypes[miINT32]
        endian = byteordercodes.to_numpy_code(dt.byteorder)
        data = self.array_flags.tostring() + self.binaryblock
        return MatlabBinaryBlock(data, endian)

               
class MatFile5Reader(MatFileReader):
    ''' Reader for Mat 5 mat files
    Adds the following attribute to base class

    uint16_codec       - char codec to use for uint16 char arrays
                          (defaults to system default codec)
   '''
    @docfiller
    def __init__(self,
                 mat_stream,
                 byte_order=None,
                 mat_dtype=False,
                 squeeze_me=False,
                 chars_as_strings=True,
                 matlab_compatible=False,
                 struct_as_record=None, # default False, for now
                 uint16_codec=None
                 ):
        '''Initializer for matlab 5 file format reader

    %(matstream_arg)s
    %(load_args)s
    %(struct_arg)s
    uint16_codec : {None, string}
        Set codec to use for uint16 char arrays (e.g. 'utf-8').
        Use system default codec if None
        '''
        # Deal with deprecations
        if struct_as_record is None:
            warnings.warn("Using struct_as_record default value (False)" +
                          " This will change to True in future versions",
                          FutureWarning, stacklevel=2)
            struct_as_record = False
        self.codecs = {}
        # Missing inputs to array reader set later (processor func
        # below, dtypes, codecs via our own set_dtype function, called
        # from parent __init__)
        self._array_reader = Mat5ArrayReader(
            mat_stream,
            None,
            None,
            None,
            None,
            struct_as_record
            )
        super(MatFile5Reader, self).__init__(
            mat_stream,
            byte_order,
            mat_dtype,
            squeeze_me,
            chars_as_strings,
            matlab_compatible,
            )
        self._array_reader.processor_func = self.processor_func
        self.uint16_codec = uint16_codec

    def get_uint16_codec(self):
        return self._uint16_codec
    def set_uint16_codec(self, uint16_codec):
        if not uint16_codec:
            uint16_codec = sys.getdefaultencoding()
        # Set length of miUINT16 char encoding
        self.codecs['uint16_len'] = len("  ".encode(uint16_codec)) \
                               - len(" ".encode(uint16_codec))
        self.codecs['uint16_codec'] = uint16_codec
        self._array_reader.codecs = self.codecs
        self._uint16_codec = uint16_codec
    uint16_codec = property(get_uint16_codec,
                            set_uint16_codec,
                            None,
                            'get/set uint16_codec')

    def set_dtypes(self):
        ''' Set dtypes and codecs '''
        self.dtypes = self.convert_dtypes(mdtypes_template)
        self.class_dtypes = self.convert_dtypes(mclass_dtypes_template)
        codecs = {}
        postfix = self.order_code == '<' and '_le' or '_be'
        for k, v in codecs_template.items():
            codec = v['codec']
            try:
                " ".encode(codec)
            except LookupError:
                codecs[k] = None
                continue
            if v['width'] > 1:
                codec += postfix
            codecs[k] = codec
        self.codecs.update(codecs)
        self.update_array_reader()

    def update_array_reader(self):
        self._array_reader.codecs = self.codecs
        self._array_reader.dtypes = self.dtypes
        self._array_reader.class_dtypes = self.class_dtypes

    def matrix_getter_factory(self):
        return self._array_reader.matrix_getter_factory()

    def guess_byte_order(self):
        ''' Guess byte order.
        Sets stream pointer to 0 '''
        self.mat_stream.seek(126)
        mi = self.mat_stream.read(2)
        self.mat_stream.seek(0)
        return mi == 'IM' and '<' or '>'

    def file_header(self):
        ''' Read in mat 5 file header '''
        hdict = {}
        hdr = self.read_dtype(self.dtypes['file_header'])
        hdict['__header__'] = hdr['description'].item().strip(' \t\n\000')
        v_major = hdr['version'] >> 8
        v_minor = hdr['version'] & 0xFF
        hdict['__version__'] = '%d.%d' % (v_major, v_minor)
        return hdict


class Mat5MatrixWriter(MatStreamWriter):
    ''' Generic matlab matrix writing class '''
    mat_tag = np.zeros((), mdtypes_template['tag_full'])
    mat_tag['mdtype'] = miMATRIX
    default_mclass = None # default class for header writing
    def __init__(self,
                 file_stream,
                 arr,
                 name,
                 is_global=False,
                 unicode_strings=False,
                 long_field_names=False,
                 oned_as='column'):
        super(Mat5MatrixWriter, self).__init__(file_stream, 
                                               arr, 
                                               name,
                                               oned_as)
        self.is_global = is_global
        self.unicode_strings = unicode_strings
        self.long_field_names = long_field_names
        self.oned_as = oned_as

    def write_dtype(self, arr):
        self.file_stream.write(arr.tostring())

    def write_element(self, arr, mdtype=None):
        ''' write tag and data '''
        if mdtype is None:
            mdtype = np_to_mtypes[arr.dtype.str[1:]]
        byte_count = arr.size*arr.itemsize
        if byte_count <= 4:
            self.write_smalldata_element(arr, mdtype, byte_count)
        else:
            self.write_regular_element(arr, mdtype, byte_count)

    def write_smalldata_element(self, arr, mdtype, byte_count):
        # write tag with embedded data
        tag = np.zeros((), mdtypes_template['tag_smalldata'])
        tag['byte_count_mdtype'] = (byte_count << 16) + mdtype
        # if arr.tostring is < 4, the element will be zero-padded as needed.
        tag['data'] = arr.tostring(order='F')
        self.write_dtype(tag)

    def write_regular_element(self, arr, mdtype, byte_count):
        # write tag, data
        tag = np.zeros((), mdtypes_template['tag_full'])
        tag['mdtype'] = mdtype
        tag['byte_count'] = byte_count
        padding = (8 - tag['byte_count']) % 8
        self.write_dtype(tag)
        self.write_bytes(arr)
        # pad to next 64-bit boundary
        self.write_bytes(np.zeros((padding,),'u1'))

    def write_header(self, mclass=None,
                     is_global=False,
                     is_complex=False,
                     is_logical=False,
                     nzmax=0,
                     shape=None):
        ''' Write header for given data options
        mclass      - mat5 matrix class
        is_global   - True if matrix is global
        is_complex  - True if matrix is complex
        is_logical  - True if matrix is logical
        nzmax        - max non zero elements for sparse arrays
        shape : {None, tuple} optional
            directly specify shape if this is not the same as for
            self.arr
        '''
        if mclass is None:
            mclass = self.default_mclass
        if shape is None:
            shape = matdims(self.arr, self.oned_as)
        self._mat_tag_pos = self.file_stream.tell()
        self.write_dtype(self.mat_tag)
        # write array flags (complex, global, logical, class, nzmax)
        af = np.zeros((), mdtypes_template['array_flags'])
        af['data_type'] = miUINT32
        af['byte_count'] = 8
        flags = is_complex << 3 | is_global << 2 | is_logical << 1
        af['flags_class'] = mclass | flags << 8
        af['nzmax'] = nzmax
        self.write_dtype(af)
        self.write_element(np.array(shape, dtype='i4'))
        # write name
        self.write_element(np.array([ord(c) for c in self.name], 'i1'))

    def update_matrix_tag(self):
        curr_pos = self.file_stream.tell()
        self.file_stream.seek(self._mat_tag_pos)
        self.mat_tag['byte_count'] = curr_pos - self._mat_tag_pos - 8
        self.write_dtype(self.mat_tag)
        self.file_stream.seek(curr_pos)

    def write(self):
        raise NotImplementedError

    def make_writer_getter(self):
        ''' Make writer getter for this stream '''
        return Mat5WriterGetter(self.unicode_strings,
                                self.long_field_names,
                                self.oned_as)


class Mat5NumericWriter(Mat5MatrixWriter):
    default_mclass = None # can be any numeric type
    def write(self):
        imagf = self.arr.dtype.kind == 'c'
        try:
            mclass = np_to_mxtypes[self.arr.dtype.str[1:]]
        except KeyError:
            if imagf:
                self.arr = self.arr.astype('c128')
            else:
                self.arr = self.arr.astype('f8')
            mclass = mxDOUBLE_CLASS
        self.write_header(mclass=mclass,is_complex=imagf)
        if imagf:
            self.write_element(self.arr.real)
            self.write_element(self.arr.imag)
        else:
            self.write_element(self.arr)
        self.update_matrix_tag()


class Mat5CharWriter(Mat5MatrixWriter):
    codec='ascii'
    default_mclass = mxCHAR_CLASS
    def write(self):
        self.arr_to_chars()
        # We have to write the shape directly, because we are going
        # recode the characters, and the resulting stream of chars
        # may have a different length
        shape = self.arr.shape
        self.write_header(shape=shape)
        # We need to do our own transpose (not using the normal
        # write routines that do this for us)
        arr = self.arr.T.copy()
        if self.arr.dtype.kind == 'U' and arr.size:
            # Recode unicode using self.codec
            n_chars = np.product(shape)
            st_arr = np.ndarray(shape=(),
                                dtype=self.arr_dtype_number(n_chars),
                                buffer=arr)
            st = st_arr.item().encode(self.codec)
            arr = np.ndarray(shape=(len(st),),
                             dtype='u1',
                             buffer=st)
        self.write_element(arr, mdtype=miUTF8)
        self.update_matrix_tag()


class Mat5UniCharWriter(Mat5CharWriter):
    codec='UTF8'


class Mat5SparseWriter(Mat5MatrixWriter):
    default_mclass = mxSPARSE_CLASS
    def write(self):
        ''' Sparse matrices are 2D
        '''
        A = self.arr.tocsc() # convert to sparse CSC format
        A.sort_indices()     # MATLAB expects sorted row indices
        is_complex = (A.dtype.kind == 'c')
        nz = A.nnz
        self.write_header(is_complex=is_complex,
                          nzmax=nz)
        self.write_element(A.indices.astype('i4'))
        self.write_element(A.indptr.astype('i4'))
        self.write_element(A.data.real)
        if is_complex:
            self.write_element(A.data.imag)
        self.update_matrix_tag()


class Mat5CellWriter(Mat5MatrixWriter):
    default_mclass = mxCELL_CLASS
    def write(self):
        self.write_header()
        self._write_items()

    def _write_items(self):
        # loop over data, column major
        A = np.atleast_2d(self.arr).flatten('F')
        MWG = self.make_writer_getter()
        for el in A:
            MW = MWG.matrix_writer_factory(self.file_stream, el)
            MW.write()
        self.update_matrix_tag()


class Mat5BinaryBlockWriter(Mat5MatrixWriter):
    ''' class to write untranslatable binary blocks '''
    def write(self):
        # check endian
        # write binary block as is
        pass

class Mat5StructWriter(Mat5CellWriter):
    ''' class to write matlab structs

    Differs from cell writing class in writing field names,
    and in mx class
    '''
    default_mclass = mxSTRUCT_CLASS

    def _write_items(self):
        # write fieldnames
        fieldnames = [f[0] for f in self.arr.dtype.descr]
        length = max([len(fieldname) for fieldname in fieldnames])+1
        max_length = (self.long_field_names and 64) or 32
        if length > max_length:
            raise ValueError(
                "Field names are restricted to %d characters"
                 % (max_length-1))
        self.write_element(np.array([length], dtype='i4'))
        self.write_element(
            np.array(fieldnames, dtype='S%d'%(length)),
            mdtype=miINT8)
        A = np.atleast_2d(self.arr).flatten('F')
        MWG = self.make_writer_getter()
        for el in A:
            for f in fieldnames:
                MW = MWG.matrix_writer_factory(self.file_stream, el[f])
                MW.write()
        self.update_matrix_tag()


class Mat5ObjectWriter(Mat5StructWriter):
    ''' class to write matlab objects

    Same as writing structs, except different mx class, and extra
    classname element after header
    '''
    default_mclass = mxOBJECT_CLASS
    def write(self):
        self.write_header()
        self.write_element(np.array(self.arr.classname, dtype='S'),
                           mdtype=miINT8)
        self._write_items()


class Mat5WriterGetter(object):
    ''' Wraps options, provides methods for getting Writer objects '''
    @docfiller
    def __init__(self, 
                 unicode_strings=True, 
                 long_field_names=False,
                 oned_as='column'):
        ''' Initialize writer getter

        Parameters
        ----------
        unicode_strings : bool
           If True, write unicode strings
        %(long_fields)s
        %(oned_as)s
        '''
        self.unicode_strings = unicode_strings
        self.long_field_names = long_field_names
        self.oned_as = oned_as

    def to_writeable(self, source):
        ''' Convert input object ``source`` to something we can write

        Parameters
        ----------
        source : object

        Returns
        -------
        arr : ndarray

        Examples
        --------
        >>> mwg = Mat5WriterGetter()
        >>> mwg.to_writeable(np.array([1])) # pass through ndarrays
        array([1])
        >>> expected = np.array([(1, 2)], dtype=[('a', '|O8'), ('b', '|O8')])
        >>> np.all(mwg.to_writeable({'a':1,'b':2}) == expected)
        True
        >>> np.all(mwg.to_writeable({'a':1,'b':2, '_c':3}) == expected)
        True
        >>> np.all(mwg.to_writeable({'a':1,'b':2, 100:3}) == expected)
        True
        >>> np.all(mwg.to_writeable({'a':1,'b':2, '99':3}) == expected)
        True
        >>> class klass(object): pass
        >>> c = klass
        >>> c.a = 1
        >>> c.b = 2
        >>> np.all(mwg.to_writeable({'a':1,'b':2}) == expected)
        True
        >>> mwg.to_writeable([])
        array([], dtype=float64)
        >>> mwg.to_writeable(())
        array([], dtype=float64)
        >>> mwg.to_writeable(None)

        >>> mwg.to_writeable('a string').dtype
        dtype('|S8')
        >>> mwg.to_writeable(1)
        array(1)
        >>> mwg.to_writeable([1])
        array([1])
        >>> mwg.to_writeable([1])
        array([1])
        >>> mwg.to_writeable(object()) # not convertable

        dict keys with legal characters are convertible

        >>> mwg.to_writeable({'a':1})['a']
        array([1], dtype=object)

        but not with illegal characters

        >>> mwg.to_writeable({'1':1}) is None
        True
        >>> mwg.to_writeable({'_a':1}) is None
        True
        '''
        if isinstance(source, np.ndarray):
            return source
        if source is None:
            return None
        # Objects that have dicts
        if hasattr(source, '__dict__'):
            source = dict((key, value) for key, value in source.__dict__.items()
                          if not key.startswith('_'))
        # Mappings or object dicts
        if hasattr(source, 'keys'):
            dtype = []
            values = []
            for field, value in source.items():
                if (isinstance(field, basestring) and 
                    not field[0] in '_0123456789'):
                    dtype.append((field,object))
                    values.append(value)
            if dtype:
                return np.array( [tuple(values)] ,dtype)
            else:
                return None
        # Next try and convert to an array
        narr = np.asanyarray(source)
        if narr.dtype.type in (np.object, np.object_) and \
           narr.shape == () and narr == source:
            # No interesting conversion possible
            return None
        return narr

    def matrix_writer_factory(self, stream, arr, name='', is_global=False):
        ''' Factory function to return matrix writer given variable to write

        Parameters
        ----------
        stream : fileobj
            stream to write to
        arr : array-like
            array-like object to create writer for
        name : string
            name as it will appear in matlab workspace
            default is empty string
        is_global : {False, True} optional
            whether variable will be global on load into matlab

        Returns
        -------
        writer : matrix writer object
        '''
        # First check if these are sparse
        if scipy.sparse.issparse(arr):
            return Mat5SparseWriter(stream, arr, name, is_global)
        # Try to convert things that aren't arrays
        narr = self.to_writeable(arr)
        if narr is None:
            raise TypeError('Could not convert %s (type %s) to array'
                            % (arr, type(arr)))
        args = (stream,
                narr,
                name,
                is_global,
                self.unicode_strings,
                self.long_field_names,
                self.oned_as)
        if isinstance(narr, MatlabBinaryBlock):
            return Mat5BinaryBlockWriter(*args)
        if isinstance(narr, MatlabObject):
            return Mat5ObjectWriter(*args)
        if narr.dtype.fields: # struct array
            return Mat5StructWriter(*args)
        if narr.dtype.hasobject: # cell array
            return Mat5CellWriter(*args)
        if narr.dtype.kind in ('U', 'S'):
            if self.unicode_strings:
                return Mat5UniCharWriter(*args)
            else:
                return Mat5CharWriter(*args)
        else:
            return Mat5NumericWriter(*args)


class MatFile5Writer(MatFileWriter):
    ''' Class for writing mat5 files '''
    @docfiller
    def __init__(self, file_stream,
                 do_compression=False,
                 unicode_strings=False,
                 global_vars=None,
                 long_field_names=False,
                 oned_as=None):
        ''' Initialize writer for matlab 5 format files 

        Parameters
        ----------
        %(do_compression)s
        %(unicode_strings)s
        global_vars : None or sequence of strings, optional
            Names of variables to be marked as global for matlab
        %(long_fields)s
        %(oned_as)s
        '''
        super(MatFile5Writer, self).__init__(file_stream)
        self.do_compression = do_compression
        if global_vars:
            self.global_vars = global_vars
        else:
            self.global_vars = []
        # deal with deprecations
        if oned_as is None:
            warnings.warn("Using oned_as default value ('column')" +
                          " This will change to 'row' in future versions",
                          FutureWarning, stacklevel=2)
            oned_as = 'column'
        self.writer_getter = Mat5WriterGetter(
            unicode_strings,
            long_field_names,
            oned_as)
        # write header
        hdr =  np.zeros((), mdtypes_template['file_header'])
        hdr['description']='MATLAB 5.0 MAT-file Platform: %s, Created on: %s' \
            % (os.name,time.asctime())
        hdr['version']= 0x0100
        hdr['endian_test']=np.ndarray(shape=(),
                                      dtype='S2',
                                      buffer=np.uint16(0x4d49))
        file_stream.write(hdr.tostring())

    def get_unicode_strings(self):
        return self.writer_getter.unicode_strings
    def set_unicode_strings(self, unicode_strings):
        self.writer_getter.unicode_strings = unicode_strings
    unicode_strings = property(get_unicode_strings,
                               set_unicode_strings,
                               None,
                               'get/set unicode strings property')

    def get_long_field_names(self):
        return self.writer_getter.long_field_names
    def set_long_field_names(self, long_field_names):
        self.writer_getter.long_field_names = long_field_names
    long_field_names = property(get_long_field_names,
                                set_long_field_names,
                                None,
                                'enable writing 32-63 character field '
                                'names for Matlab 7.6+')

    def get_oned_as(self):
        return self.writer_getter.oned_as
    def set_oned_as(self, oned_as):
        self.writer_getter.oned_as = oned_as
    oned_as = property(get_oned_as,
                       set_oned_as,
                       None,
                       'get/set oned_as property')

    def put_variables(self, mdict):
        for name, var in mdict.items():
            if name[0] == '_':
                continue
            is_global = name in self.global_vars
            if self.do_compression:
                stream = StringIO()
                mat_writer = self.writer_getter.matrix_writer_factory(
                    stream,
                    var,
                    name,
                    is_global)
                mat_writer.write()
                out_str = zlib.compress(stream.getvalue())
                tag = np.empty((), mdtypes_template['tag_full'])
                tag['mdtype'] = miCOMPRESSED
                tag['byte_count'] = len(out_str)
                self.file_stream.write(tag.tostring() + out_str)
            else: # not compressing
                mat_writer = self.writer_getter.matrix_writer_factory(
                    self.file_stream,
                    var,
                    name,
                    is_global)
                mat_writer.write()