import cv2
import numpy as np
import random

perm = [(0, 1, 2), (0, 2, 1), (1, 0, 2), (1, 2, 0), (2, 0, 1), (2, 1, 0)]
rotate = [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_180, cv2.ROTATE_90_COUNTERCLOCKWISE]


def random_crop(img0, imgt, img1, crop_size, flowt0=None, flowt1=None, distance=None):
    im_h, im_w = img0.shape[:2]
    crop_h, crop_w = crop_size, crop_size
    i = random.randint(0, im_h - crop_h)
    j = random.randint(0, im_w - crop_w)
    img0 = img0[i:i + crop_h, j:j + crop_w]
    imgt = imgt[i:i + crop_h, j:j + crop_w]
    img1 = img1[i:i + crop_h, j:j + crop_w]
    if flowt0 is not None and flowt1 is not None:
        flowt0 = flowt0[i:i + crop_h, j:j + crop_w]
        flowt1 = flowt1[i:i + crop_h, j:j + crop_w]
    if distance is not None:
        distance = distance[i:i + crop_h, j:j + crop_w]
    return img0, imgt, img1, flowt0, flowt1, distance


def random_hor_flip(img0, imgt, img1, flowt0=None, flowt1=None, distance=None):
    img0, imgt, img1 = img0[::-1, :, :], imgt[::-1, :, :], img1[::-1, :, :]
    if flowt0 is not None and flowt1 is not None:
        flowt0, flowt1 = flowt0[::-1, :, :] * np.array([1, -1]).reshape(1, 1, 2), \
                         flowt1[::-1, :, :] * np.array([1, -1]).reshape(1, 1, 2)
    if distance is not None:
        distance = distance[::-1, :, :]
    return img0, imgt, img1, flowt0, flowt1, distance


def random_ver_flip(img0, imgt, img1, flowt0=None, flowt1=None, distance=None):
    img0, imgt, img1 = img0[:, ::-1, :], imgt[:, ::-1, :], img1[:, ::-1, :]
    if flowt0 is not None and flowt1 is not None:
        flowt0, flowt1 = flowt0[:, ::-1, :] * np.array([-1, 1]).reshape(1, 1, 2), \
                         flowt1[:, ::-1, :] * np.array([-1, 1]).reshape(1, 1, 2)
    if distance is not None:
        distance = distance[:, ::-1, :]
    return img0, imgt, img1, flowt0, flowt1, distance


def random_color_permutation(img0, imgt, img1):
    perm_idx = random.randint(0, 5)
    img0, imgt, img1 = img0[:, :, perm[perm_idx]], imgt[:, :, perm[perm_idx]], img1[:, :, perm[perm_idx]]
    return img0, imgt, img1


def random_temporal_flip(img0, imgt, img1, time_step, flowt0=None, flowt1=None):
    tmp = img1
    img1 = img0
    img0 = tmp
    time_step = 1 - time_step
    if flowt0 is not None and flowt1 is not None:
        tmp = flowt0
        flowt0 = flowt1
        flowt1 = tmp
    return img0, imgt, img1, time_step, flowt0, flowt1


def random_rotation(img0, imgt, img1, degree, flowt0=None, flowt1=None, distance=None):
    if degree != 3:
        img0 = cv2.rotate(img0, rotate[degree])
        imgt = cv2.rotate(imgt, rotate[degree])
        img1 = cv2.rotate(img1, rotate[degree])
        if flowt0 is not None and flowt1 is not None:
            flowt0 = cv2.rotate(flowt0, rotate[degree])
            flowt1 = cv2.rotate(flowt1, rotate[degree])
            if degree == 0:
                flowt0 = np.concatenate((-flowt0[:, :, 1:2], flowt0[:, :, 0:1]), 2)
                flowt1 = np.concatenate((-flowt1[:, :, 1:2], flowt1[:, :, 0:1]), 2)
            elif degree == 1:
                flowt0 = -flowt0
                flowt1 = -flowt1
            elif degree == 2:
                flowt0 = np.concatenate((flowt0[:, :, 1:2], -flowt0[:, :, 0:1]), 2)
                flowt1 = np.concatenate((flowt1[:, :, 1:2], -flowt1[:, :, 0:1]), 2)
        if distance is not None:
            H,W,_ = distance.shape
            distance = cv2.rotate(distance, rotate[degree]).reshape(H,W,1)
    return img0, imgt, img1, flowt0, flowt1, distance


def random_resize(img0, imgt, img1, flowt0=None, flowt1=None, distance=None):
    '''
    img0 = cv2.resize(img0, dsize=None, fx=2.0, fy=2.0, interpolation=cv2.INTER_LINEAR)
    imgt = cv2.resize(imgt, dsize=None, fx=2.0, fy=2.0, interpolation=cv2.INTER_LINEAR)
    img1 = cv2.resize(img1, dsize=None, fx=2.0, fy=2.0, interpolation=cv2.INTER_LINEAR)
    if flowt0 is not None and flowt1 is not None:
        flowt0 = cv2.resize(flowt0, dsize=None, fx=2.0, fy=2.0, interpolation=cv2.INTER_LINEAR) * 2.0
        flowt1 = cv2.resize(flowt1, dsize=None, fx=2.0, fy=2.0, interpolation=cv2.INTER_LINEAR) * 2.0
    '''
    return img0, imgt, img1, flowt0, flowt1, distance


def read_flow(name):
    with open(name, "rb") as f:
        header = f.read(4)
        if header.decode("utf-8") != 'PIEH':
            raise Exception('Flow file header does not contain PIEH')

        width = np.fromfile(f, np.int32, 1).squeeze()
        height = np.fromfile(f, np.int32, 1).squeeze()

        flow = np.fromfile(f, np.float32, width * height * 2).reshape((height, width, 2))

    return flow.astype(np.float32)