import cv2
import numpy as np
import torch
from jaxtyping import Float


def read_img(path):
    img = cv2.imread(str(path), cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH)
    if img.ndim == 3:
        img = cv2.cvtColor(img[..., :3], cv2.COLOR_BGR2RGB)
    elif img.ndim == 2:
        img = img[..., np.newaxis]
    dinfo = np.iinfo(img.dtype)
    return (img.astype(np.float32) / dinfo.max) * 2 - 1


def write_img(path: str, data: np.ndarray):
    data = np.clip(data * 0.5 + 0.5, 0, 1)
    if data.ndim == 3 and data.shape[-1] == 3:
        data = cv2.cvtColor(data, cv2.COLOR_RGB2BGR)
    elif data.ndim == 2:
        data = data[..., np.newaxis]

    data = (data * 255).astype(np.uint8)
    cv2.imwrite(path, data)


def to_torch(img: Float[np.ndarray, "H W C"]) -> Float[torch.Tensor, "C H W"]:
    return torch.from_numpy(img).permute(2, 0, 1)


def from_torch(img: Float[torch.Tensor, "C H W"]) -> Float[np.ndarray, "H W C"]:
    return img.permute(1, 2, 0).detach().cpu().float().numpy()