#!/usr/bin/python

# This file is part of snark, a generic and flexible library
# Copyright (c) 2011 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.

__author__ = 'z.xu, v.vlaskine'

import argparse
import numpy as np
import os
import signal
import sys

description = """

read tensors as binary on stdin, load model from file, apply model, write output tensors to stdout

"""

epilog = """
examples
    main use cases
        each record contains only tensor
            cat input-tensors.bin | caffe-cat --graph graph.proto --weights weights.proto > output-tensors.bin
            
        each record in input data has header followed by tensor and then footer (e.g. containing timestamp, sequence number, or alike)
            cat header.input-tensor.footer.bin | caffe-cat --header-size 24 --footer-size 16 --graph graph.proto --weights weights.caffemodel > output-tensor.bin
            
    mnist example (you will need to download mnist dataset and train lenet on mnist, see http://caffe.berkeleyvision.org/gathered/examples/mnist.html)
        cat t10k-images-idx3-ubyte.gz | gunzip | tail -c+17 | head -c$(( 28 * 28 * 64 )) \\
            | cv-cat --input "rows=28;cols=28;type=ub;no-header" "convert-to=f,0.0039" \\
            | caffe-cat --model lenet_iter_10000.caffemodel --net lenet.prototxt \\
            | csv-from-bin 10f | csv-paste line-number -

"""

def warn( m ): print >> sys.stderr, "caffe-cat: warning:", m

def die( m ): print >> sys.stderr, "caffe-cat:", m; sys.exit( 1 )

def type_from_string( s ): # quick and dirty
    if s == 'b': return np.int8
    if s == 'ub': return np.uint8
    if s == 'w': return np.int16
    if s == 'uw': return np.uint16
    if s == 'i': return np.int32
    if s == 'ui': return np.uint32
    if s == 'l': return np.int64
    if s == 'ul': return np.uint64
    if s == 'f': return np.float32
    if s == 'd': return np.float64
    die( "expected type, got: '" + s + "'" )

def iterator( file, shape, ntype, header_size, footer_size ): # todo: quick and dirty; failed to find how to read from stdin using tf.FixedLengthRecordReader
    size = np.prod( shape )
    while True:
        header = np.fromfile( file, np.uint8, header_size )
        data = np.fromfile( file, ntype, size )
        footer = np.fromfile( file, np.uint8, footer_size )
        if data.shape[0] == size: yield header, data.reshape( shape ), footer
        else: yield None, None, None

if __name__ == '__main__':
    def handle_signal( s, f ): print >> sys.stderr, "caffe-cat: broken pipe, exit"; sys.exit( 0 )
    signal.signal( signal.SIGPIPE, handle_signal ) 
    parser = argparse.ArgumentParser( description = description, epilog = epilog, formatter_class = argparse.RawTextHelpFormatter )
    parser.add_argument( '--blob', type = str, help = 'helper parameter, see --blob-shape, --blob-type, etc' )
    parser.add_argument( '--blob-shape', action = "store_true", help = "print blob shape and exit" )
    parser.add_argument( '--blob-type', action = "store_true", help = "print blob element type and exit" )
    parser.add_argument( '--cpu', action = "store_true", help = 'run on cpu only (typically, caffe needs to be build for cpu only to make it work)' )
    parser.add_argument( '--device', '--gpu-device', '--gpu', type = int, default = 0, help = 'if in gpu mode, device id; default: %(default)s' )
    parser.add_argument( '--footer-size', type = int, default = 0, help = 'if a footer of given size is present, pass it to stdout with output tensor prepended' )
    parser.add_argument( '--header-size', type = int, default = 0, help = 'if a header of given size is present, pass it to stdout with output tensor appended' )
    parser.add_argument( '--input', '-i', type = str, help = 'input layer name; default: input name from network' )
    parser.add_argument( '--keep-input', action = "store_true", help = 'append output to input' )
    parser.add_argument( '--log-level', type = str, default = 'warning', help = 'log level: debug, info, warning, error; default: %(default)s' )
    parser.add_argument( '--model', '--weights', type = str, help = 'model filename, e.g. --model weights.caffemodel' )
    parser.add_argument( '--net', type = str, help = 'network filename, e.g. --net net.protobuf' )
    parser.add_argument( '--output', '-o', nargs='+', type = str, help = 'output layer name; multiple space- or comma-separated outputs accepted; if not present and network has only one output, use that output' )
    parser.add_argument( '--output-index', '--output-indices', '--index', '--indices', type = str, help = 'output only specific index for the given output blob, e.g. --index=0,12; a convenience/debugging option' )
    parser.add_argument( '--verbose', '-v', action = "store_true", help = 'more output' )
    args = parser.parse_args()
    outputs = args.output
    if not args.output is None and len( args.output ) == 1: outputs = args.output[0].split( ',' )
    log_levels = { "debug": "0", "info": "1", "warning": "2", "error": "3" }
    os.environ["GLOG_minloglevel"] = log_levels[ args.log_level ]
    import caffe
    major, minor = map( int, caffe.__version__.split( '-' )[0].split( '.' ) )[0:2]
    if args.net is None: die( "please specify --net" )
    if args.blob_shape or args.blob_type:
        net = caffe.Net( args.net, caffe.TEST )
        if args.blob_shape: print net.blobs[ args.blob ].data.shape
        elif args.blob_type: print net.blobs[ args.blob ].data.dtype
        sys.exit( 0 )
    if args.cpu: caffe.set_mode_cpu()
    else: caffe.set_mode_gpu(); caffe.set_device( args.device )
    if args.model is None: die( "please specify --model" )
    net = caffe.Net( args.net, caffe.TEST, weights = args.model ) if major >= 1 else caffe.Net( args.net, args.model, caffe.TEST )
    input_name = net.inputs[0] if args.input is None else args.input
    said_once = False
    indices = None
    if not args.output_index is None:
        if not outputs is None and len( outputs ) > 1: die( "--output-index: not implemented for multiple outputs" )
        indices = tuple( map( int, args.output_index.split( ',' ) ) )
    for header, data, footer in iterator( sys.stdin, net.blobs[ input_name ].data.shape, net.blobs[ input_name ].data.dtype, args.header_size, args.footer_size ):
        if data is None: break
        value = net.forward_all( **{ input_name: data } ) # quick and dirty, should we read input directly into net.blobs[ input_name ].data?
        header.tofile( sys.stdout )
        if args.keep_input: data.tofile( sys.stdout )
        footer.tofile( sys.stdout )
        if outputs is None:
            if len( value ) > 1: die( "--output not specified; expected 1 output, got", len( value ), "outputs" )
            for key in value:
                if not said_once: print >> sys.stderr, "caffe-cat: no output specified; outputting '" + str( key ) + "'"; said_once = True
                if indices is None: value[key].tofile( sys.stdout )
                else: value[key][indices].tofile( sys.stdout ) # quick and dirty 
                break;
        else:
            for output in outputs:
                if indices is None: net.blobs[output].data.tofile( sys.stdout )
                else: net.blobs[output].data[indices].tofile( sys.stdout ) # quick and dirty
        sys.stdout.flush()
