# Copyright (c) 2012, Renato Florentino Garcia <fgarcia.renato@gmail.com>
#                     Stefano Pellegrini <stefpell@ee.ethz.ch>
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#     * Redistributions of source code must retain the above copyright
#       notice, this list of conditions and the following disclaimer.
#     * 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.
#     * Neither the name of the authors nor the
#       names of its contributors may be used to endorse or promote products
#       derived from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE AUTHORS 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 AUTHORS 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 sys
for i in range(len(sys.path) - 1, 0, -1):
    if(sys.path[i].find('/usr/lib/python2.7') == 0):
        del sys.path[i]
print(sys.version)

import gdb
import matplotlib
import matplotlib.pyplot as pl
import struct
import numpy as np
import math
import cv2

import traceback

def dumpToNumpy(shape, line_step, data_address, data_symbol):
    if len(shape) == 2:
        n_channel = 1
        width, height = shape
    if len(shape) == 3:
        # dealing with planar image
        if shape[0] <= 3:
            print('Loading planar image of dim {} {} {}'.format(shape[0], shape[1], shape[2]))
            height = shape[0] * shape[1]
            width = shape[2]
            n_channel = 1
        else:
            height, width, n_channel = shape

    line_step = int(line_step)
    data_address = int(data_address)
    infe = gdb.inferiors()
    memory_data = infe[0].read_memory(data_address, line_step * height)

    # Calculate the memory padding to change to the next image line.
    # Either due to memory alignment or a ROI.
    if data_symbol in ('b', 'B'):
        elem_size = 1
    elif data_symbol in ('h', 'H'):
        elem_size = 2
    elif data_symbol in ('i', 'f'):
        elem_size = 4
    elif data_symbol == 'd':
        elem_size = 8
    padding = line_step - width * n_channel * elem_size

    # Format memory data to load into the image.
    image_data = []
    if n_channel == 1:
        fmt = '%d%s%dx' % (width, data_symbol, padding)
        print('width: {} symbol: {} padding: {}'.format(width, data_symbol, padding))
        for line in chunker(memory_data, line_step):
            image_data.extend(struct.unpack(fmt, line))
    elif n_channel == 3:
        fmt = '%d%s%dx' % (width * 3, data_symbol, padding)
        for line in chunker(memory_data, line_step):
            image_data.extend(struct.unpack(fmt, line))
    elif n_channel == 2:
        fmt = '%d%s%dx' % (width * n_channel, data_symbol, padding)
        for line in chunker(memory_data, line_step):
            image_data.extend(struct.unpack(fmt, line))
    else:
        gdb.write('Only 1, 2, or 3 channels supported\n', gdb.STDERR)
        return

    scale_alpha = 1
    scale_beta  = 0
    # Fit the opencv elemente data in the PIL element data
    min_image_data = 0
    img_range = 255
    if data_symbol == 'b':
        image_data = [i+128 for i in image_data]
    elif data_symbol == 'H':
        image_data = [i>>8 for i in image_data]
    elif data_symbol == 'h':
        image_data = [(i+32768)>>8 for i in image_data]
    elif data_symbol == 'i':
        image_data = [(i+2147483648)>>24 for i in image_data]
    elif data_symbol in ('f','d'):
        # A float image is discretized in 256 bins for display.
        max_image_data = float(max(image_data))
        min_image_data = float(min(image_data))
        img_range = max_image_data - min_image_data
        print('Image max/min - range: %1.20f / %1.20f - %1.20f'%(max_image_data, min_image_data, img_range))

        if img_range > 0.00000000001:
            print('Attempting to rescale image data')
            scale_beta = min_image_data
            scale_alpha = img_range / 255.0
            for i in range(len(image_data)):
                if(math.isnan(image_data[i])):
                    image_data[i] = 0

            image_data = [int(255 * (i - min_image_data) / img_range) \
                            for i in image_data]
        else:
            image_data = [0 for i in image_data]
    return np.array(image_data).reshape(shape)

def dumpImage(transpose, width, height, n_channel, line_step, data_address, data_symbol, swap_bgr=True):
    """ Copies the image data to a PIL image and shows it.

    Args:
        width: The image width, in pixels.
        height: The image height, in pixels.
        n_channel: The number of channels in image.
        line_step: The offset to change to pixel (i+1, j) being
            in pixel (i, j), in bytes.
        data_address: The address of image data in memory.
        data_symbol: Python struct module code to the image data type.
    """

    # Show image.
    
    img = dumpToNumpy((int(height), int(width), int(n_channel)), line_step, data_address, data_symbol)
    img = img.astype('uint8')
    if n_channel == 1:
        if(transpose):
            img = img.transpose()
        bgr = img
        rgb = bgr
    else:
        bgr = np.copy(img)
        if(swap_bgr):
            # swap opencv BGR to RGB
            bgr = img
            rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
        else:
            # img is RGB
            rgb = img
            bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
    return bgr, rgb

def get_next_variable_name(var_name):
    print('Looking for delimiters in ', var_name)
    pos1 = var_name.find('.')
    pos2 = var_name.find('->')
    pos3 = var_name.find('[')
    if(pos1 == 0):
        pos1 = var_name[1:].find('.')
        if(pos1 != -1):
            pos1 = pos1 + 1
    if(pos2 == 0):
        pos2 = var_name[2:].find('->') + 2
    if(pos3 == 0):
        pos3 = var_name[1:].find('[') + 1

    if(pos1 == -1 and pos2 == -1 and pos3 == -1):
        print('Did not find any delimiters in ', var_name)
        return [var_name, '']
    if(pos1 == -1):
        pos1 = len(var_name) + 1
    if(pos2 == -1):
        pos2 = len(var_name) + 1
    if(pos3 == -1):
        pos3 = len(var_name) + 1

    if(pos1 < pos2 and pos1 < pos3):
        print('Delimiter \'.\' found at ', pos1)
        return [var_name[0:pos1], var_name[pos1:]]
    if(pos2 < pos1 and pos2 < pos3):
        print('Delimiter \'->\' found at ', pos2)
        return [var_name[0:pos2], var_name[pos2:]]
    if(pos3 < pos1 and pos3 < pos2):
        print('Delimiter \'[\' found at ', pos3)
        return [var_name[0:pos3], var_name[pos3:]]

def index_container(container, index):
    if(str(container.type).find('std::vector') == 0 or \
    str(container.type).find('const std::vector') == 0):
        idx_pos = index.find(']')
        idx = int(index[0:idx_pos])
        print('Indexing vector element {}'.format(idx))
        obj = (container['_M_impl']['_M_start'] + idx).dereference()
        return obj, index[idx_pos+1:]
    print('Unable to index container of type \'{}\''.format(str(container.type)))

def dumpCVMat(val, **kwargs):
    flags = val['flags']
    depth = flags & 7
    channels = 1 + (flags >> 3) & 63;
    if depth == 0:
        cv_type_name = 'CV_8U'
        data_symbol = 'B'
    elif depth == 1:
        cv_type_name = 'CV_8S'
        data_symbol = 'b'
    elif depth == 2:
        cv_type_name = 'CV_16U'
        data_symbol = 'H'
    elif depth == 3:
        cv_type_name = 'CV_16S'
        data_symbol = 'h'
    elif depth == 4:
        cv_type_name = 'CV_32S'
        data_symbol = 'i'
    elif depth == 5:
        cv_type_name = 'CV_32F'
        data_symbol = 'f'
    elif depth == 6:
        cv_type_name = 'CV_64F'
        data_symbol = 'd'
    else:
        gdb.write('Unsupported cv::Mat depth\n', gdb.STDERR)
        return

    rows = val['rows']
    cols = val['cols']

    line_step = val['step']['p'][0]

    gdb.write(cv_type_name + ' with ' + str(channels) + ' channels, ' +
                str(rows) + ' rows and ' +  str(cols) +' cols\n')

    data_address = val['data']

    info = (cols, rows, channels, line_step, data_address, data_symbol)
    return dumpImage(False, *info)

def dumpGpuMat(val, token):
    print('Attempting to download gpu mat {}'.format(token))
    cv_mat = gdb.parse_and_eval("download_gpu_mat({})".format(token))
    '''
        Place this somewhere in your code to be able to access gpu mats:
        cv::Mat download_gpu_mat(const cv::cuda::GpuMat& mat)
        {
            cv::Mat h_mat(mat);
            std::cout << "Debug downloading of gpu mat complete" << std::endl;
            return h_mat;
        }

    '''
    print('Got cv mat... {}'.format(cv_mat.type.strip_typedefs()))
    return dumpCVMat(cv_mat)

def dumpEigenMatrix(val, token):
    nrows = int(gdb.parse_and_eval('{}.rows()'.format(token)))
    ncols = int(gdb.parse_and_eval('{}.cols()'.format(token)))
    stride = int(gdb.parse_and_eval('{}.outerStride()'.format(token)))
    ptr = int(gdb.parse_and_eval('{}.data()'.format(token)))
    nchannel = 1
    print('Found {} of size {} x {} ptr: {} stride: {}'.format(val.dynamic_type, nrows, ncols, ptr, stride))

    type_string = str(val.type)
    dtype = 'd'
    if(type_string[-1] == 'f' or 'float' in type_string):
        dtype = 'f'

    img_info = (nrows, ncols, 1, stride * 8, ptr, dtype)
    return dumpImage(True, **img_info)

def dumpDarknetImage(val, **kwargs):
    height = val['h']
    width = val['w']
    channels = val['c']
    data_address = val['data']

    line_step = width * 4
    arr = dumpToNumpy((channels, height, width), line_step, data_address, 'f').astype('uint8')
    rgb = arr.transpose((1,2,0))

    return cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR), rgb

class VariableParser:
    def __init__(self, type_name):
        self.type = type_name

    def __call__(self, variable):
        var_type = str(variable.type)
        if var_type.find(self.type) == 0:
            return True, variable
        return False, None

data_parsers = {
    'cv::Mat': (VariableParser('cv::Mat'), dumpCVMat),
    'cv::cuda::GpuMat': (VariableParser('cv::Mat'), dumpGpuMat),
    'Eigen::Matrix': (VariableParser('cv::Mat'), dumpEigenMatrix),
    'image': (VariableParser('image'), dumpDarknetImage)
}

def getVarHelper(var_name, variable):
    for k, v in data_parsers.items():
        success, ret = v[0](variable)
        if success:
            return ret

    print('Mat not found in variable ({}), attempting to find delimiters'.format(str(variable.type)))
    # find the next deliminating element
    var_name, rest = get_next_variable_name(var_name)

    print('Split var name to ', var_name, ' and ', rest)
    if(len(var_name)):
        if(var_name[0] == '.'):
            variable = variable[var_name[1:]]
        elif(var_name[0:2] == '->'):
            variable = variable[var_name[2:]]
    if(len(rest) == 0):
        return variable
    if(rest[0] == '.'):
        return getVarHelper(rest[1:], variable)
    if(rest[0:2] == '->'):
        return getVarHelper(rest[2:], variable)
    if(rest[0] == '['):
        print('Array indexing ', rest)
        variable, rest = index_container(variable, rest[1:])
        return getVarHelper(rest, variable)


def getVar(var_name, frame):
    print('Looking for cv::Mat named ', var_name, ' in frame ', frame.name())
    tmp_var_name, rest = get_next_variable_name(var_name)
    print('Split var name to ', tmp_var_name, ' and ', rest)
    try_this = False
    try:
        var = frame.read_var(tmp_var_name)
    except:
        try_this = True
    if(try_this):
        var = frame.read_var('this')
        var = var[tmp_var_name]
    return getVarHelper(rest, var)


def chunker(seq, size):
    assert size > 0
    return (seq[pos:pos + size] for pos in range(0, len(seq), size))

class closeWindows(gdb.Command):
    def __init__(self):
        super(closeWindows, self).__init__('closeWindows', gdb.COMMAND_SUPPORT, gdb.COMPLETE_FILENAME)

    def invoke(self, arg, from_tty):
        pl.close('all')

def handle_container(container, index, i):
    if( str(container.type).find('std::vector') == 0):
        obj = (container['_M_impl']['_M_start'] + int(index[i])).dereference()
        return handle_container(obj, index, i+1)
    if(str(container.type).find('cv::Mat') == 0):
        return container


class printMat(gdb.Command):
    def __init__(self):
        super(printMat, self).__init__('printMat', gdb.COMMAND_SUPPORT, gdb.COMPLETE_FILENAME)
        

    def invoke(self, arg, from_tty):
        # Access the variable from gdb.
        frame = gdb.selected_frame()
        val = get_var(arg, frame)

        flags = val['flags']
        depth = flags & 7
        channels = 1 + (flags >> 3) & 63;
        rows = val['rows']
        cols = val['cols']
        line_step = val['step']['p'][0]
        data_address = val['data']
        infe = gdb.inferiors()
        memory_data = infe[0].read_memory(data_address, line_step * rows)
        print('Matrix with rows: ', rows, ' cols: ', cols, ' channels: ', channels)
        if depth == 0:
            dtype='uint8'
        elif depth == 1:
            dtype='int8'
        elif depth == 2:
            dtype='uint16'
        elif depth == 3:
            dtype='int16'
        elif depth == 4:
            dtype='int32'
        elif depth == 5:
            dtype='float32'
        elif depth == 6:
            dtype='float64'
        else:
            gdb.write('Unsupported cv::Mat depth\n', gdb.STDERR)
            return
        arr = np.frombuffer(memory_data, dtype=dtype).reshape((rows,cols,channels))
        print( arr )

def getValue(val):
    val_type = str(val.type.strip_typedefs())
    if '&' in val_type:
        return val.referenced_value()
    if '*' in val_type:
        return val.referenced_value()
    return val

class imshow(gdb.Command):
    """Diplays the content of an opencv image"""

    def __init__(self):
        super(imshow, self).__init__('imshow',
                                        gdb.COMMAND_SUPPORT,
                                        gdb.COMPLETE_FILENAME)
        self.transpose = False

    def invoke (self, arg, from_tty):
        print('invoking with args: {}'.format(arg))
        try:
            transpose=False
            pos = arg.find(' ')
            flag = ''
            save = None
            tokens = arg.split(' ')
            # Access the variable from gdb.
            frame = gdb.selected_frame()
            if('block' in tokens):
                flag = 'block'
                tokens.remove('block')
            if 'save' in tokens:
                idx = tokens.index('save')
                save = tokens[idx+1]
                tokens.remove('save')
                tokens.remove(save)

            current_line = str(gdb.decode_line()[1][0])
            line = current_line[current_line.find(', line ') + 6:]
            for i in range(len(tokens)):
                token = tokens[i]
                val = getVar(token, frame)
                val = getValue(val)
                img_type = str(val.type.strip_typedefs())
                gdb.write('Got desired object {}\n'.format(img_type))
                
                bgr = None
                rgb = None
                for k,v in data_parsers.items():
                    if k in img_type:
                        bgr, rgb = v[1](val=val, token=token)
                if rgb is not None:
                    title = token + line
                    if((i == len(tokens) - 1) and flag == 'block'):
                        self.show_image(title, rgb, 'block')
                    else:
                        self.show_image(title, rgb)
                    if save is not None:
                        cv2.imwrite(save, bgr)
        except ValueError as err:
            traceback.print_exc()
    
    def show_image(self, name, img, flag = ''):

        fig = pl.figure()
        fig.canvas.set_window_title(name)
        b = fig.add_subplot(111)
        if len(img.shape) == 2:
            n_channel = 1
        else:
            _, _, n_channel = img.shape

        if n_channel == 1:
            b.imshow(img, cmap = 'jet', interpolation='nearest')
        elif n_channel == 3:
            b.imshow(img)
        elif n_channel == 2:
            pl.gca().invert_yaxis()
            b.scatter(img[:,:,0], img[:,:,1], s=2,alpha=0.5)

        def format_coord(x, y):
            col = int(x+0.5)
            row = int(y+0.5)
            if col>=0 and col<width and row>=0 and row<height:
                if n_channel == 1:
                    #if(data_symbol == 'f'):
                    #    z = float_image[row,col]
                    #else:
                    #    z = float(float(img[row,col]) * scale_alpha) + float(scale_beta)
                    z = float(float(img[row,col]) * scale_alpha) + float(scale_beta)
                    #return '(%d, %d), [%1.2f]'%(col, row, z)
                    return '(x:{}, y:{}), [{}]'.format(col, row, z)
                elif n_channel == 3:
                    z0 = img[row,col,0] * scale_alpha + scale_beta
                    z1 = img[row,col,1] * scale_alpha + scale_beta
                    z2 = img[row,col,2] * scale_alpha + scale_beta
                    return '(%d, %d), [%1.2f, %1.2f, %1.2f]'%(col, row, z0, z1, z2)
            else:
                return 'x={} ({}), y={} ({})'.format(col,width, row, height)

        b.format_coord = format_coord
        if(flag == 'block'):
            pl.show(block=True)
        else:
            pl.show(block=False)

imshow()
closeWindows()
printMat()
