import numpy as np
import torch
import cv2

def HWC3(x):
    assert x.dtype == np.uint8
    if x.ndim == 2:
        x = x[:, :, None]
    assert x.ndim == 3
    H, W, C = x.shape
    assert C == 1 or C == 3 or C == 4
    if C == 3:
        return x
    if C == 1:
        return np.concatenate([x, x, x], axis=2)
    if C == 4:
        color = x[:, :, 0:3].astype(np.float32)
        alpha = x[:, :, 3:4].astype(np.float32) / 255.0
        y = color * alpha + 255.0 * (1.0 - alpha)
        y = y.clip(0, 255).astype(np.uint8)
        return y


def resize_image(input_image, resolution):
    H, W, C = input_image.shape
    H = float(H)
    W = float(W)
    k = float(resolution) / min(H, W)
    H *= k
    W *= k
    H = int(np.round(H / 64.0)) * 64
    W = int(np.round(W / 64.0)) * 64
    img = cv2.resize(input_image, (W, H), interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA)
    return img


# normalize
def norm_normalize(norm_out):
    norm_x, norm_y, norm_z = torch.split(norm_out, 1, dim=0)
    norm = torch.sqrt(norm_x ** 2.0 + norm_y ** 2.0 + norm_z ** 2.0) + 1e-10
    final_out = torch.cat([norm_x / norm, norm_y / norm, norm_z / norm], dim=0)
    fg_mask = torch.ones_like(norm).repeat(3, 1, 1)
    fg_mask[norm.repeat(3, 1, 1) < 0.5] = 0.
    fg_mask[norm.repeat(3, 1, 1) > 1.5] = 0.

    final_out[norm.repeat(3, 1, 1) < 0.5] = -1
    final_out[norm.repeat(3, 1, 1) > 1.5] = -1
    return final_out, fg_mask.bool()


def center_crop(input_image):
    height, width = input_image.shape[:2]

    if height < width:
        min_dim = height
    else:
        min_dim = width

    center_x = width // 2
    center_y = height // 2
    half_length = min_dim // 2

    crop_x1 = center_x - half_length
    crop_x2 = center_x + half_length
    crop_y1 = center_y - half_length
    crop_y2 = center_y + half_length

    center_cropped_image = input_image[crop_y1:crop_y2, crop_x1:crop_x2]

    return center_cropped_image


def flip_x(normal):

    if isinstance(normal, np.ndarray):
        return normal.dot(np.array([[-1, 0, 0], [0, 1, 0], [0, 0, 1]])).astype(np.float32)
    else:
        trans = torch.tensor([[-1, 0, 0], [0, 1, 0], [0, 0, 1]]).float()
        return  normal @ trans