import numpy as np
import imageio
import torch
import os
import matplotlib.pyplot as plt
import cv2

from promptda.utils.logger import Log

# DEVICE = 'cuda' if torch.cuda.is_available(
# ) else 'mps' if torch.backends.mps.is_available() else 'cpu'


def to_tensor_func(arr):
    if arr.ndim == 2:
        arr = arr[:, :, np.newaxis]
    return torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0)


def to_numpy_func(tensor):
    arr = tensor.squeeze(0).permute(1, 2, 0).cpu().numpy()
    if arr.shape[2] == 1:
        arr = arr[:, :, 0]
    return arr


def ensure_multiple_of(x, multiple_of=14):
    return int(x // multiple_of * multiple_of)


def load_image(image_path, to_tensor=True, max_size=1008, multiple_of=14):
    '''
    Load image from path and convert to tensor
    max_size // 14 = 0
    '''
    image = np.asarray(imageio.imread(image_path)).astype(np.float32)
    image = image / 255.

    max_size = max_size // multiple_of * multiple_of
    if max(image.shape) > max_size:
        h, w = image.shape[:2]
        scale = max_size / max(h, w)
        tar_h = ensure_multiple_of(h * scale)
        tar_w = ensure_multiple_of(w * scale)
        image = cv2.resize(image, (tar_w, tar_h), interpolation=cv2.INTER_AREA)
    if to_tensor:
        return to_tensor_func(image)
    return image


def load_depth(depth_path, to_tensor=True):
    '''
    Load depth from path and convert to tensor
    '''
    if depth_path.endswith('.png'):
        depth = np.asarray(imageio.imread(depth_path)).astype(np.float32)
        depth = depth / 1000.
    elif depth_path.endswith('.npz'):
        depth = np.load(depth_path)['depth']
    else:
        raise ValueError(f"Unsupported depth format: {depth_path}")
    if to_tensor:
        return to_tensor_func(depth)
    return depth


def save_depth(depth,
               prompt_depth=None,
               image=None,
               output_path='data/output/depth.png',
               save_vis=True):
    '''
    Save depth to path
    '''
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    depth = to_numpy_func(depth)
    uint16_depth = (depth * 1000.).astype(np.uint16)
    imageio.imwrite(output_path, uint16_depth)

    if not save_vis:
        return

    output_path = output_path.replace('.png', '_vis.png')
    prompt_depth = to_numpy_func(prompt_depth)
    image = to_numpy_func(image)
    plt.subplot(1, 3, 1)
    plt.imshow(image)
    plt.axis('off')
    plt.subplot(1, 3, 2)
    plt.imshow(prompt_depth)
    plt.axis('off')
    plt.subplot(1, 3, 3)
    plt.imshow(depth)
    plt.axis('off')
    plt.tight_layout()
    plt.savefig(output_path)
    plt.close()
    Log.info(f'Saved depth to {output_path}')