#!/usr/bin/python

# This file is part of snark, a generic and flexible library for robotics research
# Copyright (c) 2017 The University of Sydney
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
# 3. Neither the name of the University of Sydney nor the
#    names of its contributors may be used to endorse or promote products
#    derived from this software without specific prior written permission.
#
# NO EXPRESS OR IMPLIED LICENSES TO ANY PARTY'S PATENT RIGHTS ARE
# GRANTED BY THIS LICENSE.  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT
# HOLDERS AND CONTRIBUTORS \"AS IS\" AND ANY EXPRESS OR IMPLIED
# WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR
# BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
# OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN
# IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import os
import sys
import signal
import argparse
import tempfile
import json
import comma.csv.format

def import_advice( module ):
    debian_pkg = {
        "scipy.io": "python-scipy",
        "numpy": "python-numpy"
    }
    preamble = "Install"
    if debian_pkg.has_key( module ):
        print "On debian/ubuntu install with \"apt-get install %s\"" % debian_pkg[ module ]
        preamble = "On other systems install"
    print "%s with \"pip install %s\"" % ( preamble, module )

try:
    import numpy
    import scipy.io
except ImportError as e:
    print "ImportError:", e
    module = e.args[0].split(' ')[-1]
    import_advice( module )
    sys.exit( 1 )


def message( msg, data=() ):
    msg = "mat-convert: %s" % msg
    print >> sys.stderr, msg % data


def fatal( msg, data=() ):
    message( msg, data )
    sys.exit( 1 )


def bash_completion (parser):
    print " ".join(
        [ o for o in parser._option_string_actions.keys()
          if o != "--bash-completion" ])


def parse_args():
    description="""
convert to and from Matlab save files
"""

    epilog="""
output format:
  If the variable given is a leaf node then the data is output in csv-style
  ascii or binary. If it's a compound variable then it is output in json.

file format:
  Currently supports file format v7 and earlier, which is the default for all
  current versions of Matlab (including R2017a). Use the version option in the
  Matlab save command to force a version when saving. If v7.3 support is
  required then python library h5py can help, or possibly hdf5storage.

  If you're not seeing anything it's probable that the file has been saved with
  "opaque" Matlab variables. Try ">> mynewvar = struct( myvar )" from within
  Matlab, delete myvar, and save again.

examples:
  cat data.mat | {script_name}
  cat data.mat | {script_name} --variable foo.bar
  cat data.mat | {script_name} --list
  cat data.mat | {script_name} --list --path-separator /
  cat data.mat | {script_name} --variable foo --list
  cat data.mat | {script_name} --variable foo.bar --output-format
  cat data.mat | {script_name} --variable foo.bar --to csv --binary
  cat data.mat | {script_name} | name-value-convert --no-brackets
""".format( script_name=sys.argv[0].split('/')[-1] )

    fmt=lambda prog: argparse.RawDescriptionHelpFormatter( prog, max_help_position=50 )

    parser = argparse.ArgumentParser( description=description,
                                      epilog=epilog,
                                      formatter_class=fmt )

    parser.add_argument( "--bash-completion", action="store_true", help=argparse.SUPPRESS )
    parser.add_argument( "--binary", action="store_true", help="output in binary - only supported for csv output" )
    parser.add_argument( "--from", metavar="<format>", dest='from_arg', help="convert from <format> - only json supported" )
    parser.add_argument( "--list", action="store_true", help="list variables in given file" )
    parser.add_argument( "--output-format", action="store_true", help="show format of data" )
    parser.add_argument( "--path-separator", metavar="<char>", default=".", help="separator for input and output of variable paths (default=.)" )
    parser.add_argument( "--pretty-printing", action="store_true", help="pretty print json (adds indentation)" )
    parser.add_argument( "--to", metavar="<format>", help="convert to <format> - one of csv,json" )
    parser.add_argument( "--variable", metavar="<var>", help="variable to extract from save file (default: all)" )

    args = parser.parse_args()

    if args.bash_completion:
        bash_completion( parser )
        sys.exit( 0 )

    return args


def reset_signals():
    # Reset SIGPIPE and SIGINT to their default OS behaviour.
    # This stops python dumping a stack-trace on ctrl-c or broken pipe.
    signal.signal( signal.SIGPIPE, signal.SIG_DFL )
    s = signal.signal( signal.SIGINT, signal.SIG_DFL )
    # but don't reset SIGINT if it's been assigned to something other
    # than the Python default
    if s != signal.default_int_handler:
        signal.signal( signal.SIGINT, s )


def is_compound( variable ):
    try:
        comma.csv.format.from_numpy( variable.dtype )
        return False
    except( AttributeError, ValueError ):
        return True


def step_down( mat ):
    if mat.dtype == numpy.dtype('O'):    # sometimes there's extra nesting
        mat = mat.flat[0]
    if not isinstance( mat, numpy.ndarray ):
        raise ValueError
    return mat


# convert python version of matlab data to fundamental python data types
def traverse( mat, path="", sep="" ):
    listing = ( sep != "" )
    if is_compound( mat ):
        if isinstance(mat,numpy.ndarray):
            if mat.dtype.fields is None:
                return [ traverse(j, path, sep) for j in mat.tolist() ]
            elif mat.size!=1:
                ix=0
                ls=list()
                while ix < mat.size:
                    ls.append( traverse(mat.flat[ix], "%s[%d]" % (path,ix), sep) )
                    ix+=1
                return ls
        if isinstance(mat,unicode):
            if listing:
                print path
                return
            else:
                return mat;
        # at the top level, mat is a dictionary, but as we drill down
        # it turns into nested numpy.ndarray objects
        try:
            keys = ( [ i for i in mat.keys() if i[0] != "_" ]
                     if isinstance( mat, dict )
                     else mat.dtype.fields.keys() )
        except AttributeError:
            message( "warning: didn't understand %s", path )
            return None

        d = dict()
        for k in keys:
            next_item = mat[k]

            handled = False
            if next_item.dtype == numpy.dtype("O"):    # sometimes there's extra nesting
                if next_item.size == 1 or listing:
                    next_item = next_item.flat[0]
                else:
                    d[k] = list()
                    index = 0
                    for item in next_item.flat:
                        if listing:
                            traverse( item, "%s%s%s[%s]" % ( path, sep, k, index ), sep )
                            index += 1
                        else:
                            d[k].append( traverse( item ))
                    handled = True

            if not handled:
                if isinstance( next_item, scipy.io.matlab.mio5_params.MatlabOpaque ):
                    message( "warning: didn't understand %s%s%s", (path, sep, k ))
                    continue
                if listing:
                    traverse( next_item, k if path is None else "%s%s%s" % ( path, sep, k ), sep )
                else:
                    d[k] = traverse( next_item )
        return d
    else:
        if listing:
            print path
        else:
            return mat.tolist()


def to_numpy( input ):
    if isinstance( input, dict ):
        return { to_numpy( key ): to_numpy( value ) for key, value in input.iteritems() }
    elif isinstance( input, list ):
        ls=[ to_numpy( element ) for element in input ]
        if len( input ) > 1 and isinstance( input[0], dict ):
            dt = numpy.dtype([ ( k.encode( 'utf-8' ), 'O' ) for k in input[0].keys() ])
            # build ndarray with type input[0] keys
            nar = numpy.ndarray( shape=( 1, len( input )), dtype=dt )
            ix = 0
            while ix < len( input ):
                nar[ 0, ix ] = tuple( to_numpy( input[ix].values() ))
                ix += 1
            return nar
        return ls
    elif isinstance( input, unicode ):
        return input.encode( 'utf-8' )
    elif isinstance( input, int ):
        return ( float( input ))
    else:
        return input


def to_mat( args ):
    mat_data = json.load( sys.stdin )
    mat_data = to_numpy( mat_data )

    try:
        fd, tmp_path = tempfile.mkstemp( prefix="mat-convert-" )
        tmp_file = os.fdopen( fd, "w" )
    except OSError as e:
        fatal( "failed to open temporary file - %s", e )

    try:
        scipy.io.savemat( tmp_file, mat_data, appendmat=False )
    except IOError as e:
        fatal( "failed to write to temporary file %s - %s", ( tmp_path, e ))

    tmp_file = open( tmp_path, "r" )
    sys.stdout.write( tmp_file.read() )
    tmp_file.close()
    os.remove( tmp_path )


def from_mat( args ):
    fd, tmp_path = tempfile.mkstemp( prefix="mat-convert-" )
    try:
        tmp_file = os.fdopen( fd, "w" )
        tmp_file.write( sys.stdin.read() )
        tmp_file.close()
    except IOError as e:
        fatal( "failed to write to temporary file %s - %s", ( tmp_path, e ))

    try:
        mat = scipy.io.loadmat( tmp_path )
        os.remove( tmp_path )
    except IOError as e:
        fatal( "input must be redirected from a file" )
    except NotImplementedError as e:
        fatal( "support for Matlab file format v7.3 is not yet implemented (see help)" )
    except IndexError as e:
        fatal( "invalid or unsupported matlab file format" )

    if args.variable:
        for v in args.variable.split( args.path_separator ):
            try:
                mat = step_down( mat[ v ] )
            except ( KeyError, ValueError ):
                fatal( "variable \"%s\" does not exist", args.variable )

    if args.list:
        traverse( mat, args.variable, args.path_separator )
        return

    compound_variable = is_compound( mat )

    if args.output_format:
        if compound_variable:
            if args.variable:
                message( "%s is not a data array", args.variable )
            else:
                message( "--output-format requires a non-compound variable" )
        else:
            print "%s%s" % ( mat.shape[1], comma.csv.format.from_numpy( mat.dtype ))
        return

    if not args.to is None:
        if args.to == "csv":
            if compound_variable:
                message( "only non-compound variables can be output in csv format" )
            elif len( mat.shape ) > 2:
                message( '"--to csv" only supported for 2d matrices, %s is %sd' % ( args.variable, len( mat.shape )))
            else:
                if args.binary:
                    mat.tofile( sys.stdout )
                else:
                    for t in mat:
                        t.tofile( sys.stdout, ",", "%s" )
                        print
        elif args.to == "json":
            if compound_variable:
                dump_var = mat
            else:
                d = dict()
                leaf_var_name = args.variable.split( args.path_separator )[ -1 ]
                d[ leaf_var_name ] = mat
                dump_var = d
            if args.pretty_printing:
                print json.dumps( traverse( dump_var ), indent=4, separators=(",", ": "))
            else:
                print json.dumps( traverse( dump_var ))
        else:
            fatal( "--to <format> must be one of csv or json" )
    else:
        fatal( "must have one of --to, --from, --list or --output-format" )


def main():
    reset_signals()
    args = parse_args()

    if not args.from_arg is None:
        to_mat( args )
    else:
        from_mat( args )


if __name__ == "__main__":
    main()
