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

"""

import io
import struct
import collections
import numpy as np
from scipy.io import loadmat, savemat
import crcmod

xmodem = crcmod.predefined.mkPredefinedCrcFun('xmodem')

def seek_to_next_preamble(stream, verbosity = 0, preamble_char=b'\x80', preamble_len = 16):
    """Find the next complete ensemble.
    """
    q = collections.deque(maxlen=preamble_len)
    while q.count(preamble_char) < preamble_len:
        q.append(stream.read(1))
        if verbosity > 2: print('Seeking end of preamble at position {0}.'.format(stream.tell()))
    if verbosity > 1: print('Found end of preamble at position {0}.'.format(stream.tell()))
    stream.seek(- preamble_len, 1) # move back to the start of the preamble 
    return stream.tell()


def read_rti_header(stream, verbosity=0):
    """Read the header.
    """
    header = stream.read(32)
    if header.find(b'\x80' * 16) is 0:
        i = struct.unpack(8*'i', header)
        ensemble_number = i[4]
        payload_size = i[6]
        if verbosity > 0: print('Found ensemble number {0} starting at {1} with {2}-byte payload.'.format(ensemble_number, stream.tell()-32, payload_size))
    else:
        ensemble_number = None
        payload_size = None
        stream.seek(-31,1)
        if verbosity > 0: print('Did not find valid RTI header at position {0}, have: {1}'.format(stream.tell()-1, header))
    return ensemble_number, payload_size


def read_ensemble(stream, verbosity=0):
    
    seek_to_next_preamble(stream, verbosity)

    ensemble_number, payload_size = read_rti_header(stream, verbosity)
    if ensemble_number is None:
        return None

    payload = stream.read(payload_size)
    checksum = stream.read(4)
    try: 
        checksum_read = struct.unpack('HH', checksum)[0]
    except:
        stream.seek(0,2) # give up and skip to the end of the file
        return None
    checksum_calculated = xmodem(payload)
    if verbosity > 0: print('Read ensemble {0} with {1}-byte payload; received checksum: {2}, calculated checksum {3}'.format(ensemble_number, payload_size, checksum_read, checksum_calculated))
    if checksum_read != checksum_calculated: 
        payload = None
    return payload

    
def export_ensembles_as_single_mat_file(infile, outfile=None, number_of_ensembles=None, verbosity=0):
    """Read a Rowe DVL/ADCP ensemble file (.ENS) and write ensembles
    out as a single MATLAB file (.MAT).
    
    """
    out = None
    with open(infile, 'rb') as f:
        total_size = f.seek(0,2)
        f.seek(0)
        while out is None and f.tell() < total_size:
            payload = read_ensemble(f, verbosity)
            if payload is not None:
                ensemble_size = 32 + len(payload) + 4
                if number_of_ensembles is None: # then guess the number of remaining ensembles
                    number_of_ensembles = int(np.floor((total_size - f.tell()) / ensemble_size)) + 1
                out = initialize_output_mat(payload, number_of_ensembles)

        f.seek(-ensemble_size, 1) # TODO: avoid backtracking

        for n in range(number_of_ensembles):
            payload = read_ensemble(f, verbosity)
            if payload is not None:
                for k, v in loadmat(io.BytesIO(payload)).items(): out[k][n] = v
            if f.tell() == total_size: break
            if verbosity > 1: print('Expecting next preamble to start at position {0}'.format(f.tell()))
        
        if out is None:
            print('Did not find any valid ensembles -- not writing output.')
        elif outfile is not None:
            savemat(outfile, out, appendmat=True)
            print('output saved as {0}'.format(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(io.BytesIO(payload)).items():
        shape = [number_of_ensembles]
        shape.extend(list(v.shape)) # first axis is time (or ensemble number)
        out[k] = np.empty(shape, v.dtype)
        out[k].fill(np.nan)
    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('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=None, 
        type=int, help='number of ensembles to process')
    parser.add_argument('-v','--verbosity', default=0, action='count', 
        help='increase output')
    args = parser.parse_args()
    export_ensembles_as_single_mat_file(**args.__dict__)