#!/usr/bin/env python2
#
# Example to classify faces.
# Brandon Amos
# 2015/10/11
#
# Copyright 2015-2016 Carnegie Mellon University
#
# 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 time

start = time.time()

from Database import Database

import argparse
import cv2
import os
import pickle
import sys

from operator import itemgetter

import numpy as np
np.set_printoptions(precision=2)
import pandas as pd

import openface

from sklearn.pipeline import Pipeline
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis as LDA
from sklearn.preprocessing import LabelEncoder
from sklearn.svm import SVC
from sklearn.model_selection import GridSearchCV
# from sklearn.grid_search import GridSearchCV

from sklearn.mixture import GMM
from sklearn.tree import DecisionTreeClassifier
from sklearn.naive_bayes import GaussianNB

import subprocess

fileDir = os.path.dirname(os.path.realpath(__file__))
modelDir = os.path.join(fileDir, 'batch-represent', 'models')
dlibModelDir = os.path.join(modelDir, 'dlib')
openfaceModelDir = os.path.join(modelDir, 'openface')

def preProcess(imgPath, saves):
    bgrImg = cv2.imread(imgPath)
    if bgrImg is None:
        raise Exception("Unable to load image: {}".format(imgPath))

    rgbImg = cv2.cvtColor(bgrImg, cv2.COLOR_BGR2RGB)

    bb1 = align.getLargestFaceBoundingBox(rgbImg)

    if bb1 == None:
        raise Exception("Unable to find a face: {}".format(imgPath))
        # TODO return a proper value

    reps = []
    # for bb in bbs:
    alignedFace = align.align(
        96,
        rgbImg,
        bb1,
        landmarkIndices=openface.AlignDlib.OUTER_EYES_AND_NOSE)

    if alignedFace is None:
        raise Exception("Unable to align image: {}".format(imgPath))
        # TODO return a proper value

    if (saves):
        # todo
        cv2.imwrite("uploads/preprocessed/0/0.jpg", alignedFace)
    
    return alignedFace


def train(workDir, classifier="LinearSvm", ldaDim=-1):
    print("Loading embeddings.")
    fname = "{}/labels.csv".format(workDir)
    labels = pd.read_csv(fname, header=None).as_matrix()[:, 1]
    labels = list(map(itemgetter(1),
                 map(os.path.split,
                     map(os.path.dirname, labels))))  # Get the directory.
    fname = "{}/reps.csv".format(workDir)
    embeddings = pd.read_csv(fname, header=None).as_matrix()
    le = LabelEncoder().fit(labels)
    labelsNum = le.transform(labels)

    nClasses = len(le.classes_)
    print("Training for {} classes.".format(nClasses))

    print(type(embeddings[0]))

    if classifier == 'LinearSvm':
        clf = SVC(C=1, kernel='linear', probability=True)
    elif classifier == 'GridSearchSvm':
        print("""
        Warning: In our experiences, using a grid search over SVM hyper-parameters only
        gives marginally better performance than a linear SVM with C=1 and
        is not worth the extra computations of performing a grid search.
        """)
        param_grid = [
            {'C': [1, 10, 100, 1000],
             'kernel': ['linear']},
            {'C': [1, 10, 100, 1000],
             'gamma': [0.001, 0.0001],
             'kernel': ['rbf']}
        ]
        clf = GridSearchCV(SVC(C=1, probability=True), param_grid, cv=5)
    elif classifier == 'GMM':  # Doesn't work best
        clf = GMM(n_components=nClasses)

    # ref:
    # http://scikit-learn.org/stable/auto_examples/classification/plot_classifier_comparison.html#example-classification-plot-classifier-comparison-py
    elif classifier == 'RadialSvm':  # Radial Basis Function kernel
        # works better with C = 1 and gamma = 2
        clf = SVC(C=1, kernel='rbf', probability=True, gamma=2)
    elif classifier == 'DecisionTree':  # Doesn't work best
        clf = DecisionTreeClassifier(max_depth=20)
    elif classifier == 'GaussianNB':
        clf = GaussianNB()

    # ref: https://jessesw.com/Deep-Learning/
    elif classifier == 'DBN':
        from nolearn.dbn import DBN
        clf = DBN([embeddings.shape[1], 500, labelsNum[-1:][0] + 1],  # i/p nodes, hidden nodes, o/p nodes
                  learn_rates=0.3,
                  # Smaller steps mean a possibly more accurate result, but the
                  # training will take longer
                  learn_rate_decays=0.9,
                  # a factor the initial learning rate will be multiplied by
                  # after each iteration of the training
                  epochs=300,  # no of iternation
                  # dropouts = 0.25, # Express the percentage of nodes that
                  # will be randomly dropped as a decimal.
                  verbose=1)

    if ldaDim > 0:
        clf_final = clf
        clf = Pipeline([('lda', LDA(n_components=ldaDim)),
                        ('clf', clf_final)])

    clf.fit(embeddings, labelsNum)

    fName = "{}/classifier.pkl".format(workDir)
    print("Saving classifier to '{}'".format(fName))
    with open(fName, 'wb') as f:
        pickle.dump((le, clf), f)


def confineVectorData(returned_output):
    returned_output = (str(returned_output))
    startIndex = returned_output.find('VECTOR__') + 8
    endIndex = returned_output.find('__END') -1
    vectorStr = returned_output[startIndex : endIndex]
    vectorArr = vectorStr.split('|')
    print(len(vectorArr))
    for i in range(0,len(vectorArr)):
        vectorArr[i] = float(vectorArr[i])

    return vectorArr


def infer(img):
    db = Database()
    db.connect_db()

    result = db.where('status', 'is_my_model=0')
    if (len(result) > 0):
        classifierModel = fileDir + "/batch-represent/models/openface/celeb-classifier.nn4.small2.v1.pkl"
    else:
        classifierModel = fileDir + "/batch-represent/reps/classifier.pkl"

    with open(classifierModel, 'rb') as f:
        if sys.version_info[0] < 3:
            (le, clf) = pickle.load(f)
        else:
            (le, clf) = pickle.load(f, encoding='latin1')

    alignedFace = preProcess(img, True)
    
    scriptPath = '/home/hamid/CP/final project/faceRecognition-backend/image_processing_core/batch-represent/src/main.lua'
    dataPath = '/home/hamid/CP/final project/faceRecognition-backend/uploads/preprocessed'
    returned_output = subprocess.check_output(['luajit', scriptPath , '-outDir', 'out', '-data', dataPath, '-infer'])

    rep = confineVectorData(returned_output)
    rep = np.asarray(rep, dtype=np.float32)

    # rep = net.forward(alignedFace)
    # # TODO what is it for?
    # reshapedRep = rep.reshape(1, -1)

    print(rep)


    predictions = clf.predict_proba(rep).ravel()
    maxI = np.argmax(predictions)
    person = le.inverse_transform(maxI)
    confidence = predictions[maxI]

    # print ("RESULT__d-prsn__" + person.decode('utf-8') + "__d-conf__{:.2f}".format(confidence))
    print ("RESULT__d-prsn__" + person + "__d-conf__{:.2f}".format(confidence))

    if isinstance(clf, GMM):
        dist = np.linalg.norm(rep - clf.means_[maxI])
        print("  + Distance from the mean: {}".format(dist))


if __name__ == "__main__":

    parser = argparse.ArgumentParser()

    parser.add_argument(
        '--dlibFacePredictor',
        type=str,
        help="Path to dlib's face predictor.",
        default=os.path.join(
            dlibModelDir,
            "shape_predictor_68_face_landmarks.dat"))
    parser.add_argument(
        '--networkModel',
        type=str,
        help="Path to Torch network model.",
        default=os.path.join(
            openfaceModelDir,
            'nn4.small2.v1.t7'))
            
    parser.add_argument('--cuda', action='store_true')

    subparsers = parser.add_subparsers(dest='mode', help="Mode")
    trainParser = subparsers.add_parser('train',
                                        help="Train a new classifier.")

    inferParser = subparsers.add_parser(
        'infer', help='Predict who an image contains from a trained classifier.')

    inferParser.add_argument(
        'img',
        type=str,
        help="")

    args = parser.parse_args()

    align = openface.AlignDlib(args.dlibFacePredictor)
    net = openface.TorchNeuralNet(args.networkModel, imgDim=96,
                                  cuda=args.cuda)

    if args.mode == 'train':
        train(fileDir + '/batch-represent/reps')
    elif args.mode == 'infer':
        infer(args.img)
    else:
        print("bad mode")

