from skimage.measure import label
#Scientific computing 
import numpy as np
import os
import matplotlib.pyplot as plt
import torch.nn.functional as F
from torch.nn.functional import one_hot
import cv2
import torch
import random
#Pytorch packages

def random_sum_to(n, num_terms = None):
    '''
    generate num_tersm with sum as n
    '''
    num_terms = (num_terms or r.randint(2, n)) - 1
    a = random.sample(range(1, n), num_terms) + [0, n]
    list.sort(a)
    return [a[i+1] - a[i] for i in range(len(a) - 1)]



def get_first_prompt(mask_cls,dist_thre_ratio=0.3,prompt_num=5,max_prompt_num=15,region_type='random'):
    '''
    if region_type = random, we random select one region and generate prompt
    if region_type = all, we generate prompt at each object region
    if region_type = largest_k, we generate prompt at largest k region, k <10
    '''
    if prompt_num==-1:
        prompt_num = random.randint(1, max_prompt_num)
    # Find all disconnected regions
    label_msk, region_ids = label(mask_cls, connectivity=2, return_num=True)
    #print('num of regions found', region_ids)
    ratio_list, regionid_list = [], []
    for region_id in range(1, region_ids+1):
        #find coordinates of points in the region
        binary_msk = np.where(label_msk==region_id, 1, 0)

        # clean some region that is abnormally small
        r = np.sum(binary_msk) / np.sum(mask_cls)
        #print('curr mask over all mask ratio', r)
        ratio_list.append(r)
        regionid_list.append(region_id)
    if len(ratio_list)>0:
        ratio_list, regionid_list = zip(*sorted(zip(ratio_list, regionid_list)))
        regionid_list = regionid_list[::-1]
    
        if region_type == 'random':
            prompt_num = 1
            regionid_list = [random.choice(regionid_list)] # random choose 1 region
            prompt_num_each_region = [1]
        elif region_type[:7] == 'largest':
            region_max_num = int(region_type.split('_')[-1])
            #print(region_max_num,prompt_num,len(regionid_list))
            valid_region = min(region_max_num,len(regionid_list))
            if valid_region<prompt_num:
                prompt_num_each_region = random_sum_to(prompt_num,valid_region)
            else:
                prompt_num_each_region = prompt_num*[1]
            regionid_list = regionid_list[:min(valid_region,prompt_num)]
            #print(prompt_num_each_region)
        else:
            prompt_num_each_region = len(regionid_list)*[1]


        prompt = []
        mask_curr = np.zeros_like(label_msk)
        

        for reg_id in range(len(regionid_list)):
            binary_msk = np.where(label_msk==regionid_list[reg_id], 1, 0)
            mask_curr = np.logical_or(binary_msk,mask_curr)


            padded_mask = np.uint8(np.pad(binary_msk, ((1, 1), (1, 1)), 'constant'))
            dist_img = cv2.distanceTransform(padded_mask, distanceType=cv2.DIST_L2, maskSize=5).astype(np.float32)[1:-1, 1:-1]

            # sort the distances 
            dist_array=sorted(dist_img.copy().flatten())[::-1]
            dist_array = np.array(dist_array)
            # find the threshold:
            dis_thre = max(dist_array[int(dist_thre_ratio*np.sum(dist_array>0))],1)
            #print(np.max(dist_array))
            #print(dis_thre)
            cY, cX = np.where(dist_img>=dis_thre)
            while prompt_num_each_region[reg_id]>0:
                # random select one prompt
                random_idx = np.random.randint(0, len(cX))
                cx, cy = int(cX[random_idx]), int(cY[random_idx])
                prompt.append((cx,cy,1))
                prompt_num_each_region[reg_id] -=1

        while len(prompt)<max_prompt_num: # repeat prompt to ensure the same size
            prompt.append((cx,cy,1))
    else: # if this image doesn't have target object
        prompt = [(0,0,-1)]
        mask_curr = np.zeros_like(label_msk)
        while len(prompt)<max_prompt_num: # repeat prompt to ensure the same size
            prompt.append((0,0,-1))
    prompt = np.array(prompt) 
    mask_curr = np.array(mask_curr,dtype=int)
    return prompt,mask_curr


def get_top_boxes(mask_cls,dist_thre_ratio=0.10,prompt_num=15,region_type='largest_15'):
    # Find all disconnected regions
    label_msk, region_ids = label(mask_cls, connectivity=2, return_num=True)
    #print('num of regions found', region_ids)
    ratio_list, regionid_list = [], []
    for region_id in range(1, region_ids+1):
        #find coordinates of points in the region
        binary_msk = np.where(label_msk==region_id, 1, 0)

        # clean some region that is abnormally small
        r = np.sum(binary_msk) / np.sum(mask_cls)
        #print('curr mask over all mask ratio', r)
        ratio_list.append(r)
        regionid_list.append(region_id)
    if len(ratio_list)>0:
        # sort the region from largest to smallest
        ratio_list, regionid_list = zip(*sorted(zip(ratio_list, regionid_list)))
        regionid_list = regionid_list[::-1]

        if region_type == 'random':
            prompt_num = 1
            regionid_list = [random.choice(regionid_list)] # random choose 1 region
        elif region_type[:7] == 'largest':
            region_max_num = int(region_type.split('_')[-1])
            regionid_list = regionid_list[:min(region_max_num,len(regionid_list))]

        prompt = []
        mask_curr = np.zeros_like(label_msk)
        for reg_id in range(len(regionid_list)):
            binary_msk = np.where(label_msk==regionid_list[reg_id], 1, 0)
            mask_curr = np.logical_or(binary_msk,mask_curr)
            box = MaskToBoxSimple(binary_msk,dist_thre_ratio)
            prompt.append(box)

        while len(prompt)<prompt_num: # repeat prompt to ensure the same size
            prompt.append(box)
        prompt = np.array(prompt) 
        mask_curr = np.array(mask_curr,dtype=int)
    else:
        prompt = [[0,0,0,0]]
        mask_curr = np.zeros_like(label_msk)
        while len(prompt)<prompt_num:
            prompt.append(prompt[0])
    return prompt,mask_curr
        
def MaskToBoxSimple(mask,random_thre=0.05):
    '''
    random_thre, the randomness at each side of box
    '''
    mask = mask.squeeze()
    
    y_max,x_max = mask.shape[0],mask.shape[1]
    
    #find coordinates of points in the region
    row, col = np.argwhere(mask).T
    # find the four corner coordinates
    y0,x0 = row.min(),col.min()
    y1,x1 = row.max(),col.max()
    
    y_thre = (y1-y0)*random_thre
    x_thre = (x1-x0)*random_thre
    
    x0 = max(0,x0-x_thre*random.random())
    x1 = min(x_max,x1+x_thre*random.random())
    
    y0 = max(0,y0-y_thre*random.random())
    y1 = min(y_max,y1+y_thre*random.random())
    

    return [x0,y0,x1,y1]

def min_max_normalize(tensor,p=0.01):
    p_min = torch.quantile(tensor,p)
    p_max = torch.quantile(tensor,1-p)
    tensor = torch.clamp(tensor,p_min,p_max)
    return tensor