import cv2
import numpy as np
from sklearn.preprocessing import normalize


class FaceModel():
    def __init__(self,model_prefix, model_epoch, ctx_id=7 ):
        self.gpu_id=ctx_id
        self.image_size = (112, 112)
        self.model_prefix=model_prefix
        self.model_epoch=model_epoch
        self.model=self._get_model(ctx=ctx_id,image_size=self.image_size,prefix=self.model_prefix,epoch=self.model_epoch,layer='fc1')
    def _get_model(self, ctx, image_size, prefix, epoch, layer):
        pass

    def _getFeatureBlob(self,input_blob):
        pass

    def get_feature(self, image_path):
        image = cv2.imread(image_path)
        image = cv2.resize(image, (112, 112))
        a = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        a = np.transpose(a, (2, 0, 1))
        input_blob = np.expand_dims(a, axis=0)
        emb=self._getFeatureBlob(input_blob)
        emb = normalize(emb.reshape(1, -1))
        #print(type(emb), emb.shape)
        return emb

    def get_batch_feature(self, image_path_list, batch_size=64, flip=0, color="BGR"):
        count = 0
        num_batch =  int(len(image_path_list) / batch_size)
        features = []
        for i in range(0, len(image_path_list), batch_size):

            if count < num_batch:
                tmp_list = image_path_list[i : i+batch_size]
            else:
                tmp_list = image_path_list[i :]
            count += 1

            images = []
            for image_path in tmp_list:
                image = cv2.imread(image_path)
                image = cv2.resize(image, (112, 112))
                if (color == "RGB"):
                    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
                if flip == 1:
                    image = cv2.flip(image, 1)      # flip image horizontally
                # a = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
                a = np.transpose(image, (2, 0, 1))
                images.append(a)
            input_blob = np.array(images)

            emb = self._getFeatureBlob(input_blob)
            features.append(emb)
            print("batch"+str(i))
        features = np.vstack(features)
        features = normalize(features)
        #print('output features:', features.shape)
        return features
