from .imagefunc import *

import torch.nn as nn
from torchvision import transforms
from .BiRefNet_legacy.baseline import BiRefNet
from .BiRefNet_legacy.config import Config

class BiRefNet_img_processor:
    def __init__(self, config):
        self.config = config
        self.data_size = (config.size, config.size)
        self.transform_image = transforms.Compose([
            transforms.Resize(self.data_size),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ])

    def __call__(self, _image: np.array):
        _image_rs = cv2.resize(_image, (self.config.size, self.config.size), interpolation=cv2.INTER_LINEAR)
        _image_rs = Image.fromarray(np.uint8(_image_rs*255)).convert('RGB')
        image = self.transform_image(_image_rs)
        return image

class BiRefNetRemoveBackground:
    def __init__(self):
        self.ready = False

    def load(self, weight_path, device):
        # load model
        self.model = BiRefNet()
        state_dict = torch.load(weight_path, map_location='cpu')
        unwanted_prefix = '_orig_mod.'
        for k, v in list(state_dict.items()):
            if k.startswith(unwanted_prefix):
                state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
        self.model.load_state_dict(state_dict)
        self.model = self.model.to(device)
        self.model.eval()
        # load processor
        self.processor = BiRefNet_img_processor(Config())
        self.ready = True

  
    def generate_mask(self, image:Image) -> Image:

        if torch.backends.mps.is_available():
            device = "mps"
        elif torch.cuda.is_available():
            device = "cuda"
        else:
            device = "cpu"

        if not self.ready:
            model_folder_name = 'BiRefNet'
            model_name = 'BiRefNet-ep480.pth'
            model_file_path = ""
            try:
                model_file_path = os.path.join(
                    os.path.normpath(folder_paths.folder_names_and_paths[model_folder_name][0][0]), model_name)
            except:
                pass
            if not os.path.exists(model_file_path):
                model_file_path = os.path.join(folder_paths.models_dir, model_folder_name, model_name)
            self.load(model_file_path, device=device)

        i = pil2tensor(image)
        orig_image = image.convert('RGB')
        np_image = i.squeeze().numpy()
        img = self.processor(np_image)
        inputs = img[None, ...].to(device)
        with torch.no_grad():
            scaled_preds = self.model(inputs)[-1].sigmoid()
        _mask = nn.functional.interpolate(scaled_preds[0].unsqueeze(0),
                                          size=np_image.shape[:2],
                                          mode='bilinear',
                                          align_corners=True
                                          )[0]

        brightness_image = ImageEnhance.Brightness(tensor2pil(_mask))

        return brightness_image.enhance(factor=1.01)