# -*- coding: UTF-8 -*-
'''=================================================
@Project -> File   pram -> common
@IDE    PyCharm
@Author fx221@cam.ac.uk
@Date   29/01/2024 15:05
=================================================='''
import os
import torch
import json
import yaml
import cv2
import numpy as np
from typing import Tuple
from copy import deepcopy


def load_args(args, save_path):
    with open(save_path, "r") as f:
        args.__dict__ = json.load(f)


def save_args_yaml(args, save_path):
    with open(save_path, 'w') as f:
        yaml.dump(args, f)


def merge_tags(tags: list, connection='_'):
    out = ''
    for i, t in enumerate(tags):
        if i == 0:
            out = out + t
        else:
            out = out + connection + t
    return out


def torch_set_gpu(gpus):
    if type(gpus) is int:
        gpus = [gpus]

    cuda = all(gpu >= 0 for gpu in gpus)

    if cuda:
        os.environ['CUDA_VISIBLE_DEVICES'] = ','.join([str(gpu) for gpu in gpus])
        # print(os.environ['CUDA_VISIBLE_DEVICES'])
        assert cuda and torch.cuda.is_available(), "%s has GPUs %s unavailable" % (
            os.environ['HOSTNAME'], os.environ['CUDA_VISIBLE_DEVICES'])
        torch.backends.cudnn.benchmark = True  # speed-up cudnn
        torch.backends.cudnn.fastest = True  # even more speed-up?
        print('Launching on GPUs ' + os.environ['CUDA_VISIBLE_DEVICES'])

    else:
        print('Launching on CPU')

    return cuda


def resize_img(img, nh=-1, nw=-1, rmax=-1, mode=cv2.INTER_NEAREST):
    assert nh > 0 or nw > 0 or rmax > 0
    if nh > 0:
        return cv2.resize(img, dsize=(int(img.shape[1] / img.shape[0] * nh), nh), interpolation=mode)
    if nw > 0:
        return cv2.resize(img, dsize=(nw, int(img.shape[0] / img.shape[1] * nw)), interpolation=mode)
    if rmax > 0:
        oh, ow = img.shape[0], img.shape[1]
        if oh > ow:
            return cv2.resize(img, dsize=(int(img.shape[1] / img.shape[0] * rmax), rmax), interpolation=mode)
        else:
            return cv2.resize(img, dsize=(rmax, int(img.shape[0] / img.shape[1] * rmax)), interpolation=mode)

    return cv2.resize(img, dsize=(nw, nh), interpolation=mode)


def resize_image_with_padding(image: np.array, nw: int, nh: int, padding_color: Tuple[int] = (0, 0, 0)) -> np.array:
    """Maintains aspect ratio and resizes with padding.
    Params:
        image: Image to be resized.
        new_shape: Expected (width, height) of new image.
        padding_color: Tuple in BGR of padding color
    Returns:
        image: Resized image with padding
    """
    original_shape = (image.shape[1], image.shape[0])  # (w, h)
    ratio_w = nw / original_shape[0]
    ratio_h = nh / original_shape[1]

    if ratio_w == ratio_h:
        image = cv2.resize(image, (nw, nh), interpolation=cv2.INTER_NEAREST)

    ratio = ratio_w if ratio_w < ratio_h else ratio_h

    new_size = tuple([int(x * ratio) for x in original_shape])
    image = cv2.resize(image, new_size, interpolation=cv2.INTER_NEAREST)
    delta_w = nw - new_size[0] if nw > new_size[0] else new_size[0] - nw
    delta_h = nh - new_size[1] if nh > new_size[1] else new_size[1] - nh

    left, right = delta_w // 2, delta_w - (delta_w // 2)
    top, bottom = delta_h // 2, delta_h - (delta_h // 2)

    # print('top, bottom, left, right: ', top, bottom, left, right)
    image = cv2.copyMakeBorder(image, top, bottom, left, right, cv2.BORDER_CONSTANT, value=padding_color)
    return image


def puttext_with_background(image, text, org=(0, 0), fontFace=cv2.FONT_HERSHEY_SIMPLEX,
                            fontScale=1, text_color=(0, 0, 255),
                            thickness=2, lineType=cv2.LINE_AA, bg_color=None):
    out_img = deepcopy(image)
    if bg_color is not None:
        (text_width, text_height), baseline = cv2.getTextSize(text,
                                                              fontFace,
                                                              fontScale=fontScale,
                                                              thickness=thickness)
        box_coords = (
            (org[0], org[1] + baseline),
            (org[0] + text_width + 2, org[1] - text_height - 2))

        cv2.rectangle(out_img, box_coords[0], box_coords[1], bg_color, cv2.FILLED)
    out_img = cv2.putText(img=out_img, text=text,
                          org=org,
                          fontFace=fontFace,
                          fontScale=fontScale, color=text_color,
                          thickness=thickness, lineType=lineType)
    return out_img