import os
import argparse
from collections import defaultdict
import dlib
import cv2#
#from sklearn.externals import joblib
import joblib

import numpy as np
import pandas as pd
from sklearn import metrics
from pipeline import face_utils
from sklearn.metrics import precision_recall_fscore_support as score
from sklearn import metrics
import matplotlib.pyplot as plt
import pickle 

def parse_args():
    """Parses input arguments."""
    parser = argparse.ArgumentParser()
    parser.add_argument('-i', '--input', dest='input',
                        help='Path to input image or folder containting multiple images.')
    parser.add_argument('-o', '--output', dest='output', help='Path to save outputs.',
                        default='./output')
    parser.add_argument('-p', '--pipeline', choices=['gan', 'deepfake', 'face2face'], dest='pipeline')
    parser.add_argument('-f', '--features', dest='save_features', action='store_true',
                        help='Set flag to save features, e.g. to fit classifier.',
                        default=False)
    args = parser.parse_args()
    return args


def load_classifiers(pipeline):
    """Loads classifiers for specified pipeline."""
    classifiers = None
    if pipeline == 'gan':
        classifiers = joblib.load('models/gan/bagging_knn.pkl')
    elif pipeline == 'deepfake':
    
        classifier_mlp = joblib.load('models/deepfake/mlp_df.pkl')
        classifier_logreg = joblib.load('models/deepfake/logreg_df.pkl')
        classifiers = [classifier_mlp, classifier_logreg]
    elif pipeline == 'face2face':
        classifier_mlp = joblib.load('models/face2face/mlp_f2f.pkl')
        classifier_logreg = joblib.load('models/face2face/logreg_f2f.pkl')
        classifiers = [classifier_mlp, classifier_logreg]
    else:
        print ('Unknown pipeline argument.')
        exit(-1)

    return classifiers


def load_facedetector():
    """Loads dlib face and landmark detector."""
    # download if missing http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2
    if not os.path.isfile('shape_predictor_68_face_landmarks.dat'):
        print ('Could not find shape_predictor_68_face_landmarks.dat.')
        exit(-1)
    face_detector = dlib.get_frontal_face_detector()
    sp68 = dlib.shape_predictor('shape_predictor_68_face_landmarks.dat')

    return face_detector, sp68


def main(input_path, output_path, pipeline, save_features):
    """Main function to process input files with selected pipeline.
    Given a path to single image or folder and a output path,
    images are processed with selected pipeline.
    Outputs are saved as .csv file. If selected, the computed feature vectors
    for classification are saved as .npy.
    The scores.csv file contains the output score of the different classifiers. The
    'Valid' value indicates if the face detection and segmentation was successful.
    Args:
        input_path: Path to image, or folder containing multiple images.
        output_path: Path to save outputs.
        pipeline: Selected pipeline for processing. Options: 'gan', 'deepfake', 'face2face'
        save_features: Boolean flag. If set true, feature vectors will be saved as single .npy file.
    """
    # transfer image names to list
    pathorg = input_path
    result_dict = defaultdict(list)
    result_dict_labels = defaultdict(list)
    feature_vec_list = []
    score_clfList = []
    score_hsvList = []
    labs = []

    for pth in os.listdir(pathorg):
        input_path = os.path.join(pathorg,pth)
        print(input_path)
        if os.path.isdir(input_path):
            file_list = [name for name in os.listdir(input_path) if os.path.isfile(os.path.join(input_path, name))]
        else:
            file_list = [os.basename(input_path)]
            input_path = os.dirname(input_path)

        if len(file_list) == 0:
            print ('No files at given input path.')
            exit(-1)

        # load classifiers, sanity check
        classifiers = load_classifiers(pipeline)

        # setup face detector and landmark predictor
        face_detector, sp68 = load_facedetector()

        # create save folder
        if not os.path.exists(output_path):
            os.makedirs(output_path)

        files_processed = 0

        for input_file in file_list:
            # load image
            img = cv2.imread(os.path.join(input_path, input_file))
            if img is None or img is False:
                print ("Could not open image file: %s" % os.join(input_path, input_file))
                continue
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

            if pipeline == 'gan':
                # detect and crop faces
                face_crops, final_landmarks = face_utils.get_crops_landmarks(face_detector, sp68, img)
                # call feature extraction, classifier pipeline
                from pipeline import eyecolor
                score_clf, score_HSV, feature_vector, valid_seg = eyecolor.process_faces(classifiers,
                                                                                         face_crops,
                                                                                         final_landmarks,
                                                                                         scale=768)
                # add results to output dict
                result_dict['Filename'].append(input_file)
                result_dict['Score_kNN'].append(score_clf)
                result_dict['Score_HSV'].append(score_HSV)
                result_dict['Valid'].append(int(valid_seg))
                if int(valid_seg) == 0:
                    score_clfList.append(score_clf)
                    score_hsvList.append(score_clf)
                else:
                    score_clfList.append(score_clf)
                    score_hsvList.append(score_clf)
                
                result_dict_labels['Filename'].append(input_file)
                if pth == 'REAL0':
                    result_dict_labels['Label'].append(int(0))
                    labs.append(0)
                else:
                    result_dict_labels['Label'].append(int(1))
                    labs.append(1)
                    
                if save_features:
                    # add feature vector to list, take care might be None
                    if feature_vector is None:
                        feature_vector = np.full(6, np.nan)
                    feature_vec_list.append(feature_vector)

            elif pipeline == 'deepfake' or pipeline == 'face2face':
                if pipeline == 'face2face':
                    extend_roi = 0.1
                else:
                    extend_roi = 0.0
                # detect and crop faces
                face_crops, final_landmarks = face_utils.get_crops_landmarks(face_detector, sp68, img,
                                                                             roi_delta=extend_roi)
                # call feature extraction, classifier pipeline
                from pipeline import texture
                scores, feature_vectors, valid_seg = texture.process_faces(classifiers,
                                                                           face_crops,
                                                                           final_landmarks,
                                                                           pipeline,
                                                                           scale=256)
                if int(valid_seg) == 1:
                    score_clfList.append(scores[0])
                    score_hsvList.append(scores[1])
                    if pth == 'REAL0':
                        result_dict_labels['Label'].append(int(0))
                        labs.append(0.0)
                    else:
                        result_dict_labels['Label'].append(int(1))
                        labs.append(1.0)



                # add results to output dict
                result_dict['Filename'].append(input_file)
                result_dict['Score_MLP'].append(scores[0])
                result_dict['Score_LogReg'].append(scores[1])
                result_dict['Valid'].append(int(valid_seg))

                result_dict_labels['Filename'].append(input_file)
                if pth == 'REAL0':
                    result_dict_labels['Label'].append(int(0))
                else:
                    result_dict_labels['Label'].append(int(1))

                if save_features:
                    # add feature vector to list, save invalid as nan
                    if feature_vectors[0] is None:
                        feature_vector = np.full(18, np.nan)
                    else:
                        feature_vector = feature_vectors[0]
                    feature_vec_list.append(feature_vector)

            files_processed += 1
            print ("Files processed: ", files_processed, " of ", len(file_list))



    tol_pred = []
    for i in range(len(score_hsvList)):
        if labs[i] == 0.0:
            if score_hsvList[i] < 0.5:
                tol_pred.append(0.0)
            else:
                tol_pred.append(1.0)
        else:
            if score_hsvList[i] > 0.5:
                tol_pred.append(1.0)
            else:
                tol_pred.append(0.0)
    
    precision, recall, fscore, support = score(labs, tol_pred)
    print('precision: {}'.format(precision))
    print('recall: {}'.format(recall))
    print('fscore: {}'.format(fscore))
    print('support: {}'.format(support))
    print(metrics.accuracy_score(labs, tol_pred))
    
    # save outputs
    tol_pred = []
    for i in range(len(score_clfList)):
        if labs[i] == 0.0:
            if score_clfList[i] < 0.5:
                tol_pred.append(0.0)
            else:
                tol_pred.append(1.0)
        else:
            if score_clfList[i] > 0.5:
                tol_pred.append(1.0)
            else:
                tol_pred.append(0.0)
    
    precision, recall, fscore, support = score(labs, tol_pred)
    print('precision: {}'.format(precision))
    print('recall: {}'.format(recall))
    print('fscore: {}'.format(fscore))
    print('support: {}'.format(support))
    print(metrics.accuracy_score(labs, tol_pred))


    y_true = labs
    y_pred = tol_pred
    result = metrics.classification_report(y_true, y_pred, labels=None, target_names=None, sample_weight=None, digits=4, output_dict=False, zero_division='warn')    
    print(result)
    print(f'ACC is {metrics.accuracy_score(y_true, y_pred)}')
    y_true, y_pred = np.array(y_true),np.array(y_pred)
    print(y_true.shape, y_pred.shape)
    fpr = dict()
    tpr = dict()
    roc_auc = dict()
    for i in range(2):
        fpr[i], tpr[i], _ = metrics.roc_curve(y_true,y_pred)
        roc_auc[i] = metrics.auc(fpr[i], tpr[i])

    # Compute micro-average ROC curve and ROC area
    fpr["micro"], tpr["micro"], _ = metrics.roc_curve(y_true,y_pred)
    roc_auc["micro"] = metrics.auc(fpr["micro"], tpr["micro"])

    plt.figure()
    lw = 2
    print(roc_auc[1])
    plt.plot(fpr[1], tpr[1], color='darkred', lw=lw, label='ROC curve ({:.2f})'.format(roc_auc[1]))
    plt.plot([0, 1], [0, 1], color='navy', lw=lw, linestyle='--')
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title('EfficientNet-B0')
    plt.legend(loc="lower right")
    plt.show()




    result_df = pd.DataFrame.from_dict(result_dict, orient="columns")
    result_df_labels = pd.DataFrame.from_dict(result_dict_labels, orient="columns")
    if pipeline == 'gan':
        result_df.to_csv(os.path.join(output_path, "scores.csv"), sep=',',
                         columns=['Filename', 'Score_kNN', 'Score_HSV', 'Valid'], index=False)
    elif pipeline == 'deepfake' or pipeline == 'face2face':
        result_df.to_csv(os.path.join(output_path, "scores.csv"), sep=',',
                         columns=['Filename', 'Score_MLP', 'Score_LogReg', 'Valid'], index=False)

    if save_features:
        result_df_labels.to_csv(os.path.join(output_path, "labels.csv"), sep=',',
                         columns=['Filename', 'Label'], index=False)
        feature_vec_list = np.asarray(feature_vec_list, dtype=np.float32)
        np.save(os.path.join(output_path, 'features.npy'), feature_vec_list)


if __name__ == '__main__':
    args_in = parse_args()
    main(args_in.input, args_in.output, args_in.pipeline, args_in.save_features)