import tvm
import logging
import numpy as np
from collections.abc import Iterable
from .namespace import MSIR_COLLECTION,MSIR_NAME,MSIR_TARGET


def _get_logger():
    if not MSIR_COLLECTION.get(MSIR_NAME.LOGGER):
        MSIR_COLLECTION.set(MSIR_NAME.LOGGER, logging.getLogger("MSIR"))
    return MSIR_COLLECTION.get(MSIR_NAME.LOGGER)


def info(msg,level=0):
    env_level = MSIR_COLLECTION.get(MSIR_NAME.VERBOSE_LEVEL, 0)
    if level > env_level:
        _get_logger().info("[LV{}]{}".format(level,msg))


def warning(msg):
    _get_logger().warning(msg)


def debug(msg):
    _get_logger().debug(msg)


def check_type(obj,r_type):
    assert isinstance(obj,r_type), "Object {}({}) is not {}".format(obj,type(obj),r_type)


def check_iterable_type(obj,r_type,length=-1):
    assert isinstance(obj,Iterable), "object {}({}) is not iterable".format(obj,type(obj))
    assert all(isinstance(o,r_type) for o in obj),"Some of the object {} is not class of {}".format(obj,r_type)
    if length>0:
        assert len(obj)==length, "Object length {} mismatch with required {}".format(len(obj),length)


def get_version(target = MSIR_TARGET.MSIR):
    if target == MSIR_TARGET.TORCH:
        import torch
        version=torch.__version__
        if '+cu' in version:
            version=version.split('+cu')[0]
    elif target == MSIR_TARGET.TF:
        import tensorflow
        version = tensorflow.__version__
    else:
        raise Exception("Unexpected target " + str(target))
    return version


def _cast_array(array):
    if isinstance(array,tvm.nd.NDArray):
        return "tvm.ndarray", array.asnumpy()
    try:
        import torch
        if isinstance(array, torch.Tensor):
            return "torch.Tensor", array.detach().cpu().numpy()
    except:
        pass
    assert isinstance(array, np.ndarray), "Unexpected array type " + str(type(array))
    return "ndarray", array


def array_info(array):
    array_type, array = _cast_array(array)
    return "<{}>S:{}, D:{}, Max:{:g}, MIN:{:g}, SUM:{:g}".format(array_type, array.shape, array.dtype, array.max(), array.min(), array.sum())


def show_array(array, name="array"):
    print("{}:{}".format(name, array_info(array)))