# Copyright 2019 The Vearch Authors. 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.
# ==============================================================================

import config
import util
import exceptions


class Base(object):
    pass


class Face(Base):

    def __init__(self):
        super(Face, self).__init__()
        model_name = config.face_config['modelname']
        self.face = util.get_model(model_name)
        self.model = self.face.load_model(config.face_config)

    def pre_process(self, image, boundingbox):
        if boundingbox:
            bbox = list(map(int, boundingbox.split(',')))
            image = util.crop(image, bbox)
        return image

    def encode(self, url):
        try:
            image = util.read_image(url)
        except Exception as err:
            raise exceptions.ImageError(url)
        # image = self.pre_process(image)
        res = self.model.encode(image)
        return res


class ImageSearch(Base):

    def __init__(self):
        super(ImageSearch, self).__init__()
        model_name = config.image_config['modelname']
        detect_name = config.image_config['detectname']
        self.extract_model = util.get_model(model_name).load_model()
        self.detect_model = util.get_model(detect_name).load_model() if detect_name else None

    def pre_process(self, image):
        if self.detect_model:
            bbox = self.detect_model.detect(image)
            image = util.crop(image, bbox)
        return image

    def encode(self, url):
        try:
            image = util.read_image(url)
        except Exception as err:
            raise exceptions.ImageError(f'read {url} failed!')
        image = self.pre_process(image)
        feat = self.extract_model.forward(image)
        assert len(feat) > 0, 'No detect object.'
        return feat[0]


class Text(Base):

    def __init__(self):
        super(Text, self).__init__()
        model_name = config.text['modelname']
        text = util.get_model(model_name)
        self.model = text.load_model(config.text)

    def encode(self, text):
        return self.model.encode([text])[0]


def load_model(model_name):
    if model_name == 'face_retrieval':
        model = Face()
    elif model_name == 'image_retrieval':
        model = ImageSearch()
    elif model_name == 'text':
        model = Text()
    elif model_name == 'audio':
        raise NotImplementedError()
    else:
        raise exceptions.LoadModelError(f'{model_name} is not existed')

    return model