#! /usr/bin/env python
# -*- coding: utf-8 -*-
'''
Created on 2022/8/31 20:14:34

@author: LiuKuan
@copyright: Apache License, Version 2.0
'''
import sys, types
import inspect
from pprint import pprint
# ===============================================================================
r'''提供一些调试方便的函数
'''
# ===============================================================================
__all__ = ['pprint', 'print_args','print_eval',
           'print_type', 'print_repr', 'print_signature',
           'print_numpy_array', 'print_torch_tensor'
           ]

def print_args(*args, **kws):
    print('args:', args)
    print('kws:', kws)

def print_eval(st):
    '''字符串打印并执行'''
    _frm = sys._getframe(1)
    val = eval(st, _frm.f_globals, _frm.f_locals)
    print( f"{st} = {val}")

def print_type(t):
    '''打印类别'''
    _type_name = None
    for tname in dir(types):
        if tname[0].isupper() and isinstance(t, getattr(types, tname)):
            _type_name = tname
            break
    if _type_name is None:
        _type_name = f"{type(t)}"
    print(_type_name)

def print_repr(v):
    '''打印 repr '''
    print(repr(v))

def print_signature(func):
    sig = inspect.signature(func)
    print(f'{func.__name__}{str(sig)}')

def print_numpy_array(data, title=None, show_data=True):
    '''print detail data information
    '''
    import numpy as np
    if title is None:
        msg = f"===="
    else:
        msg = f"====[{title}] "

    if isinstance(data, np.ndarray):
        _type = 'ndarray'
    else:
        _type = type(data)

    msg += f'type[{_type}]'
    if isinstance(data, np.ndarray):
        msg += f' shape[{data.shape}]  dtype[{data.dtype}]'

    msg += f' id[{id(data)}]\n'

    if show_data:
        msg += f'{data}'

    print(msg)
    return data

def print_torch_tensor(data, title=None, show_data=True):
    '''print detail data information
    '''
    import torch
    if title is None:
        msg = f"===="
    else:
        msg = f"====[{title}] "

    if isinstance(data, torch.Tensor):
        _type = 'Tensor'
    else:
        _type = type(data)

    def _fmt_size(sz):
        return '[' + ", ".join(str(i) for i in sz) + ']'

    def _fmt_dtype(dt):
        return str(dt).replace('torch.', '')

    msg += f'type[{_type}]'
    if isinstance(data, torch.Tensor):
        msg += f' shape{_fmt_size(data.shape)} dim[{data.ndim}] dtype[{_fmt_dtype(data.dtype)}] device[{data.device}]'

    msg += f' id[{id(data)}]\n'

    if show_data:
        msg += f'{data}'

    print(msg)
    return data

