#!/usr/bin/python

import os
import sys
import binascii
import types
import numpy
import scipy.io.mio
import re
import getopt


from lcm import EventLog

class LCMTypeDatabase:
    def __init__(self, package_names,verbose):
        for p in lcm_packages:
            try:
                __import__(p)
            except:
                if verbose:
                    sys.stderr.write("couldn't load module %s\n" % p)
                    lcm_packages.remove(p)

        self.klasses = {}
        for pkg in [ sys.modules[n] for n in package_names ]:
            for mname in dir(pkg):
                module = getattr(pkg, mname)
                #print mname
                if type(module) != types.ModuleType:
                    continue
                klassname = module.__name__.split(".")[-1]
                klass = getattr(module, klassname)
                self.klasses[klass._get_packed_fingerprint()] = klass

    def find_type(self, packed_fingerprint):
        return self.klasses.get(packed_fingerprint,None)


flatteners = {}
data = {}

def make_simple_accessor(fieldname):
    return lambda lst, x: lst.append(getattr(x, fieldname))

def make_numpy_array_accessor(fieldname):
    return lambda lst, x: lst.extend(numpy.array(getattr(x, fieldname)).ravel())

def make_obj_accessor(fieldname, func):
    return lambda lst, x: func(lst, getattr(x, fieldname))

def make_obj_list_accessor(fieldname, func):
    return lambda lst, x: lst.extend(map(lambda item: func(lst, item), getattr(x, fieldname)))

def make_lcmtype_accessor(msg):
    funcs = []

    for fieldname in getattr(msg, '__slots__'):
        m = getattr(msg, fieldname)

        if type(m) in [ types.IntType, types.LongType, types.FloatType,
                types.BooleanType ]:
            # scalar
            accessor = make_simple_accessor(fieldname)
            funcs.append(accessor)
        elif type(m) in [ types.ListType, types.TupleType ]:
            # convert to a numpy array
            arr = numpy.array(m)

            # check the data type of the array
            if arr.dtype.kind in "bif":
                # numeric data type
                funcs.append(make_numpy_array_accessor(fieldname))
            elif arr.dtype.kind == "O":
                # compound data type
                typeAccess = make_lcmtype_accessor(m[0])
                make_obj_list_accessor(fieldname, typeAccess)#untested
                #pass
        elif type(m) in types.StringTypes:
            # ignore strings
            pass
        else:
            funcs.append(make_obj_accessor(fieldname, make_lcmtype_accessor(m)))

    def flatten(lst, m):
        for func in funcs:
            func(lst, m)
    return flatten

def make_flattener(msg):
    accessor = make_lcmtype_accessor(msg)
    def flattener(m):
        result = []
        accessor(result, m)
        return result
    return flattener




def make_lcmtype_string(msg,base=True):
    typeStr = []
    count = 0
    for fieldname in getattr(msg, '__slots__'):
        m = getattr(msg, fieldname)

        if type(m) in [ types.IntType, types.LongType, types.FloatType,types.BooleanType ]:
            count = count+1
            if base:
                typeStr.append("%d- %s" % (count,fieldname))
            else:
                typeStr.append(fieldname)
        elif type(m) in [ types.ListType, types.TupleType ]:
            # convert to a numpy array
            arr = numpy.array(m)
            # check the data type of the array
            if arr.dtype.kind in "bif":
                # numeric data type
                
                if base:
                    typeStr.append("%d- %s(%d)" % (count+1,fieldname,len(arr.ravel())))
                else:
                    typeStr.append("%s(%d)" % (fieldname,len(arr.ravel())))
                count = count+len(arr.ravel())
            elif arr.dtype.kind == "O":
                # compound data type
                subStr,subCount = make_lcmtype_string(m[0],False)
                numSub = len(m)
                if base:
                    subStr = "%d- %s<%s>(%d)" % (count+1, fieldname, ", ".join(subStr), numSub)
                else:
                    subStr = "%s<%s>(%d)" % (fieldname, ", ".join(subStr), numSub)
                typeStr.append(subStr)
                count = count+numSub*subCount
                #pass
        elif type(m) in types.StringTypes:
            # ignore strings
            pass
        else:
            subStr, subCount = make_lcmtype_string(m,False);
            if base:
                for s in subStr:
                    count= count+1
                    typeStr.append("%d- %s.%s" %(count,fieldname , s))
            else:
                count= count+subCount
                for s in subStr:
                    typeStr.append(fieldname + "." + s)

    return typeStr, count

def deleteStatusMsg(statMsg):
    if statMsg:
        sys.stderr.write("\r")
        sys.stderr.write( " " * (len(statMsg)))
        sys.stderr.write("\r")
    return ""

longOpts = ["help", "print","format","separator","channels","ignore","outfile","lcm_packages","extract_jpgs"]
def usage():
    pname, sname = os.path.split(sys.argv[0])
    sys.stderr.write("usage: % s %s < filename > \n" %  (sname,str(longOpts)))
    print """
    -h --help                 print this message
    -p --print                Output log data to stdout instead of to .mat
    -f --format               print the data format to stderr
    -s --seperator=sep        print data with separator [sep] instead of default to ["" ""]
    -c --channels=chan        Parse channels that match Python regex [chan] defaults to [".*"]
    -i --ignore=chan          Ignore channels that match Python regex [chan]
                              both -c and -i results in undefined behavior!
    -o --outfile=ofname       output data to [ofname] instead of default [filename.mat or stdout]
    -l --lcmtype_pkgs=pkgs    load python modules from list of packages [pkgs] defaults to ["botlcm", "carmen3d"]
    -j --extract_jpgs         extract the jpg image data from log
    -v                        Verbose

    """


    sys.exit()


try:
    opts, args = getopt.gnu_getopt(sys.argv[1:], "hpvfs:c:i:o:l:j", longOpts)
except getopt.GetoptError, err:
    # print help information and exit:
    print str(err) # will print something like "option -a not recognized"
    usage()
if len(args)!=1:
    usage()
#default options
fname = args[0]
lcm_packages = [ "botlcm", "carmen3d"]
outFname = fname + ".mat"
printFname = "stdout"
printFile = sys.stdout
verbose = False
printOutput=False
printFormat=False
ignoreChannels=False
channels= ".*"
separator=' '
extract_jpgs=False
for o, a in opts:
    if o == "-v":
        verbose = True
    elif o in ("-h", "--help"):
        usage()
    elif o in ("-p", "--print"):
        printOutput = True
    elif o in ("-f", "--format"):
        printFormat = True
    elif o in ("-s", "--separator="):
        separator = a
    elif o in ("-o", "--outfile="):
        outFname = a
        printFname = a
    elif o in ("-c", "--channels="):
        channels = a
        ignoreChannels=False;
    elif o in ("-i", "--ignore="):
        channels = a
        ignoreChannels=True;
    elif o in ("-l", "--lcm_packages="):
        lcm_packages = a.split(",")
    elif o in ("-j", "--extract_jpgs="):
        extract_jpgs = True
    else:
        assert False, "unhandled option"
        
fullPathName = os.path.abspath(outFname)
dirname = os.path.dirname(fullPathName)
outBaseName = os.path.basename(outFname).split(".")[0]
fullBaseName = dirname + "/" + outBaseName

imageDir = fullBaseName + "_images"

type_db = LCMTypeDatabase(lcm_packages,verbose)

channels = re.compile(channels)
imageChannels = re.compile("IMAGE_.*")
log = EventLog(fname, "r")

if printOutput:
    sys.stderr.write("opened % s, printing output to %s \n" % (fname, printFname))
    if printFname == "stdout":
        printFile = sys.stdout
    else:
        printFile = open(printFname,"w")
else:
    sys.stderr.write("opened % s, outputing to % s\n" % (fname,outFname))

channels_to_ignore = []
msgCount = 0
statusMsg = ""
startTime =0
imageFnames = {}

for e in log:
    if msgCount==0: 
        startTime=e.timestamp
        
    if e.channel in channels_to_ignore:
        continue
    m = channels.match(e.channel)
    if (ignoreChannels and m) or (not ignoreChannels and not m):
        if verbose:
            sys.stderr.write("ignoring channel %s\n" % e.channel)
        channels_to_ignore.append(e.channel)
        continue

    lcmtype = type_db.find_type(e.data[:8])
    if not lcmtype:
        if verbose:
            sys.stderr.write("ignoring channel %s\n" % e.channel)
        channels_to_ignore.append(e.channel)
        continue
    isImage = imageChannels.match(e.channel)

    msgCount=msgCount+1
    if (msgCount % 5000) ==0:
        statusMsg = deleteStatusMsg(statusMsg)
        statusMsg = "read % d messages, % d %% done" % (msgCount,  log.tell()/float(log.size())*100)
        sys.stderr.write(statusMsg)
        sys.stderr.flush()

    msg = lcmtype.decode(e.data)

    if e.channel in flatteners:
        flattener = flatteners[e.channel]
    else:
        flattener = make_flattener(msg)
        flatteners[e.channel] = flattener
        data[e.channel]=[]
        if printFormat:
            statusMsg = deleteStatusMsg(statusMsg)
            sys.stderr.write( "\n%s  %s : \n" % (e.channel,lcmtype))
            typeStr,fieldCount = make_lcmtype_string(msg)
            typeStr.append("%d- log_timestamp" % (fieldCount+1))
            sc =0;
            sys.stderr.write("[\n")
            sys.stderr.write("\n".join(typeStr))
            sys.stderr.write("\n]\n")
        if isImage and extract_jpgs:
            imageFnames[e.channel] = []
            if not os.path.isdir(imageDir):
                os.mkdir(imageDir)
        


    a = flattener(msg)
    a.append((e.timestamp-startTime)/1e6)
    if printOutput:
        printFile.write("%s%s%s\n" % (e.channel,separator, separator.join([str(k) for k in a])))
    else:
        data[e.channel].append(a)
        
    if isImage and extract_jpgs:
        imFileName = "%s/%s_%07d.jpg"%(imageDir,e.channel,len(data[e.channel]))
        imageFnames[e.channel].append(imFileName)
        imFile = open(imFileName,"wb")
        imFile.write(msg.data)
        imFile.close()
    
        
    

deleteStatusMsg(statusMsg)
if not printOutput:
    sys.stderr.write("loaded all %d messages, saving to % s\n" % (msgCount,outFname))
    scipy.io.mio.savemat(outFname,data)
    
    mfile = open(dirname +"/"+ outBaseName + ".m","w")
    loadFunc =  """function [d imFnames]=%s()
load('%s');
s=whos('-file', '%s');
for i = 1:length(s)
    d.(s(i).name)=eval(s(i).name);
end
imFnames = struct();
""" % (outBaseName,fullPathName,fullPathName)

    for k, v in imageFnames.items():
        loadFunc += "imFnames.%s ={'%s'}; \n" %(k,"',\n'".join(v))
    
    
    mfile.write(loadFunc);
    mfile.close()
