from transformers import AutoModelForImageSegmentation, AutoConfig
from transformers.dynamic_module_utils import get_class_from_dynamic_module
from transformers.models.auto.auto_factory import add_generation_mixin_to_remote_model
import torch
from torchvision import transforms
import numpy as np
from PIL import Image
import torch.nn.functional as F
import os
from scipy.ndimage import binary_erosion
from pymatting.alpha.estimate_alpha_cf import estimate_alpha_cf
from pymatting.foreground.estimate_foreground_ml import estimate_foreground_ml
from pymatting.util.util import stack_images
from .rmbg.session_factory import new_session
from folder_paths import models_dir, add_model_folder_path, get_folder_paths

torch.set_float32_matmul_precision(["high", "highest"][0])

transform_image = transforms.Compose(
    [
        transforms.Resize((1024, 1024)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    ]
)

current_path  = os.getcwd()

## ComfyUI portable standalone build for Windows 

add_model_folder_path('birefnet', os.path.join(models_dir, "birefnet"))
model_path = "BiRefNet"

max_gpu_size = 3

def clear_memory():
    import gc
    # Cleanup
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.ipc_collect()

def tensor2pil(image):
    return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))

def pil2tensor(image):
    return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)

def resize_image(image):
    image = image.convert('RGB')
    model_input_size = (1024, 1024)
    image = image.resize(model_input_size, Image.BILINEAR)
    return image


def alpha_matting_cutout(
    img: Image,
    mask: Image,
    foreground_threshold: int,
    background_threshold: int,
    erode_structure_size: int,
) -> Image:
    """
    Perform alpha matting on an image using a given mask and threshold values.

    This function takes a PIL image `img` and a PIL image `mask` as input, along with
    the `foreground_threshold` and `background_threshold` values used to determine
    foreground and background pixels. The `erode_structure_size` parameter specifies
    the size of the erosion structure to be applied to the mask.

    The function returns a PIL image representing the cutout of the foreground object
    from the original image.
    """
    if img.mode == "RGBA" or img.mode == "CMYK":
        img = img.convert("RGB")

    img_array = np.asarray(img)
    mask_array = np.asarray(mask)

    is_foreground = mask_array > foreground_threshold
    is_background = mask_array < background_threshold

    structure = None
    if erode_structure_size > 0:
        structure = np.ones(
            (erode_structure_size, erode_structure_size), dtype=np.uint8
        )

    is_foreground = binary_erosion(is_foreground, structure=structure)
    is_background = binary_erosion(is_background, structure=structure, border_value=1)

    trimap = np.full(mask_array.shape, dtype=np.uint8, fill_value=128)
    trimap[is_foreground] = 255
    trimap[is_background] = 0

    img_normalized = img_array / 255.0
    trimap_normalized = trimap / 255.0

    alpha = estimate_alpha_cf(img_normalized, trimap_normalized)
    foreground = estimate_foreground_ml(img_normalized, alpha)
    cutout = stack_images(foreground, alpha)

    cutout = np.clip(cutout * 255, 0, 255).astype(np.uint8)
    cutout = Image.fromarray(cutout)

    return cutout


def naive_cutout(img: Image, mask: Image) -> Image:
    """
    Perform a simple cutout operation on an image using a mask.

    This function takes a PIL image `img` and a PIL image `mask` as input.
    It uses the mask to create a new image where the pixels from `img` are
    cut out based on the mask.

    The function returns a PIL image representing the cutout of the original
    image using the mask.
    """
    empty = Image.new("RGBA", (img.size), 0)
    cutout = Image.composite(img, empty, mask)
    return cutout

def putalpha_cutout(img: Image, mask: Image) -> Image:
    """
    Apply the specified mask to the image as an alpha cutout.

    Args:
        img (PILImage): The image to be modified.
        mask (PILImage): The mask to be applied.

    Returns:
        PILImage: The modified image with the alpha cutout applied.
    """
    img.putalpha(mask)
    return img

colors = ["transparency", "green", "white", "red", "yellow", "blue", "black", "pink", "purple", "brown", "violet", "wheat", "whitesmoke", "yellowgreen", "turquoise", "tomato", "thistle", "teal", "tan", "steelblue", "springgreen", "snow", "slategrey", "slateblue", "skyblue", "orange"]

def get_device_by_name(device):
    """
    "device": (["auto", "cuda", "cpu", "mps", "xpu", "meta"],{"default": "auto"}), 
    """
    if device == 'auto':
        try:
            device = "cpu"
            if torch.cuda.is_available():
                device = "cuda"
            elif torch.backends.mps.is_available():
                device = "cpu"
            elif torch.xpu.is_available():
                device = "xpu"
        except:
                raise AttributeError("What's your device(到底用什么设备跑的)？")
    print("\033[93mUse Device(使用设备):", device, "\033[0m")
    return device

def get_device_provider(device):
    providers = []
    cuda_provider_options = {
        "gpu_mem_limit": 9 * 1024 * 1024 * 1024, # 9GB.
        "arena_extend_strategy": "kSameAsRequested",
        "cudnn_conv_use_max_workspace": "0"
    }
    if device == 'cuda' and torch.cuda.is_available():
        providers.append(("CUDAExecutionProvider", cuda_provider_options))
    
    if len(providers) == 0:
        providers.append(("CPUExecutionProvider", {}))
    
    return providers
        

_birefnet_model = None
_birefnet_model_name = ""
_birefnet_onnx_model = None
_birefnet_model_name_onnx = ""

class BiRefNet_onnx:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "device": (["cpu", "cuda"],{"default": "cpu"}),
                "model_name": ([
                    "birefnet-general",
                    "birefnet-general-lite",
                    "birefnet-portrait",
                    "birefnet-cod",
                    "birefnet-dis",
                    "birefnet-hrsod",
                    "birefnet-massive"
                ], {"default": "birefnet-general"}),
                "cutout_func": (["putalpha", "naive", "alpha_matting"],{"default": "putalpha"}),
                # "cached": ("BOOLEAN", {"default": True}),
                "alpha_matting_foreground_threshold": ("INT", {"default": 240}),
                "alpha_matting_background_threshold": ("INT", {"default": 10}),
                "alpha_matting_erode_size": ("INT", {"default": 10}),
            },
        }
    RETURN_TYPES = ("IMAGE", "MASK",)
    RETURN_NAMES = ("image", "mask",)
    FUNCTION = "background_remove"
    CATEGORY = "🔥BiRefNet"
    
    def background_remove(self, 
                          image, 
                          device,
                          model_name,
                          cutout_func,
                        #   cached,
                          alpha_matting_foreground_threshold,
                          alpha_matting_background_threshold,
                          alpha_matting_erode_size,
    ):
        providers = get_device_provider(device)
        processed_images = []
        processed_masks = []
        global _birefnet_onnx_model, _birefnet_model_name_onnx
        
        ## TODO: cache model
        if _birefnet_onnx_model is not None and _birefnet_model_name_onnx == model_name:
            session = _birefnet_onnx_model
            print("Loading from cache...")
        else:
            session = new_session(model_name, providers=providers)
            if _birefnet_onnx_model is not None:
                _birefnet_onnx_model = None
                _birefnet_model_name_onnx = ''
                clear_memory()
            # if cached:
            #     _birefnet_onnx_model = session
            #     _birefnet_model_name_onnx = model_name
        
        for image in image:
            orig_image = tensor2pil(image)
            masks = session.predict(orig_image)
            pil_im = masks[0]
            if cutout_func == 'putalpha':
                new_im = putalpha_cutout(orig_image, pil_im)
            elif cutout_func == 'naive':
                new_im = naive_cutout(orig_image, pil_im)
            elif cutout_func == 'alpha_matting':
                new_im = alpha_matting_cutout(
                    orig_image,
                    pil_im,
                    foreground_threshold=alpha_matting_foreground_threshold,
                    background_threshold=alpha_matting_background_threshold,
                    erode_structure_size=alpha_matting_erode_size
                )
            
            new_im_tensor = pil2tensor(new_im)
            pil_im_tensor = pil2tensor(pil_im)
            processed_images.append(new_im_tensor)
            processed_masks.append(pil_im_tensor)
        
        clear_memory()

        new_ims = torch.cat(processed_images, dim=0)
        new_masks = torch.cat(processed_masks, dim=0)

        return new_ims, new_masks

class BiRefNet_Lite:
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(cls):
        
        return {
            "required": {
                "image": ("IMAGE",),
                "load_local_model": ("BOOLEAN", {"default": True}),
                "device": (["auto", "cuda", "cpu", "mps", "xpu", "meta"],{"default": "auto"}),
                "cutout_func": (["putalpha", "naive", "alpha_matting"],{"default": "putalpha"}),
                "mask_precision_threshold": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 1.0}),
                "cached": ("BOOLEAN", {"default": True}),
                "cpu_size": ("FLOAT",{"default": 0}),
                "alpha_matting_foreground_threshold": ("INT", {"default": 240}),
                "alpha_matting_background_threshold": ("INT", {"default": 10}),
                "alpha_matting_erode_size": ("INT", {"default": 10}),
            },
            "optional": {
                "local_model_path": ("STRING", {"default": model_path}),
            }
        }

    RETURN_TYPES = ("IMAGE", "MASK",)
    RETURN_NAMES = ("image", "mask",)
    FUNCTION = "background_remove"
    CATEGORY = "🔥BiRefNet"
  
    def background_remove(self, 
                          image, 
                          load_local_model,
                          device,
                          cutout_func,
                          mask_precision_threshold,
                          cached,
                          cpu_size,
                          alpha_matting_foreground_threshold,
                          alpha_matting_background_threshold,
                          alpha_matting_erode_size,
                          *args, **kwargs
                          ):
        processed_images = []
        processed_masks = []
       
        device = get_device_by_name(device)
        birefnet = None
        global _birefnet_model
        global _birefnet_model_name
        
        local_model_path = kwargs.get("local_model_path", model_path)
        if cached and _birefnet_model is not None and _birefnet_model_name == local_model_path:
            birefnet = _birefnet_model
        else:
            if _birefnet_model is not None:
                _birefnet_model = None
                _birefnet_model_name = None
                clear_memory()
            
            if load_local_model:
                local_model_path = os.path.join(get_folder_paths('birefnet')[-1], local_model_path)
                # 判断是否开启双卡支持
                # 获取不同显卡的内存可用量
                # 显卡不足时，自动往CPU调度
                spare_params = {}
                config, kwargs = AutoConfig.from_pretrained(
                    local_model_path,
                    return_unused_kwargs=True,
                    trust_remote_code=True,
                    code_revision=None,
                    _commit_hash=None,
                    **kwargs,
                )
                
                class_ref = config.auto_map["AutoModelForImageSegmentation"]
                model_class = get_class_from_dynamic_module(
                    class_ref, local_model_path, code_revision=None, **kwargs
                )
                if cpu_size > 0 and device == 'cuda':
                    spare_params['device_map'] = 'auto'
                    spare_params['max_memory'] = {0: f"{max_gpu_size}GiB", "cpu": f"{cpu_size}GiB"}
                    setattr(model_class, '_no_split_modules', ["Decoder", "SwinTransformer"])
                else:
                    try:
                        if hasattr(model_class, "_no_split_modules"):
                            delattr(model_class, "_no_split_modules")
                    except Exception as e:
                        print('No need to delete:', e)
                    
                AutoModelForImageSegmentation.register(config.__class__, model_class, exist_ok=True)
                model_class = add_generation_mixin_to_remote_model(model_class)
                birefnet = model_class.from_pretrained(
                    local_model_path, config=config, **spare_params
                )
                # birefnet = AutoModelForImageSegmentation.from_pretrained(local_model_path,trust_remote_code=True, **spare_params)
                if cached:
                    _birefnet_model = birefnet
                    _birefnet_model_name = local_model_path
            else:
                birefnet = AutoModelForImageSegmentation.from_pretrained(
                    "ZhengPeng7/BiRefNet", trust_remote_code=True
                )
                # 远程加载不启动分块
                cpu_size = 0
            if cpu_size == 0:
                birefnet.to(device)
        
        for image in image:
            orig_image = tensor2pil(image)
            w,h = orig_image.size
            image = resize_image(orig_image)
            im_tensor = transform_image(image).unsqueeze(0)
            im_tensor=im_tensor.to(device)
            with torch.no_grad():
                result = birefnet(im_tensor)[-1].sigmoid().cpu()
            result = torch.squeeze(F.interpolate(result, size=(h,w)))
            ma = torch.max(result)
            mi = torch.min(result)
            result = (result-mi)/(ma-mi)    
            im_array = (result*255).cpu().data.numpy().astype(np.uint8)
            pil_im = Image.fromarray(np.squeeze(im_array))
            if cutout_func == 'putalpha':
                new_im = putalpha_cutout(orig_image, pil_im)
            elif cutout_func == 'naive':
                new_im = naive_cutout(orig_image, pil_im)
            elif cutout_func == 'alpha_matting':
                new_im = alpha_matting_cutout(
                    orig_image,
                    pil_im,
                    foreground_threshold=alpha_matting_foreground_threshold,
                    background_threshold=alpha_matting_background_threshold,
                    erode_structure_size=alpha_matting_erode_size
                )
            new_im_tensor = pil2tensor(new_im)
            pil_im_tensor = pil2tensor(pil_im)
            if mask_precision_threshold > 0:
                pil_im_tensor[pil_im_tensor <= mask_precision_threshold] = 0.0
                pil_im_tensor[pil_im_tensor > mask_precision_threshold] = 1.0
            
            processed_images.append(new_im_tensor)
            processed_masks.append(pil_im_tensor)
        
        clear_memory()

        new_ims = torch.cat(processed_images, dim=0)
        new_masks = torch.cat(processed_masks, dim=0)

        return new_ims, new_masks


NODE_CLASS_MAPPINGS = {
    "BiRefNet_Super": BiRefNet_Lite,
    "BiRefNet_Lite": BiRefNet_Lite,
    "BiRefNet_onnx": BiRefNet_onnx
}

# A dictionary that contains the friendly/humanly readable titles for the nodes
NODE_DISPLAY_NAME_MAPPINGS = {
    "BiRefNet_Lite": "🔥BiRefNet_Lite",
    "BiRefNet_Super": "🔥BiRefNet_Super",
    "BiRefNet_onnx": "🔥BiRefNet_onnx"
}