# 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 numpy
import towhee
import sys
from pathlib import Path
from torchvision import transforms
from towhee.types.image_utils import to_pil

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

@register(output_schema=['vec'])
class Dolg(NNOperator):
    """
    DOLG Embedding Operator
    """
    def __init__(self, img_size, input_dim, hidden_dim, output_dim):
        super().__init__()
        sys.path.append(str(Path(__file__).parent))
        from dolg_impl import DolgNet
        self.model = DolgNet(img_size, input_dim, hidden_dim, output_dim)
        normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225])
        self.tfms = transforms.Compose([transforms.Resize([img_size, img_size]),
                                        transforms.Scale([img_size, img_size]),
                                        transforms.ToTensor(),
                                        normalize])

    @arg(1, to_image_color('RGB'))
    def __call__(self, img: numpy.ndarray):
        img = self.tfms(to_pil(img)).unsqueeze(0)
        self.model.eval()
        features = self.model(img)
        feature_vector = features.flatten().detach().numpy()
        return feature_vector