import os
import cv2
import numpy as np
from keras.models import Model
from keras.models import load_model
from sklearn.externals import joblib
from tensorflow.keras.preprocessing.image import img_to_array

os.environ["CUDA_VISIBLE_DEVICES"] = '0'
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'


class FeatureExtractor(object):
    def __init__(self, input_size=224, out_put_layer='avg_pool', model_path='FerPlus3.h5'):
        self.model = load_model(model_path)
        self.input_size = input_size
        self.model_inter = Model(inputs=self.model.input, outputs=self.model.get_layer(out_put_layer).output)

    def inference(self, image):
        image = cv2.resize(image, (self.input_size, self.input_size))
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        image = image.astype("float") / 255.0
        image = img_to_array(image)
        image = np.expand_dims(image, axis=0)
        feature = self.model_inter.predict(image)[0]
        return feature


def features2feature(pics_features):

        pics_features = np.array(pics_features)
        fea_mean = pics_features.mean(axis=0)
        fea_max = np.amax(pics_features, axis=0)
        fea_min = np.amin(pics_features, axis=0)
        fea_std = pics_features.std(axis=0)

        return np.concatenate((fea_mean, fea_max, fea_min, fea_std), axis=1).reshape(1, -1)


def start_filter(config):

    cls_emotion_path = config['MODEL']['CLS_EMOTION']
    face_feature_dir = config['VIDEO']['FACE_FEATURE_DIR']
    frame_list_dir = config['VIDEO']['FRAME_LIST_DIR']
    result_file_name = config['EMOTION']['RESULT_FILE']

    svm_clf = joblib.load(cls_emotion_path)

    result_file_path = os.path.join(frame_list_dir, result_file_name)
    result_file = open(result_file_path, 'w')

    face_feature_names = os.listdir(face_feature_dir)
    for face_feature in face_feature_names:
        face_feature_path = os.path.join(face_feature_dir, face_feature)

        features_np = np.load(face_feature_path, allow_pickle=True)

        feature = features2feature(features_np)
        res = svm_clf.predict_proba(feature)
        proba = np.squeeze(res)
        # class_pre = svm_clf.predict(feature)

        result_file.write(face_feature[:-4] + ' ')
        result_file.write(str(proba[0]) + ',' + str(proba[1]) + ',' + str(proba[2]) + '\n')

    result_file.close()