#!/usr/bin/env python
"""Log extractor for Rowe Technology DVL/ADCP ensemble files.

"""

from cStringIO import StringIO
import numpy as np
from scipy.io import loadmat, savemat

def export_ensembles_as_single_mat_file(infile, outfile=None, number_of_ensembles=9):
    """Read a Rowe DVL/ADCP ensemble file (.ENS) and write several ensembles
    out as separate MATLAB files (.MAT).
    
    """    
    with open(infile, 'rb') as f:
        raw = f.read() # reads entire file. There's probably a smarter option.
        delimiter = '\x80'*16 # look for first 16 bytes of header
        ensembles = raw.split(delimiter)[2:]
        payloads = [ensemble[16:-4] for ensemble in ensembles[:number_of_ensembles]]
        out = initialize_output_mat(payloads[0], len(payloads))
        for ensemble_number, payload in enumerate(payloads):
            try: 
                data = loadmat(StringIO(payload))
            except Exception as e:
                print 'Exception raised on ensemble number', ensemble_number
                # raise e
            for k, v in data.iteritems():
                try: 
                    out[k][ensemble_number] = v
                except: 
                    print "problem with key", k, "in ensemble", ensemble_number
        if outfile is not None:
            savemat(outfile, out, appendmat=True)
            print 'output saved as', outfile
        else:
            savemat(infile, out, appendmat=True)
            print 'output saved as {0}.mat'.format(infile)
            

def initialize_output_mat(payload, number_of_ensembles=9):
    """Initialize dictionary for mat file output.
    
    Use supplied ensemble payload as template.
    
    """
    out = dict()
    for k, v in loadmat(StringIO(payload)).iteritems():
        shape = [number_of_ensembles]
        shape.extend(list(v.shape)) # first axis is time (or ensemble number)
        out[k] = np.empty(shape, v.dtype)
    return out


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description='RTI ensemble extractor')
    parser.add_argument('-V', '--version', action='version', 
        version='%(prog)s 0.0.1',
        help='display version information and exit')
    parser.add_argument('-i', '--infile', metavar='filename', 
        type=str, help='input file (RTI .ENS)')
    parser.add_argument('-o', '--outfile', metavar='filename', default=None,
        type=str, help='output file (MATLAB .mat)')
    parser.add_argument('-n','--number-of-ensembles', default=9, 
        type=int, help='number of ensembles to split')
    args = parser.parse_args()
    export_ensembles_as_single_mat_file(**args.__dict__)