import numpy as np
import kneron_preprocessing
"""Get the affine transform matrix, given the center/scale/rot/output_size.

    Args:
        center (np.ndarray[2, ]): Center of the bounding box (x, y).
        scale (np.ndarray[2, ]): Scale of the bounding box
            wrt [width, height].
        output_size (array): output size of network input

    Returns:
        np.ndarray: The transform matrix.
    """
def get_affine_transform(center, scale, output_size):
    if not isinstance(scale, np.ndarray) and not isinstance(scale, list):
        scale = np.array([scale, scale])
    src_w = scale[0]
    dst_w = output_size[1]
    dst_h = output_size[0]
    src_dir =[0.0, src_w * -0.5]
    dst_dir = np.array([0, dst_w * -0.5], np.float32)
    src = np.zeros((3, 2), dtype=np.float32)
    dst = np.zeros((3, 2), dtype=np.float32)
    src[0, :] = center
    src[1, :] = center + src_dir
    dst[0, :] = [dst_w * 0.5, dst_h * 0.5]
    dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir
    src[2:, :] = get_3rd_point(src[0, :], src[1, :])
    dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :])
    trans = kneron_preprocessing.similarity_transform(np.float32(dst).flatten().tolist(),np.float32(src).flatten().tolist(),type='float')
    return trans
def get_3rd_point(a, b):
    direct = a - b
    return b + np.array([-direct[1], direct[0]], dtype=np.float32)
def _bbox_to_center_and_scale( bbox):
        x, y, w, h = bbox

        center = np.zeros(2, dtype=np.float32)
        center[0] = x + w / 2.0
        center[1] = y + h / 2.0

        scale = np.array([w , h],
                dtype=np.float32)
        return center, scale

'''
    return preprocess affine tranfor, 
    :param input_img: image path
    :param bbox: bounding box of size 4+ [x,y,w,h,.....]
    :param image_size: determines size of output
    :return: preprocessed input of size (1, image_size 0, image_size 1, 3)
'''
def postprocess_(input_img,bbox,image_size=[384,288],scale_ext=[1.09,1.135],**kwargs):
    aspect_ratio = image_size[1]/image_size[0]
    if isinstance(input_img, str):
        try:
            data_numpy = kneron_preprocessing.API.load_image(input_img)
        except:
            try:
                data_numpy = kneron_preprocessing.API.load_bin(input_img, **kwargs)
            except:
                print('input format error')
                assert 0
    else:
        assert isinstance(image, np.ndarray)
    if data_numpy is None:
        raise ValueError('=> fail to read {}'.format(image_file))
    center,scale = _bbox_to_center_and_scale(bbox)
    scale[0] *= scale_ext[0]
    scale[1] *= scale_ext[1]

    # fit the ratio
    if scale[0] > aspect_ratio * scale[1]:
        scale[1] = scale[0] * 1.0 / aspect_ratio
    else:
        scale[0] = scale[1] * 1.0 * aspect_ratio
    trans = get_affine_transform(center, scale, image_size)

    height = data_numpy.shape[0]
    width = data_numpy.shape[1]
    pre_info = {'centers':[center.tolist()],'scales':[scale.tolist()],'height':height,'width':width}
    img = kneron_preprocessing.warpAffine(data_numpy,Matrix=trans,warp_size= (int(image_size[1]), int(image_size[0]))) #reversed in affine for new one
    img = kneron_preprocessing.API.norm(img)
    img = img[np.newaxis,:,:,:]
    return img,pre_info