#!/usr/bin/env python2
#
# 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 os
import sys
fileDir = os.path.dirname(os.path.realpath(__file__))
sys.path.append(os.path.join(fileDir, "..", ".."))
import copy
import txaio
txaio.use_twisted()

from autobahn.twisted.websocket import WebSocketServerProtocol, \
    WebSocketServerFactory
from twisted.internet import task, defer
from twisted.internet.ssl import DefaultOpenSSLContextFactory

from twisted.python import log

import argparse
import cv2
import imagehash
import json
from PIL import Image
import numpy as np
import os
import StringIO
import urllib
import base64

from sklearn.decomposition import PCA
from sklearn.grid_search import GridSearchCV
from sklearn.manifold import TSNE
from sklearn.svm import SVC
from sklearn.cluster import KMeans

import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.cm as cm

import openface
import pickle, pprint
import json
from numpy import genfromtxt
import requests
import pandas as pd


modelDir = os.path.join(fileDir, '..', '..', 'models')
dlibModelDir = os.path.join(modelDir, 'dlib')
openfaceModelDir = os.path.join(modelDir, 'openface')
trainDir=os.path.join(fileDir,'..','..','Train_Image')
classifierDir=os.path.join(fileDir,'..','..')
imagesDir=os.path.join(fileDir,'..','..')
peopleDir=os.path.join(fileDir,'..','..')
featureDir=os.path.join(fileDir,'..','..','Feature_gui')
Feature_unknown=os.path.join(fileDir,'..','..','Feature_unknown')
alignDir=os.path.join(fileDir,'..','..')
# For TLS connections
tls_crt = os.path.join(fileDir, 'tls', 'server.crt')
tls_key = os.path.join(fileDir, 'tls', 'server.key')

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('--imgDim', type=int,
                    help="Default image dimension.", default=96)
parser.add_argument('--cuda', action='store_true')
parser.add_argument('--unknown', type=bool, default=False,
                    help='Try to predict unknown people')
parser.add_argument('--port', type=int, default=9000,
                    help='WebSocket Port')

args = parser.parse_args()

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



def loadDist():
    try:
        with open(featureDir+'/classifier.pkl', 'rb') as f:
            # if sys.version_info[0] < 3:
            if sys.version_info[0] < 3:
                (le, clf) = pickle.load(f)
                return (le,clf)
            else:
                (le, clf) = pickle.load(f, encoding='latin1')
                return (le,clf)
    except Exception as e:
        return None

def loadUnknown():
    try:
        with open(Feature_unknown+'/classifier.pkl', 'rb') as f:
            # if sys.version_info[0] < 3:
            if sys.version_info[0] < 3:
                (le, clf) = pickle.load(f)
                return (le,clf)
            else:
                (le, clf) = pickle.load(f, encoding='latin1')
                return (le,clf)
    except Exception as e:
        return None




def retrain():
    os.chdir(alignDir)
    os.remove(alignDir+'/Aligned_data'+'/cache.t7')
    os.system('python align-dlib.py'+' '+trainDir+' '+'align outerEyesAndNose Aligned_data/ --size 96')
    os.system('./batch-represent/main.lua -outDir Feature_gui/ -data Aligned_data/')
    os.system('python classifier.py train Feature_gui/ --classifier RadialSvm')
    os.remove(alignDir+'/Aligned_data_unknown'+'/cache.t7')
    os.system('cp -r Feature_gui/classifier.pkl Feature_dir/')
    os.system('cp -r  Unknown/ Train_Image/')
    os.system('python align-dlib.py Train_Image/ align outerEyesAndNose Aligned_data_unknown/ --size 96')
    os.system('./batch-represent/main.lua -outDir Feature_unknown/ -data Aligned_data_unknown/')
    os.system('python classifier.py train Feature_unknown/ --classifier RadialSvm')
    os.remove(alignDir+'/Train_Image/'+'Unknown/')


    os.chdir(fileDir)

    try:
        with open(self.featureDir+'/classifier.pkl', 'rb') as f:
        # if sys.version_info[0] < 3:
            if sys.version_info[0] < 3:
                (le, clf) = pickle.load(f)
                return (le,clf)
            else:
                (le, clf) = pickle.load(f, encoding='latin1')
                return (le,clf)
    except Exception as e:
        return None





def loadImages():
    try:
        with open('images.pkl', 'rb') as f:
            # if sys.version_info[0] < 3:
            images = pickle.load(f)
        return images
    except Exception as e:
        return {}


#my_data = genfromtxt('people.csv', delimiter=',')
def loadModel():
    # model = open('model.pkl', 'r')
    # svm_persisted = pickle.load('model.pkl')
    # output.close()
    # return svm_persisted
    # return True
    try:
        with open('model.pkl', 'rb') as f:
            # if sys.version_info[0] < 3:
            mod = pickle.load(f)
            return mod
    except Exception as e:
        return None

def loadPeople():
    try:
        with open('people.pkl', 'rb') as f:
            mod = pickle.load(f)
            return mod
    except Exception as e:
        return []

class Face:

    def __init__(self, rep, identity):
        self.rep = rep
        self.identity = identity

    def __repr__(self):
        return "{{id: {}, rep[0:5]: {}}}".format(
            str(self.identity),
            self.rep[0:5]
        )


class OpenFaceServerProtocol(WebSocketServerProtocol):
    def __init__(self):
        super(OpenFaceServerProtocol, self).__init__()
        self.images = {}
        self.training = True
        self.people = []
        self.svm = None
        (self.le_dist,self.clf_dist)=loadDist()
        (self.le_unknown,self.clf_unknown)=loadUnknown()
        self.centroids_dist=pd.read_csv(featureDir+'/centroids_csv.csv')
        self.centroids_dist.drop(self.centroids_dist.columns[0],axis=1,inplace=True)
        self.embeddings_dist=pd.read_csv(featureDir+'/embeddings.csv')
        self.embeddings_unknown=pd.read_csv(Feature_unknown+'/embeddings.csv')
        self.centroids_unknown=pd.read_csv(Feature_unknown+'/centroids_csv.csv')
        self.centroids_unknown.drop(self.centroids_unknown.columns[0],axis=1,inplace=True)
        self.KMeans=None
        #self.centroids.drop(self.centroids.columns[0],axis=1,inplace=True)
        #self.retrain=False
        self.classifier="Distance"
        if args.unknown:
            self.unknownImgs = np.load("./examples/web/unknown.npy")

    def onConnect(self, request):
        print("Client connecting: {0}".format(request.peer))
        self.training = True

    def onOpen(self):
        print("WebSocket connection open.")

    def onMessage(self, payload, isBinary):
        raw = payload.decode('utf8')
        msg = json.loads(raw)
        print("Received {} message of length {}.".format(
            msg['type'], len(raw)))
        if msg['type'] == "ALL_STATE":
            self.loadState(msg['images'], msg['training'], msg['people'])
        elif msg['type'] == "NULL":
            self.sendMessage('{"type": "NULL"}')
        elif msg['type'] == "FRAME":
            self.processFrame(msg['dataURL'], msg['identity'])
            self.sendMessage('{"type": "PROCESSED"}')
        elif msg['type'] == "TRAINING":
            self.training = msg['val']
            #if not self.training:
            #    self.trainSVM()
        elif msg['type'] == "ADD_PERSON":
            if msg['val'].encode('ascii','ignore') not in self.people:


               self.people.append(msg['val'].encode('ascii', 'ignore'))
               self.people=self.people
            #np.savetxt("people.csv", self.people, delimiter=",")
               with open(peopleDir+'/people.pkl', 'w') as f:
                    pickle.dump(self.people, f)

            print(self.people)
            if not os.path.isdir(trainDir+"/"+msg['val'].encode('ascii','ignore')):
                os.mkdir(trainDir+"/"+msg['val'].encode('ascii','ignore'))
            os.chdir(trainDir+"/"+msg['val'].encode('ascii','ignore'))

        elif msg['type'] == "UPDATE_IDENTITY":
            h = msg['hash'].encode('ascii', 'ignore')
            if h in self.images:
                self.images[h].identity = msg['idx']
                if not self.training:
                    self.trainSVM()
            else:
                print("Image not found.")
        elif msg['type'] == "REMOVE_IMAGE":
            h = msg['hash'].encode('ascii', 'ignore')
            if h in self.images:
                del self.images[h]
                if not self.training:
                    self.trainSVM()
            else:
                print("Image not found.")
        elif msg['type'] == 'REQ_TSNE':
            print "Received message TSNE"
            self.sendTSNE(msg['people'])
        elif msg['type']== 'RE-TRAIN':
            print self.classifier
            self.retrain()
        elif msg['type']=='DISTANCE':
            #(self.le,self.clf)=loadDist()
            self.classifier="Distance"
            self.centroids=pd.read_csv(featureDir+'/centroids_csv.csv')
            self.centroids.drop(self.centroids.columns[0],axis=1,inplace=True)

            print self.classifier
        elif msg['type']=="UNKNOWN":
            #(self.le,self.clf)=loadUnknown()
            self.classifier="Unknown"
            self.centroids=pd.read_csv(Feature_unknown+'/centroids_csv.csv')
            self.centroids.drop(self.centroids.columns[0],axis=1,inplace=True)
            print self.classifier


        else:
            print("Warning: Unknown message type: {}".format(msg['type']))

    def onClose(self, wasClean, code, reason):
        print("WebSocket connection closed: {0}".format(reason))

    def loadState(self, jsImages, training, jsPeople):
        self.training = training

        for jsImage in jsImages:
            h = jsImage['hash'].encode('ascii', 'ignore')
            self.images[h] = Face(np.array(jsImage['representation']),
                                  jsImage['identity'])

        for jsPerson in jsPeople:
            self.people.append(jsPerson.encode('ascii', 'ignore'))

        if not training:
            self.trainSVM()

    def getData(self):
        X = []
        y = []
        for img in self.images.values():
            X.append(img.rep)
            y.append(img.identity)

        numIdentities = len(set(y + [-1])) - 1
        if numIdentities == 0:
            return None

        if args.unknown:
            numUnknown = y.count(-1)
            numIdentified = len(y) - numUnknown
            numUnknownAdd = (numIdentified / numIdentities) - numUnknown
            if numUnknownAdd > 0:
                print("+ Augmenting with {} unknown images.".format(numUnknownAdd))
                for rep in self.unknownImgs[:numUnknownAdd]:
                    # print(rep)
                    X.append(rep)
                    y.append(-1)

        X = np.vstack(X)
        y = np.array(y)
        return (X, y)

    def sendTSNE(self, people):
        #if self.classifier=='Distance':
        #    d=self.embeddings_dist
        #else:
        #    d=self.embeddings_unknown
        d = pd.read_csv(featureDir+'/embeddings.csv')
        if d is None:
            return
        else:
            print d.columns
            (X, y) = (copy.copy(d),copy.copy(d['labels']))
            X.drop('labels',axis=1, inplace=True)

        X_pca = PCA(n_components=50).fit_transform(X, X)