# Copyright 2021 Zilliz. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import logging
import os
import numpy
from typing import Union, List
from pathlib import Path

import towhee
from towhee.operator.base import NNOperator, OperatorFlag
from towhee.types.arg import arg, to_image_color
from towhee import register
from towhee.models import isc

import torch
from torch import nn
from torchvision import transforms
from PIL import Image as PILImage
import timm

import warnings

warnings.filterwarnings('ignore')
log = logging.getLogger()


@register(output_schema=['vec'])
class Isc(NNOperator):
    """
    The operator uses pretrained ISC model to extract features for an image input.

    Args:
        skip_preprocess (`bool = False`):
            Whether skip image transforms.
    """

    def __init__(self, timm_backbone: str = 'tf_efficientnetv2_m_in21ft1k',
                 skip_preprocess: bool = False, checkpoint_path: str = None, device: str = None) -> None:
        super().__init__()
        if device is None:
            device = 'cuda' if torch.cuda.is_available() else 'cpu'
        self.device = device
        self.skip_tfms = skip_preprocess
        if checkpoint_path is None:
            checkpoint_path = os.path.join(str(Path(__file__).parent), 'checkpoints', timm_backbone + '.pth')

        backbone = timm.create_model(timm_backbone, features_only=True, pretrained=True)
        self.model = isc.create_model(pretrained=True, checkpoint_path=checkpoint_path, device=self.device,
                                      backbone=backbone, p=3.0, eval_p=1.0)
        self.model.eval()

        self.tfms = transforms.Compose([
                transforms.Resize((512, 512)),
                transforms.ToTensor(),
                transforms.Normalize(mean=backbone.default_cfg['mean'],
                                     std=backbone.default_cfg['std'])
        ])

    def __call__(self, data: Union[List[towhee._types.Image], towhee._types.Image]):
        if not isinstance(data, list):
            imgs = [data]
        else:
            imgs = data
        img_list = []
        for img in imgs:
            img = self.convert_img(img)
            img = img if self.skip_tfms else self.tfms(img)
            img_list.append(img)
        inputs = torch.stack(img_list)
        inputs = inputs.to(self.device)
        features = self.model(inputs)
        features = features.to('cpu').flatten(1)

        if isinstance(data, list):
            vecs = list(features.detach().numpy())
        else:
            vecs = features.squeeze(0).detach().numpy()
        return vecs

    @arg(1, to_image_color('RGB'))
    def convert_img(self, img: towhee._types.Image):
        img = PILImage.fromarray(img.astype('uint8'), 'RGB')
        return img


# if __name__ == '__main__':
#     from towhee import ops
#
#     path = 'https://github.com/towhee-io/towhee/raw/main/towhee_logo.png'
#
#     decoder = ops.image_decode.cv2()
#     img = decoder(path)
#
#     op = Isc()
#     out = op(img)
#     assert out.shape == (256,)