from typing import List, Optional, Union
from torchvision import transforms
from PIL import Image
from transformers.image_processing_utils import BaseImageProcessor
from transformers import PretrainedConfig, PreTrainedModel, AutoConfig, AutoImageProcessor, AutoModel
import os
from huggingface_hub import hf_hub_download
import torch
import torch.nn as nn
from transformers.pipelines import PIPELINE_REGISTRY
from transformers.utils import add_end_docstrings
from transformers.pipelines.base import Pipeline, build_pipeline_init_args
class SscdImageProcessor(BaseImageProcessor):
    def __init__(
            self,
            do_resize: bool = True,
            size: int = 288,
            image_mean: Optional[Union[float, List[float]]] = None,
            image_std: Optional[Union[float, List[float]]] = None,
            do_convert_rgb: bool = True,
            **kwargs,
    ) -> None:
        super().__init__(**kwargs)
        self.size = size
        self.image_mean = image_mean if image_mean is not None else [0.485, 0.456, 0.406]
        self.image_std = image_std if image_std is not None else [0.229, 0.224, 0.225]
        self.do_convert_rgb = do_convert_rgb
        self.do_resize = do_resize

    def preprocess(
            self,
            image: Image,
            do_resize: bool = None,
            **kwargs,
    ):
        size_transforms = [
            transforms.ToTensor(),
            transforms.Normalize(
                mean=self.image_mean, std=self.image_std,
            ),
        ]
        if do_resize is None:
            do_resize = self.do_resize
        if do_resize:
            size_transforms.append(transforms.Resize(self.size))
        preprocess = transforms.Compose([
            transforms.Resize(self.size),
            transforms.ToTensor(),
            transforms.Normalize(
                mean=self.image_mean, std=self.image_std,
            ),
        ])
        if self.do_convert_rgb:
            image = image.convert('RGB')
        return preprocess(image).unsqueeze(0)


class SscdConfig(PretrainedConfig):
    model_type = 'sscd-copy-detection'

    def __init__(self, model_path: str = None, **kwargs):
        if model_path is None:
            model_path = 'sscd_disc_mixup.torchscript.pt'
        super().__init__(model_path=model_path, **kwargs)


class SscdModel(PreTrainedModel):
    config_class = SscdConfig

    def __init__(self, config, model_path: str = None):
        super().__init__(config)
        self.dummy_param = nn.Parameter(torch.zeros(0))
        if model_path is None:
            model_path = config.model_path
        is_local = os.path.isdir(config.name_or_path)
        if is_local:
            config.base_path = config.name_or_path
        else:
            file_path = hf_hub_download(repo_id=config.name_or_path, filename=model_path)
            config.base_path = os.path.dirname(file_path)
        model_path = config.base_path + '/' + model_path
        if model_path is not None:
            self.model = torch.jit.load(model_path)

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
        return cls(AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs))

    def forward(self, inputs):
        return self.model(inputs)[0, :]



@add_end_docstrings(build_pipeline_init_args(has_image_processor=True))
class SscdPipeline(Pipeline):
    def __init__(self, model, **kwargs):
        self.device_id = kwargs['device']
        super().__init__(model=model, **kwargs)

    def _sanitize_parameters(self, **kwargs):
        return {}, {}, {}

    def preprocess(self, input):
        return self.image_processor.preprocess(input)

    def _forward(self, inputs):
        return self.model(inputs)

    def postprocess(self, model_outputs):
        return model_outputs


AutoConfig.register('sscd-copy-detection', SscdConfig)
AutoModel.register(SscdConfig, SscdModel)
AutoImageProcessor.register(SscdConfig, slow_image_processor_class=SscdImageProcessor)
models = AutoModel.from_pretrained('m3/sscd-copy-detection')

PIPELINE_REGISTRY.register_pipeline(
    task='sscd-copy-detection',
    pipeline_class=SscdPipeline,
    pt_model=SscdModel
)