""" 
sewrite provides a class 'startfile' for writing SE files. This class
is built on the script se.py, and extends some of the features. This
class can write header attributes, data columns, collumn attrabutes and
tables.

FEB2014: creation, Christian Ritter, Luke Siemens

"""

import numpy
import se
import tables
from itertools import izip as izip

class startfile():
    """ 
    startfile provides some methods for writing SE file.

    Methods:
        write_hattr -- write header attributes
        write_table -- write a data table
        write_cattr -- write column attributes to a cycle
        write_dcol -- wtire column data to a cycle

    Example:
    >>> import sewrite
    >>> cyc = 10
    >>> f=sewrite.startfile('file_name')
    >>> hattr_name = ['rotini', 'zini']
    >>> hattr_data = [0.1, 0.2]
    >>> f.write_hattr(hattr_name,hattr_data)
    >>> A_data = [1,2,3]
    >>> f.write_table('A', A_data)
    >>> dcol_name = ['mass']
    >>> dcol_data = [[0.3]]
    >>> f.write_dcol(cyc,dcol_name,dcol_data)
    >>> cattr_name = ['age']
    >>> cattr_data = [0.4]
    >>> f.write_cattr(cyc,cattr_name,cattr_data)

    """
    output_file = ""

    def __init__(self, file_name):
        """ 
        Initalizes the class with the name of the file to create.

        Input:
            file_name -- name of the file to be created

        """
        self.output_file = file_name

    def write_hattr(self, hattr_name, hattr_val):
        """ 
        writes the header attributes of the file.

        Input:
            hattr_name -- a list containing the names of the attributes
            hattr_val -- a list containing the values associated with
                         the header attribute

        The length of hattr_name must be the same as the length of
        hattr_val, if it is not an error will be reported and the
        function will exit..

        """
        if len(hattr_name) != len(hattr_val):
            print "Error: Write_hattr, list of names and values do not match."
        else:
            for i in xrange(len(hattr_name)):
                se.writeattr(self.output_file,-1,hattr_name[i],hattr_val[i])

    def write_table(self, table_name, table_val):
        """ 
        Writes a table to the SE file.

        Input:
            table_name -- the name of the table
            table_val -- an array or list of the table data

        Exceptions:
            NodeError -- if the SE file already has a node with the
                         name table_name an exception will be raised

        """
        if not (hasattr(table_val,"shape") and hasattr(table_val,"dtype")):
            table_val = numpy.array(table_val)

        try:
            sefile = tables.openFile(self.output_file, mode = "r+")
        except IOError:
            sefile = tables.openFile(self.output_file, mode = "a")
            sefile.setNodeAttr("/", "SE_version", "1.2")

        table_val_type = tables.Float64Col(shape=table_val.shape[1:]) #table_val_type is float by default
        if (str(table_val.dtype) == "int32") or (str(table_val.dtype) == "int64"):
            table_val_type = tables.Int64Col(shape=table_val.shape[1:])

        description = {"data": table_val_type}
        table = sefile.createTable('/', table_name, description,
                                filters=tables.Filters(complevel = 6))

        row = table.row
        for i in xrange(len(table_val)):
            row["data"] = table_val[i]
            row.append()
        sefile.close()

    def write_cattr(self, cyc, cattr_name, cattr_val):
        """ 
        Writes column attributes for a cycle.

        Input:
            cyc -- the number of cycle to write
            cattr_name -- a list containing the names of the column
                          attributes
            cattr_val -- a list contining the values associated with
                         the column attribute

       	The length of cattr_name must be the same as the length of
       	cattr_val, if it is not an error will be reported and the
       	function will exit.

        The column attributes should be writen after the column data.

        """
        cattr_list = []
        if len(cattr_name) != len(cattr_val):
            print "Error: Write_cattr, list of names and value list do not match."
        else:
            for name, val in izip(cattr_name, cattr_val):
                if name == "age":
                    val = numpy.array([val])#h5T.py h5File() requires that age
                                            #is a numpy array
                se.writeattr(self.output_file, int(cyc), name, val)

    def write_dcol(self, cyc, dcol_name, dcol_val):
        """ 
        Write the data for a cycle.

        Input:
            cyc -- the number of cycle to write
            dcol_name -- a list containing the names of the data to 
                         write
            dcol_val -- a list contining the lists or arrays of data

        Exceptions:
            NodeError -- if the SE file already has a node with a name
                         in dcol_name an exception will be raised

       	The length of dcol_name must be the same as the length of
       	dcol_val, if it is not an error will be reported and the
       	function will exit.

        """
        dcol_list = []
        if len(dcol_name) != len(dcol_val):
            print "Error: Write_dcol, list of names and value list do not match."
        else:
            for name, val_list in izip(dcol_name, dcol_val):
                if not (hasattr(val_list,"shape") and hasattr(val_list,"dtype")):
                    val_list = numpy.array(val_list)

                val_type = tables.Float64Col(shape=val_list.shape[1:]) #val_type is float by default
                if (str(val_list.dtype) == "int32") or (str(val_list.dtype) == "int64"):
                    val_type = tables.Int64Col(shape=val_list.shape[1:])

                dcol_list.append((name,val_type,val_list))#list of tuples
            se.write(self.output_file, int(cyc), dcol_list)