import math

import cv2
import numpy as np
import torch

from backend import memory_management


def prepare_free_memory(aggressive=False):
    if aggressive:
        memory_management.unload_all_models()
        print("Cleanup all memory...")
        return

    memory_management.free_memory(memory_required=memory_management.minimum_inference_memory(), device=memory_management.get_torch_device())
    print("Cleanup minimal inference memory...")


def apply_circular_forge(model, tiling_enabled=False):
    if not model.is_webui_legacy_model():
        return
    if model.tiling_enabled == tiling_enabled:
        return

    model.tiling_enabled = tiling_enabled

    unet: torch.nn.Module = model.forge_objects.unet.model.diffusion_model
    for layer in [layer for layer in unet.modules() if isinstance(layer, torch.nn.Conv2d)]:
        layer.padding_mode = "circular" if tiling_enabled else "zeros"

    vae: torch.nn.Module = model.forge_objects.vae.first_stage_model
    for layer in [layer for layer in vae.modules() if isinstance(layer, torch.nn.Conv2d)]:
        layer.padding_mode = "circular" if tiling_enabled else "zeros"

    print(f"Tiling: {tiling_enabled}")


def HWC3(x: np.ndarray) -> np.ndarray:
    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


@torch.inference_mode()
def pytorch_to_numpy(x: torch.Tensor) -> np.ndarray:
    return [np.clip(255.0 * y.cpu().numpy(), 0, 255).astype(np.uint8) for y in x]


@torch.inference_mode()
def numpy_to_pytorch(x: np.ndarray) -> torch.Tensor:
    y = x.astype(np.float32) / 255.0
    y = y[None]
    y = np.ascontiguousarray(y.copy())
    y = torch.from_numpy(y).float()
    return y


def pad64(x: int) -> int:
    return int(math.ceil(x / 64.0) * 64 - x)


def safer_memory(x: np.ndarray) -> np.ndarray:
    # Fix many macOS / AMD problems
    return np.ascontiguousarray(x.copy()).copy()


def resize_image_with_pad(input_image: np.ndarray, resolution: int, *, skip_hwc3: bool = False):
    img = input_image if skip_hwc3 else HWC3(input_image)
    H_raw, W_raw, _ = img.shape
    k = float(resolution) / float(min(H_raw, W_raw))
    interpolation = cv2.INTER_CUBIC if k > 1 else cv2.INTER_AREA
    H_target = int(np.round(float(H_raw) * k))
    W_target = int(np.round(float(W_raw) * k))
    img = cv2.resize(img, (W_target, H_target), interpolation=interpolation)
    H_pad, W_pad = pad64(H_target), pad64(W_target)
    img_padded = np.pad(img, [[0, H_pad], [0, W_pad], [0, 0]], mode="edge")

    def remove_pad(x):
        return safer_memory(x[:H_target, :W_target])

    return safer_memory(img_padded), remove_pad
