from . import io
from . import utils
from . import general_funcs as funcs
from ctypes import c_float
import numpy as np
import math
from PIL import Image

def inproc_720_flow(image,return_float=False,raw_fmt=None,raw_size=None ,npu_size=None, crop_box=None, pad_mode=None, resize_size=None, pad_length=None, matrix=None,bias=None, shift_1_bit=False, sub_128=False, radix=8, pad_limit_on=True):
    """
    inproc_720 python simulator

    Parameters
    ----------
    image : np.ndarray / string
        Image input, can be image file, bin or hex file

    return_float : bool
        If False, return RGBA fix format
        If True, return RGB float format
        default = False

    raw_fmt : string
        If image is raw file, set the raw fmt

    raw_size : tuple
        (src_w, src_h)

    npu_size : tuple
        (target_w, target_h)

    crop_box : tuple
        (x1, y1, x2, y2), will skip crop if set None

    pad_mode : int
        0: pad 2 sides, 1: pad 1 side, 2: no pad. default = 0

    matrix : list / tuple
        Matirx for the matrix operator, must be a list/tuple for 9 elements, default = [256,0,0,0,256,0,0,0,256]

    bias : list / tuple
        Bias for the matrix operator, must be a list/tuple for 3 elements, default = [0,0,0]

    shift_1_bit : bool
        Shift 1 bit for the matrix operator, default = False

    sub_128 : bool
        Sub 128 for the matrix operator, default = False

    radix : int
        default = 8

    pad_limit_on : bool
        turn on the hw pad limit(up to 255 pre side) or not, default = True(on)

    Returns
    ----------
    out : np.ndarray 

    Examples
    ----------
        >>> image = './w512xh375_rgb565.raw'
        >>> image_data = kp.inproc_720(image,raw_fmt='rgb565',raw_size=(512,375),npu_size=(416,416),pad_mode=1)
    """
    # ============================================================================================
    # assert
    # ============================================================================================
    assert(npu_size is not None)
    assert(isinstance(npu_size,tuple))
    assert(len(npu_size) == 2)

    # ============================================================================================
    # load image
    # ============================================================================================
    ## if input is jpg/bin
    if not isinstance(image, np.ndarray):
        ## assume input is jpg
        if raw_size == None:
            image = Image.open(image).convert("RGB")
            image = np.array(image).astype('uint8')
            raw_fmt = 'rgb888'
        ## assume input is bin
        else:
            assert raw_fmt is not None
            image = io.load_bin(input_file=image, src_w=raw_size[0], src_h=raw_size[1], image_fmt=raw_fmt)
    else:
        assert(image.ndim >= 2)
        if raw_fmt is None:
            if image.ndim == 2:
                raw_fmt = 'nir'
            elif image.shape[2] == 4:
                raw_fmt = 'rgba'
            else:
                raw_fmt = 'rgb'

    # ============================================================================================
    # init para
    # ============================================================================================
    ## Crop
    if crop_box != None:
        crop_fisrt = True
    else:
        crop_fisrt = False

    ## Resize & Pad
    # if pad_length is not None:
    #     assert isinstance(pad_length, list) | isinstance(pad_length, tuple)
    #     if pad_length[1] != 0:
    #         assert ( (pad_length[2]==0) & (pad_length[3]==0) )
    #         if pad_length[0] != 0:
    #             pad_mode = 0
    #         else:
    #             pad_mode = 1
    #     elif pad_length[3] != 0:
    #         assert ( (pad_length[0]==0) & (pad_length[1]==0) )
    #         if pad_length[2] != 0:
    #             pad_mode = 0
    #         else:
    #             pad_mode = 1
    #     else:
    #         pad_mode = 2
    # else:
    pad_mode = utils.str2int(pad_mode)

    if (pad_mode == 0):
        resize_keep_ratio = True
    elif (pad_mode == 1):
        resize_keep_ratio = True
    else:
        resize_keep_ratio = False

    # if isinstance(pad_val,int) | isinstance(pad_val,float):
    #     pad_val = [pad_val,pad_val,pad_val,pad_val]
    # assert isinstance(pad_val, list) | isinstance(pad_val, tuple)
    # assert len(pad_val) == 4

    ## matrix
    if matrix is None:
        matrix = [256,0,0,0,256,0,0,0,256]
    assert(isinstance(matrix,list) | isinstance(matrix,tuple))
    assert(len(matrix) == 9)
    if bias is None:
        bias = [0,0,0]
    assert(isinstance(bias,list)| isinstance(bias,tuple))
    assert(len(bias) == 3)
    matrix_c00 = matrix[0]
    matrix_c01 = matrix[1]
    matrix_c02 = matrix[2]
    matrix_c10 = matrix[3]
    matrix_c11 = matrix[4]
    matrix_c12 = matrix[5]
    matrix_c20 = matrix[6]
    matrix_c21 = matrix[7]
    matrix_c22 = matrix[8]
    vector_b00 = bias[0]
    vector_b01 = bias[1]
    vector_b02 = bias[2]
    if sub_128:
        subvalue = 128
        pad_val = (-128, -128, -128, -128)
    else:
        subvalue = 0
        pad_val = (0, 0, 0, 0)

    # ============================================================================================
    # main flow
    # ============================================================================================
    #crop
    if crop_fisrt:
        image = funcs.crop(image=image, box=crop_box)
          
    #color
    image = funcs.color_convert(image=image, input_fmt=raw_fmt, output_fmt='rgb')

    #resize
    # condider hw constraint: pad length
    hw_Constraint_pad_length = 255
    if resize_keep_ratio:
        cropped_h = image.shape[0]
        cropped_w = image.shape[1]
        new_w, new_h = funcs.calculate_keep_ratio_size(tar_size=npu_size, ori_size=(cropped_w,cropped_h), platform_type='720', pad_mode=pad_mode)
    else:
        new_w = npu_size[0]
        new_h = npu_size[1]

    if resize_size is not None:
        assert ( (new_w == resize_size[0]) & (new_h == resize_size[1]))

    image = funcs.resize_fixed(image=image, size=(new_w, new_h), platform_type='720')

    #matrix
    h, w, _ = image.shape
    image_f = image[:, :, 0:3].reshape((h * w, 3))
    if shift_1_bit:
        matrix_c = np.array([[128, 0, 0],
                                [0, 128, 0],
                                [0, 0, 128]])
        b = np.array([[0], [0], [0]])
    else:
        matrix_c = np.array([[matrix_c00, matrix_c01, matrix_c02],
                                [matrix_c10, matrix_c11, matrix_c12],
                                [matrix_c20, matrix_c21, matrix_c22]])
        b = np.array([[vector_b00], [vector_b01], [vector_b02]])

    calculated_image_f = np.zeros(image_f.shape, dtype=np.int32)
    for i in range(h*w):
        pt = np.swapaxes(image_f[np.newaxis, i, :], 0, 1)
        matrix_pt = np.floor(np.multiply((matrix_c @ pt), 1/np.power(2, 1)))
        matrix_pt.astype(int)
        result = np.floor(np.multiply(np.add(matrix_pt, b), 1/np.power(2, 7)))
        result.astype(int)
        result = twos_complement_pix(result)

        # if shiftvalue == 1:
        #     result = clip_pix(np.add(result, -128 * np.ones(result.shape)), -128, 127)
        # else:
        #     result = clip_pix(result, 0, 255)
        result = clip_pix(result, 0, 255)

        result = result - np.array([[subvalue], [subvalue], [subvalue]])
        calculated_image_f[i, :] = utils.clip_ary(np.squeeze(result))
    calculated_image = calculated_image_f.reshape(image[:, :, 0:3].shape)
    image = calculated_image

    #padding
    if pad_mode == 0:
        pad_l,pad_r,pad_t,pad_b = funcs.calculate_pad_length_center(tar_size=npu_size,ori_size=(image.shape[1],image.shape[0]))
        if pad_length is not None:
            assert ( (pad_length[0]==pad_l) & (pad_length[1]==pad_r) & (pad_length[2]==pad_t) & (pad_length[3]==pad_b))
        image = funcs.pad(image=image, pad_l=pad_l, pad_r=pad_r, pad_t=pad_t,pad_b=pad_b, pad_val=pad_val)
    elif pad_mode == 1:
        pad_l,pad_r,pad_t,pad_b = funcs.calculate_pad_length_corner(tar_size=npu_size,ori_size=(image.shape[1],image.shape[0]))
        if pad_length is not None:
            assert ( (pad_length[0]==pad_l) & (pad_length[1]==pad_r) & (pad_length[2]==pad_t) & (pad_length[3]==pad_b))        
        image = funcs.pad(image=image, pad_l=pad_l, pad_r=pad_r, pad_t=pad_t,pad_b=pad_b, pad_val=pad_val)
    else:
        image = np.concatenate((image, np.zeros((image.shape[0], image.shape[1], 1), dtype=np.uint8) ), axis=2)

    image=image.astype('uint8')

    if return_float:
        image = image[:,:,:3]
        image = (image + subvalue).astype('uint8')
        if shift_1_bit:
            image = image<<1
            norm = 2 ** (radix+1)
        else:
            norm = 2 ** radix
        image = (image.astype('float')-subvalue) / norm

    ## 256/256 - 0.5 -> 0.5 ~ -0.5   =   (256-128) / 256    radix = 8 => 128 ~ -128, shift = Fasle, sub = 128
    ## 256/256       ->   1 ~  0     =   (256)     / 256    radix = 7 => 128 ~    0, shift = True,  sub =   0
    ## 256/128 -   1 ->   1 ~ -1     =   (256-128) / 128    radix = 7 => 128 ~ -128, shift = Fasle, sub = 128
    ## 256/128       ->   2 ~  0     =   (256)     / 128    radix = 6 => 128 ~    0, shift = True,  sub =   0
    ## 256/64  -   2 ->   2 ~ -2     =   (256-128) /  64    radix = 6 => 128 ~ -128, shift = False, sub = 128

    return image

def twos_complement(value):
    value = int(value)
    # msb = (value & 0x8000) * (1/np.power(2, 15))
    msb = (value & 0x8000) >> 15
    if msb == 1:
        if (((~value) & 0xFFFF) + 1) >= 0xFFFF:
            result = ((~value) & 0xFFFF)
        else:
            result = (((~value) & 0xFFFF) + 1)
        result = result * (-1)
    else:
        result = value

    return result


def twos_complement_pix(value):
    h, _ = value.shape
    for i in range(h):
        value[i, 0] = twos_complement(value[i, 0])

    return value

def clip(value, mini, maxi):
    if value < mini:
        result = mini
    elif value > maxi:
        result = maxi
    else:
        result = value

    return result

def clip_pix(value, mini, maxi):
    h, _ = value.shape
    for i in range(h):
        value[i, 0] = clip(value[i, 0], mini, maxi)

    return value