#Module for reading and writing MNE .stc files

#This module developed by Padraig Kitterick
# Last edit 30 October 2008

#!/usr/bin/env python
import struct
import numpy
import sys

def read_stc(filename):
    """Reads in an .stc file.
    
        v,data,start_ms,sample_ms = read_stc(filename)
    
        v           List of vertex IDs relative to original surface
        data        (vertices x samples) data matrix
        start_ms    Onset time of data in ms (e.g. -200)
        sample_ms   Length of each sample in ms (e.g. 1.475)
    """
    fp = open(filename,'rb')
    
    floatsize = struct.calcsize('>f')
    uintsize = struct.calcsize('>I')
    
    # Read time info
    start_ms = struct.unpack('>f',fp.read(floatsize))[0]
    sample_ms = struct.unpack('>f',fp.read(floatsize))[0]
    
    # Next, the vertex count and their vertex numbers 
    vnum = struct.unpack('>I',fp.read(uintsize))[0]
    
    # Set up output arrays
    v = numpy.zeros((vnum),numpy.int32())
    v[:] = struct.unpack(">%s" % ('I' * vnum),fp.read(uintsize*vnum))
    
    # Read the number samples, and the sample data
    numsamples = struct.unpack('>I',fp.read(uintsize))[0]
    
    datapoints = vnum*numsamples
    data = numpy.zeros((datapoints,1))
    data[:,0] = struct.unpack(">%s" % ('f' * datapoints),fp.read(floatsize*datapoints))
    
    # Reshape the data into the right format
    data = data.reshape(vnum,numsamples,order='F')
    
    fp.close()

    return v,data,start_ms,sample_ms

def write_stc(filename,vertices,data,start_ms,sample_ms):
    """Writes out an .stc file.
    
        write_stc(filename,vertices,data,start_ms,sample_ms)
    
        vertices    List of vertex IDs relative to original surface
        data        (vertices x samples) data matrix
        start_ms    Onset time of data in ms (e.g. -200)
        sample_ms   Length of each sample in ms (e.g. 1.475)
    """
    fp = open(filename,'wb')
    # Write out time info
    fp.write(struct.pack('>ff',start_ms,sample_ms))

    # Next, the vertex count and their vertex numbers 
    fp.write(struct.pack('>I',len(vertices)))

    for i in range(len(vertices)):
        fp.write(struct.pack('>I',vertices[i]))

    # Write out the number samples, and the sample data
    data = data.astype(numpy.float32)
    numvert,numsamples = data.shape
    fp.write(struct.pack('>I',numsamples))

    for ti in range(numsamples):
        vcolumn = data[:,ti]
        
        if sys.byteorder == 'big':
            vstring = vcolumn.tostring()
        else:
            # byteswap for big endian-ness
            vstring = vcolumn.byteswap().tostring()
        fp.write(vstring)

    fp.close()