"""
Source url: https://github.com/OPHoperHPO/image-background-remove-tool
Author: Nikita Selin (OPHoperHPO)[https://github.com/OPHoperHPO].
License: Apache License 2.0
"""
import pathlib
from typing import List, Union

import PIL.Image
import numpy as np
import torch
import torchvision.transforms as transforms
from PIL import Image
from pathlib import Path
from io import BytesIO

import requests

import urllib3
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)


from carvekit.ml.arch.tracerb7.efficientnet import EfficientEncoderB7
from carvekit.ml.arch.tracerb7.tracer import TracerDecoder
from carvekit.ml.files.models_loc import tracer_b7_pretrained
from carvekit.utils.image_utils import load_image, convert_image
from carvekit.utils.models_utils import get_precision_autocast, cast_network
from carvekit.utils.pool_utils import thread_pool_processing, batch_generator

from urllib.parse import urlparse  
import requests
import logging
logger = logging.getLogger(__name__)

__all__ = ["TracerUniversalB7"]


class TracerUniversalB7(TracerDecoder):
    """TRACER B7 model interface"""

    def __init__(
        self,
        device="cpu",
        input_image_size: Union[List[int], int] = 640,
        batch_size: int = 4,
        load_pretrained: bool = True,
        fp16: bool = False,
        model_path: Union[str, pathlib.Path] = None,
    ):
        """
        Initialize the TRACER model

        Args:
            layers_cfg: neural network layers configuration
            device: processing device
            input_image_size: input image size
            batch_size: the number of images that the neural network processes in one run
            load_pretrained: loading pretrained model
            fp16: use fp16 precision

        """
        if model_path is None:
            model_path = tracer_b7_pretrained()
        super(TracerUniversalB7, self).__init__(
            encoder=EfficientEncoderB7(),
            rfb_channel=[32, 64, 128],
            features_channels=[48, 80, 224, 640],
        )

        self.fp16 = fp16
        self.device = device
        self.batch_size = batch_size
        if isinstance(input_image_size, list):
            self.input_image_size = input_image_size[:2]
        else:
            self.input_image_size = (input_image_size, input_image_size)

        self.transform = transforms.Compose(
            [
                transforms.ToTensor(),
                transforms.Resize(self.input_image_size),
                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
            ]
        )
        self.to(device)
        if load_pretrained:
            # TODO remove edge detector from weights. It doesn't work well with this model!
            self.load_state_dict(
                torch.load(model_path, map_location=self.device), strict=False
            )
        self.eval()

    def data_preprocessing(self, data: PIL.Image.Image) -> torch.FloatTensor:
        """
        Transform input image to suitable data format for neural network

        Args:
            data: input image

        Returns:
            input for neural network

        """

        return torch.unsqueeze(self.transform(data), 0).type(torch.FloatTensor)

    @staticmethod
    def data_postprocessing(
        data: torch.tensor, original_image: PIL.Image.Image
    ) -> PIL.Image.Image:
        """
        Transforms output data from neural network to suitable data
        format for using with other components of this framework.

        Args:
            data: output data from neural network
            original_image: input image which was used for predicted data

        Returns:
            Segmentation mask as PIL Image instance

        """
        output = (data.type(torch.FloatTensor).detach().cpu().numpy() * 255.0).astype(
            np.uint8
        )
        output = output.squeeze(0)
        mask = Image.fromarray(output).convert("L")
        mask = mask.resize(original_image.size, resample=Image.BILINEAR)
        return mask
    
    def uri_validator(self, url):
        try:
            result = urlparse(url)
            return all([result.scheme, result.netloc])
        except:
            return False
    
        
    def load_image(self, data) -> PIL.Image.Image:
        """
        Loads an image file for other processing
        :param data: Path to image file or PIL image
        :return: image tensor, original pil image
        """
        image_size = 320  # Size of the input and output image for the model
        if isinstance(data, str) or isinstance(data, Path):
            try:
                is_url = self.uri_validator(data)
                if is_url is True:
                    headers = {'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/96.0.4664.110 Safari/537.36'}
                    response = requests.get(data, verify=False, timeout=10, headers=headers)
                    image = Image.open(BytesIO(response.content))   # Load image if there is a path
                else:
                    image = Image.open(data)  # Load image if there is a path

            except IOError:
                raise Exception("load_image fail is not str or path")

        elif isinstance(data, bytes):
            image = Image.open(BytesIO(data))
            image = image.convert("RGB")

        elif isinstance(data, Image):
            image = data.convert("RGB")
        else:
            raise Exception("load_image fail unknown imge type")

        h, w = image.size
        if h < 2 or w < 2:
            raise Exception("Image is too small. Minimum size 2x2")

        return image
    
    def __call__(
        self, images: List[Union[str, pathlib.Path, PIL.Image.Image]]
    ) -> List[PIL.Image.Image]:
        """
        Passes input images though neural network and returns segmentation masks as PIL.Image.Image instances

        Args:
            images: input images

        Returns:
            segmentation masks as for input images, as PIL.Image.Image instances

        """
        collect_masks = []
        autocast, dtype = get_precision_autocast(device=self.device, fp16=self.fp16)
        with autocast:
            cast_network(self, dtype)
            for image_batch in batch_generator(images, self.batch_size):
                converted_images = thread_pool_processing(
                    lambda x: convert_image(load_image(x)), image_batch
                )
                batches = torch.vstack(
                    thread_pool_processing(self.data_preprocessing, converted_images)
                )
                with torch.no_grad():
                    batches = batches.to(self.device)
                    masks = super(TracerDecoder, self).__call__(batches)
                    masks_cpu = masks.cpu()
                    del batches, masks
                masks = thread_pool_processing(
                    lambda x: self.data_postprocessing(
                        masks_cpu[x], converted_images[x]
                    ),
                    range(len(converted_images)),
                )
                collect_masks += masks

        return collect_masks