import cv2
import numpy as np
import torch


def save_img(img, name, gamma=False):
    if gamma:
        img = np.power(img, 1/2.2)
    img = np.clip(img, 0, 1)
    img = (img * 65535).astype(np.uint16)
    if img.ndim == 3:
        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    cv2.imwrite(name, img)


def compute_y(img: np.ndarray) -> np.ndarray:
    y = 0.299 * img[:, :, 0] + 0.587 * img[:, :, 1] + 0.114 * img[:, :, 2]
    return y


def compute_raw_y(img: np.ndarray) -> np.ndarray:
    g1 = img[..., 1]
    g2 = img[..., 2]
    ret = (g1 + g2) / 2
    return ret


def pack_raw(im):
    # pack Bayer image to 4 channels
    if isinstance(im, torch.Tensor):
        im = torch.unsqueeze(im, dim=-1)
        img_shape = im.shape
        H = img_shape[0]
        W = img_shape[1]

        out = torch.cat((im[0:H:2, 0:W:2, :], im[0:H:2, 1:W:2, :],
                         im[1:H:2, 1:W:2, :], im[1:H:2, 0:W:2, :]),
                        dim=-1)
    elif isinstance(im, np.ndarray):
        im = np.expand_dims(im, axis=-1)
        img_shape = im.shape
        H = img_shape[0]
        W = img_shape[1]

        out = np.concatenate((im[0:H:2, 0:W:2, :], im[0:H:2, 1:W:2, :],
                              im[1:H:2, 1:W:2, :], im[1:H:2, 0:W:2, :]),
                             axis=-1)
    return out


def depack_raw(im):
    # unpack 4 channels to Bayer image
    img_shape = im.shape
    H = img_shape[0]
    W = img_shape[1]
    if isinstance(im, torch.Tensor):
        output = torch.zeros((H * 2, W * 2), dtype=im.dtype)
    elif isinstance(im, np.ndarray):
        output = np.zeros((H * 2, W * 2), dtype=im.dtype)
    img_shape = output.shape
    H = img_shape[0]
    W = img_shape[1]

    output[0:H:2, 0:W:2] = im[:, :, 0]
    output[0:H:2, 1:W:2] = im[:, :, 1]
    output[1:H:2, 1:W:2] = im[:, :, 2]
    output[1:H:2, 0:W:2] = im[:, :, 3]

    return output