import tqdm
import torch
from torchvision.transforms.functional import to_tensor
import numpy as np
import random
import cv2

def gen_dilate(alpha, min_kernel_size, max_kernel_size): 
    kernel_size = random.randint(min_kernel_size, max_kernel_size)
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size,kernel_size))
    fg_and_unknown = np.array(np.not_equal(alpha, 0).astype(np.float32))
    dilate = cv2.dilate(fg_and_unknown, kernel, iterations=1)*255
    return dilate.astype(np.float32)

def gen_erosion(alpha, min_kernel_size, max_kernel_size): 
    kernel_size = random.randint(min_kernel_size, max_kernel_size)
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size,kernel_size))
    fg = np.array(np.equal(alpha, 255).astype(np.float32))
    erode = cv2.erode(fg, kernel, iterations=1)*255
    return erode.astype(np.float32)

@torch.inference_mode()
@torch.cuda.amp.autocast()
def matanyone(processor, frames_np, mask, r_erode=0, r_dilate=0, n_warmup=10):
    """
    Args:
        frames_np: [(H,W,C)]*n, uint8
        mask: (H,W), uint8
    Outputs:
        com: [(H,W,C)]*n, uint8
        pha: [(H,W,C)]*n, uint8
    """

    # print(f'===== [r_erode] {r_erode}; [r_dilate] {r_dilate} =====')
    bgr = (np.array([120, 255, 155], dtype=np.float32)/255).reshape((1, 1, 3))
    objects = [1]

    # [optional] erode & dilate on given seg mask
    if r_dilate > 0:
        mask = gen_dilate(mask, r_dilate, r_dilate)
    if r_erode > 0:
        mask = gen_erosion(mask, r_erode, r_erode)

    mask = torch.from_numpy(mask).cuda()

    frames_np = [frames_np[0]]* n_warmup + frames_np

    frames = []
    phas = []
    for ti, frame_single in tqdm.tqdm(enumerate(frames_np)):
        image = to_tensor(frame_single).cuda().float()

        if ti == 0:
            output_prob = processor.step(image, mask, objects=objects)      # encode given mask
            output_prob = processor.step(image, first_frame_pred=True)      # clear past memory for warmup frames
        else:
            if ti <= n_warmup:
                output_prob = processor.step(image, first_frame_pred=True)  # clear past memory for warmup frames
            else:
                output_prob = processor.step(image)

        # convert output probabilities to an object mask
        mask = processor.output_prob_to_mask(output_prob)

        pha = mask.unsqueeze(2).cpu().numpy()
        com_np = frame_single / 255. * pha + bgr * (1 - pha)
        
        # DONOT save the warmup frames
        if ti > (n_warmup-1):
            frames.append((com_np*255).astype(np.uint8))
            phas.append((pha*255).astype(np.uint8))
    
    return frames, phas