import sys
from typing import List

import numpy as np  # type: ignore

from code_loader.contract.datasetclasses import SectionCallableInterface, SubsetResponse


def to_numpy_return_wrapper(encoder_function: SectionCallableInterface) -> SectionCallableInterface:
    def numpy_encoder_function(idx: int, samples: SubsetResponse) -> np.ndarray:
        result = encoder_function(idx, samples)
        numpy_result = np.array(result)
        return numpy_result

    return numpy_encoder_function


def get_root_traceback(exc_tb):
    return_traceback = exc_tb
    while return_traceback.tb_next is not None:
        return_traceback = return_traceback.tb_next
    return return_traceback


def get_root_exception_line_number():
    exc_tb = sys.exc_info()[2]
    root_traceback = get_root_traceback(exc_tb)
    return root_traceback.tb_lineno


def get_shape(raw_result) -> List[int]:
    np_result = np.array(raw_result)
    np_shape = np_result.shape
    # fix single result shape viewing
    if np_shape is ():
        np_shape = (1,)
    shape = list(np_shape)
    return shape
