import numpy as np
import pkg_resources
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.sequence import pad_sequences
import pickle

class IRIS():
    def __init__(self):
        # Initialise emotion classes and index
        self.classes = ['joy', 'fear', 'anger', 'sadness', 'love', 'surprise']
        self.class_to_index = dict((c,i) for i, c in enumerate(self.classes))
        self.index_to_class = dict((v, k) for k, v in self.class_to_index.items())

        # Initialize SA-Model-Final-v8
        self.model_path = pkg_resources.resource_filename('iris_emotion', 'SA_Model_Final_v8/')
        self.model= load_model(self.model_path)

        # Initialize Tokenizer
        self.tk_path = pkg_resources.resource_filename('iris_emotion', 'tokenizer.pickle')
        with open(self.tk_path, 'rb') as handle:
            self.tokenizer = pickle.load(handle)

    def get_sequences(self, msg):
        # Preprocess message
        maxlen = 50
        sequences = self.tokenizer.texts_to_sequences([msg])
        padded = pad_sequences(sequences, truncating='post' , padding='post', maxlen = maxlen)

        # Return tokenized message
        return padded

    def get_emotion(self, msg):
        # Tokenize message
        msg_seq = self.get_sequences(msg)

        # Predict emotion
        p = self.model.predict(msg_seq)[0]
        pred_class = self.index_to_class[np.argmax(p).astype('uint8')]

        # Return predicted emotion
        return pred_class


# coding: UTF-8
import numpy as np
import cv2
import sys
import time
import os
import subprocess
import ConfigParser

class TemplateImage():
    def __init__(self, path):
        self.name = path
        self.img = cv2.imread(path, 0)
        self.w, self.h = self.img.shape[::-1]
        self.last_result = None

    def match(self, img):
        res = cv2.matchTemplate(img, self.img, eval("cv2.TM_CCOEFF_NORMED"))
        min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(res)
        top_left = max_loc
        bottom_right = (top_left[0] + self.w, top_left[1] + self.h)
        self.last_result = (top_left, bottom_right), max_val
        return self.last_result

    def imshow(self, name = ""):
        if name == "":
            name = self.name
        cv2.imshow(name, self.img)

class VideoLoader():
    def __init__(self, path):
        self.name = path
        self.cap = cv2.VideoCapture(path)
        self.frame_count = int(self.cap.get(cv2.cv.CV_CAP_PROP_FRAME_COUNT))
        self.fps = self.cap.get(cv2.cv.CV_CAP_PROP_FPS)
        return

    def __del__(self):
        self.cap.release()

    def next(self):
        return self.cap.read()

    def curpos(self):
        return int(self.cap.get(cv2.cv.CV_CAP_PROP_POS_FRAMES))

    def seek_frame(self, num):
        result = True
        frame = self.curpos() + num
        if frame > self.frame_count:
            frame = self.frame_count
            result = False
        self.cap.set(cv2.cv.CV_CAP_PROP_POS_FRAMES, frame)
        return result

    def available(self):
        return self.cap.isOpened()

class CUIProgressBar():
    def __init__(self, width=20):
        self.width = width

    def update(self, ratio):
        hashes = '#' * int(round(ratio * self.width))
        spaces = ' ' * (self.width - len(hashes))
        sys.stdout.write("\rProgress : [{0}] {1}%".format(hashes + spaces, int(round(ratio * 100))))
        sys.stdout.flush()

class SubProcessCut():
    def __init__(self, file_path):
        self.file_path = file_path
        self.p = None

    def run(self, begin_sec, length_sec, output_file_name):
        self.command = [
            'ffmpeg',
            '-y',
            '-ss', str(begin_sec),
            '-i', self.file_path,
            '-t', str(length_sec),
            '-c:v', 'copy',
            '-c:a', 'copy',
            output_file_name
        ]
        print "popen", self.command
        with open("stderr.log", "wb") as logfile:
            self.p = subprocess.Popen(self.command, stdout=subprocess.PIPE, stderr=logfile, shell=False)

    def wait(self):
        if self.p:
            self.p.wait()

def main(preset_path, movie_file):
    conf = ConfigParser.SafeConfigParser()
    conf.read(os.path.join(preset_path, 'config.txt'))
    start_offset = int(conf.get('offset', 'start'))
    end_offset = int(conf.get('offset', 'end'))
    video = VideoLoader(movie_file)
    start = TemplateImage(os.path.join(preset_path, 'start.png'))
    end = TemplateImage(os.path.join(preset_path, 'end.png'))
    start_th = float(conf.get('thratio', 'start'))
    end_th = float(conf.get('thratio', 'end'))
    behind_start = int(conf.get('option', 'behind_start'))

    if not video.available():
        print "video load error"
        raw_input()
        return

    between = False
    section_list = []
    section_start = 0
    section_end = 0
    bar = CUIProgressBar()

    cutman = None
    number = 1

    #cv2.namedWindow("image", cv2.WINDOW_NORMAL)
    while(video.available()):
        ret, frame = video.next()

        if not ret:
            break
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

        if behind_start or not between:
            rect, max_value = start.match(gray)
            print "s", max_value,
            if max_value >= start_th:
                section_start = video.curpos() + video.fps * start_offset
                print "section start", video.curpos(), section_start
                #cv2.imshow("image", frame)
                if section_start < 0:
                    section_start = 0
                between = True

        if between:
            rect, max_value = end.match(gray)
            print "e", max_value,
            if max_value >= end_th:
                print "section end"
                section_end = video.curpos() + video.fps * end_offset
                if section_end > video.frame_count:
                    section_end = video.frame_count
                section_list.append((section_start, section_end))
                base_name, ext = os.path.splitext(movie_file)
                cutman = SubProcessCut(movie_file)
                cutman.run( float(section_start) / video.fps,
                            float(section_end - section_start) / video.fps,
                            base_name + '_%03d' % number + ext)
                number += 1
                section_start, section_end = 0, 0
                print "section_list", section_list
                between = False
        bar.update(float(video.curpos()) / video.frame_count)
        if not video.seek_frame(video.fps):
            break
        #cv2.waitKey(1)
    if cutman:
        cutman.wait()
        print "waiting cut process..."
    #cv2.destroyAllWindows()

if __name__ == '__main__':
    script_path = os.path.dirname(sys.argv[0])
    preset_path = os.path.join(script_path, sys.argv[1])
    movie_file = os.path.abspath(sys.argv[2])
    main(preset_path, movie_file)


# Copyright (c) 2018 Alaa BEN JABALLAH

from __future__ import print_function
import sys
import os
import time
import cv2
import numpy as np
import theano
import theano.tensor as T
import lasagne

size = 96
input_var = T.tensor4('inputs')

def load_img(img_path):
    img = cv2.imread(img_path)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    faceCascade = cv2.CascadeClassifier("models/haarcascade_frontalface_default.xml")
    faces = faceCascade.detectMultiScale(
        gray,
        scaleFactor=1.1,
        minNeighbors=5,
        minSize=(30, 30),
        flags= cv2.CASCADE_SCALE_IMAGE #cv2.cv.CV_HAAR_SCALE_IMAGE
    )
    for (x, y, w, h) in faces:
        small = cv2.resize(gray[y:y+h, x:x+w], (size,size))
        arr = np.array(small / np.float(256))
        arr = arr.reshape(1,1,size,size)
        return arr
    return np.array([])

# modul przygotowywania danych
def train_net(datadir, imagedir, labeldir, network, epochs=1000, save_each=False):
    import gzip
    def load(img_path):
        img = cv2.imread(img_path)
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        faceCascade = cv2.CascadeClassifier("models/haarcascade_frontalface_default.xml")
        faces = faceCascade.detectMultiScale(
            gray,
            scaleFactor=1.1,
            minNeighbors=5,
            minSize=(30, 30),
            flags=cv2.cv.CV_HAAR_SCALE_IMAGE
        )
        for (x, y, w, h) in faces:
            small = cv2.resize(gray[y:y+h, x:x+w], (size,size))
            return small / np.float(256)

    def load_images(path):
        paths = os.listdir(path)
        files = []
        for p in paths :
            if os.path.isfile(path + "\\" + p) and (p.endswith(".jpg") or p.endswith(".png")):
                files.append(p)
        data = []
        for f in files :
            face = load(path + "\\" + f)
            if face is not None :
                data.append(face)
        return data

    def iterate_minibatches(inputs, targets, batchsize, shuffle=False):
        assert len(inputs) == len(targets)
        if len(inputs) < batchsize :
            batchsize = len(inputs)
        if shuffle:
            indices = np.arange(len(inputs))
            np.random.shuffle(indices)
        for start_idx in range(0, len(inputs) - batchsize + 1, batchsize):
            if shuffle:
                excerpt = indices[start_idx:start_idx + batchsize]
            else:
                excerpt = slice(start_idx, start_idx + batchsize)
            out1 = []
            out2 = []
            for i in excerpt :
                out1.append(inputs[i])
                out2.append(targets[i])
            out1_arr = np.array(out1)
            out1_arr = out1_arr.reshape(-1,1,size,size)
            yield out1_arr, np.array(out2)

    # ladowanie danych
    print("Loading data...")
    data = []
    labels = []
    imagepaths = os.listdir(datadir + imagedir)
    labelpaths = os.listdir(datadir + labeldir)
    counter = 0
    for path in labelpaths :
        count = 100 * counter / len(labelpaths)
        out = '#' * (count / 5) + ' ' * (20 - (count / 5)) + ' ' + str(count) + '%'
        sys.stdout.write('\r'+out)
        if os.path.isdir(datadir + labeldir + path):
            in_dir = os.listdir(datadir + labeldir + path)
            for in_path in in_dir :
                if os.path.isdir(datadir + labeldir + path + '\\' + in_path):
                    txt_dir = os.listdir(datadir + labeldir + path + '\\' + in_path)
                    for textfile in txt_dir:
                        lab = 0
                        if os.path.isfile(datadir + labeldir + path + '\\' + in_path + '\\' + textfile) and textfile.endswith(".txt"):
                            with open(datadir + labeldir + path + '\\' + in_path + '\\' + textfile) as f:
                                for line in f :
                                    lab = int(float(line))
                            imgs = load_images(datadir + imagedir + path + '\\' + in_path + '\\')
                            for i in range (0, len(imgs)):
                                labels.append(lab - 1)
                            data.extend(imgs)
                        break
        counter += 1

    data_arr = np.array(data)
    print (len(data_arr))
    data_arr = data_arr.reshape(-1,1,size,size)
    labels_arr = np.array(labels)
    print ('\ndata loaded')

    target_var = T.ivector('targets')

    # funkcja bledu
    prediction = lasagne.layers.get_output(network)
    loss = lasagne.objectives.categorical_crossentropy(prediction, target_var)
    loss = loss.mean()

    # funkcja dobierajaca kolejne wagi Stochastic Gradient Descent (SGD) with Nesterov momentum
    params = lasagne.layers.get_all_params(network, trainable=True)
    updates = lasagne.updates.nesterov_momentum(
            loss, params, learning_rate=0.01, momentum=0.9)

    # funkcja bledu
    train_fn = theano.function([input_var, target_var], loss, updates=updates)

    # petla trenujaca
    print("Starting training...")
    for epoch in range(epochs):
        train_err = 0
        train_batches = 0
        start_time = time.time()
        for batch in iterate_minibatches(data, labels, 50, shuffle=True):
            inputs, targets = batch
            train_err += train_fn(inputs, targets)
            train_batches += 1

        print("Epoch {} of {} took {:.3f}s".format(
            epoch + 1, epochs, time.time() - start_time))
        print("  training loss:\t\t{:.6f}".format(train_err / train_batches))

        if save_each:
            name = "custom_model%s.npz" % epoch
            np.savez(name, *lasagne.layers.get_all_param_values(network))

    np.savez('custom_model.npz', *lasagne.layers.get_all_param_values(network))

# modul tworzenia architektury sieci
def build_cnn(model = None):
    network = lasagne.layers.InputLayer(shape=(None, 1, size, size),
                                        input_var=input_var)

    network = lasagne.layers.Conv2DLayer(
            network, num_filters=32, filter_size=(5, 5),
            nonlinearity=lasagne.nonlinearities.rectify,
            W=lasagne.init.GlorotUniform())

    network = lasagne.layers.MaxPool2DLayer(network, pool_size=(2, 2))

    network = lasagne.layers.Conv2DLayer(
            network, num_filters=32, filter_size=(5, 5),
            nonlinearity=lasagne.nonlinearities.rectify)
    network = lasagne.layers.MaxPool2DLayer(network, pool_size=(2, 2))

    network = lasagne.layers.DenseLayer(
            lasagne.layers.dropout(network, p=.5),
            num_units=256,
            nonlinearity=lasagne.nonlinearities.rectify)

    network = lasagne.layers.DenseLayer(
            lasagne.layers.dropout(network, p=.5),
            num_units=7,
            nonlinearity=lasagne.nonlinearities.softmax)

    # ladowanie wag sieci jezeli podano
    if model != None:
        with np.load(model) as f:
            param_values = [f['arr_%d' % i] for i in range(len(f.files))]
        lasagne.layers.set_all_param_values(network, param_values)

    return network

# ewaluacja sieci
# 0=anger, 1=contempt, 2=disgust, 3=fear, 4=happy, 5=sadness, 6=surprise
def evaluate(network, faces_matrix):
    if faces_matrix.ndim != 4 :
        return np.array([[.0,.0,.0,.0,.0,.0,.0]])

    out = lasagne.layers.get_output(network, faces_matrix)
    return out.eval()


"""
HNSW module
"""

import numpy as np

# pylint: disable=E0611
from hnswlib import Index

from .base import ANN

class HNSW(ANN):
    """
    Builds an ANN model using the hnswlib library.
    """

    def load(self, path):
        # Load index
        self.model = Index(dim=self.config["dimensions"], space=self.config["metric"])
        self.model.load_index(path)

    def index(self, embeddings):
        # Inner product is equal to cosine similarity on normalized vectors
        self.config["metric"] = "ip"

        # Lookup index settings
        efconstruction = self.setting("efconstruction", 200)
        m = self.setting("m", 16)
        seed = self.setting("randomseed", 100)

        # Create index
        self.model = Index(dim=self.config["dimensions"], space=self.config["metric"])
        self.model.init_index(max_elements=embeddings.shape[0], ef_construction=efconstruction, M=m, random_seed=seed)

        # Add items
        self.model.add_items(embeddings, np.arange(embeddings.shape[0]))

        # Update id offset and set delete counter
        self.config["offset"] = embeddings.shape[0]
        self.config["deletes"] = 0

    def append(self, embeddings):
        new = embeddings.shape[0]

        # Resize index
        self.model.resize_index(self.config["offset"] + new)

        # Append new ids
        self.model.add_items(embeddings, np.arange(self.config["offset"], self.config["offset"] + new))

        # Update id offset
        self.config["offset"] += new

    def delete(self, ids):
        # Mark elements as deleted to omit from search results
        for uid in ids:
            try:
                self.model.mark_deleted(uid)
                self.config["deletes"] += 1
            except RuntimeError:
                # Ignore label not found error
                continue

    def search(self, queries, limit):
        # Set ef query param
        ef = self.setting("efsearch")
        if ef:
            self.model.set_ef(ef)

        # Run the query
        ids, distances = self.model.knn_query(queries, k=limit)

        # Map results to [(id, score)]
        results = []
        for x, distance in enumerate(distances):
            # Convert distances to similarity scores
            scores = [1 - d for d in distance]

            results.append(list(zip(ids[x], scores)))

        return results

    def count(self):
        return self.model.get_current_count() - self.config["deletes"]

    def save(self, path):
        # Write index
        self.model.save_index(path)


import argparse
import os
import sys
import math
import cv2
import numpy as np
import multiprocessing
from sklearn.metrics import confusion_matrix

sys.path.append('.')
from pyActionRecog import parse_directory
from pyActionRecog import parse_split_file

from pyActionRecog.utils.video_funcs import default_aggregation_func

parser = argparse.ArgumentParser()
parser.add_argument('dataset', type=str, choices=['ucf101', 'hmdb51'])
parser.add_argument('split', type=int, choices=[1, 2, 3],
                    help='on which split to test the network')
parser.add_argument('modality', type=str, choices=['rgb', 'flow'])
parser.add_argument('frame_path', type=str, help="root directory holding the frames")
parser.add_argument('net_proto', type=str)
parser.add_argument('net_weights', type=str)
parser.add_argument('--rgb_prefix', type=str, help="prefix of RGB frames", default='img_')
parser.add_argument('--flow_x_prefix', type=str, help="prefix of x direction flow images", default='flow_x_')
parser.add_argument('--flow_y_prefix', type=str, help="prefix of y direction flow images", default='flow_y_')
parser.add_argument('--num_frame_per_video', type=int, default=25,
                    help="prefix of y direction flow images")
parser.add_argument('--save_scores', type=str, default=None, help='the filename to save the scores in')
parser.add_argument('--num_worker', type=int, default=1)
parser.add_argument("--caffe_path", type=str, default='./lib/caffe-action/', help='path to the caffe toolbox')
parser.add_argument("--gpus", type=int, nargs='+', default=None, help='specify list of gpu to use')
args = parser.parse_args()

print args

sys.path.append(os.path.join(args.caffe_path, 'python'))
from pyActionRecog.action_caffe import CaffeNet

# build neccessary information
print args.dataset
split_tp = parse_split_file(args.dataset)
f_info = parse_directory(args.frame_path,
                         args.rgb_prefix, args.flow_x_prefix, args.flow_y_prefix)

gpu_list = args.gpus

eval_video_list = split_tp[args.split - 1][1]

score_name = 'fc-action'

def build_net():
    global net
    my_id = multiprocessing.current_process()._identity[0] \
        if args.num_worker > 1 else 1
    if gpu_list is None:
        net = CaffeNet(args.net_proto, args.net_weights, my_id-1)
    else:
        net = CaffeNet(args.net_proto, args.net_weights, gpu_list[my_id - 1])

def eval_video(video):
    global net
    label = video[1]
    vid = video[0]

    video_frame_path = f_info[0][vid]
    if args.modality == 'rgb':
        cnt_indexer = 1
    elif args.modality == 'flow':
        cnt_indexer = 2
    else:
        raise ValueError(args.modality)
    frame_cnt = f_info[cnt_indexer][vid]

    stack_depth = 0
    if args.modality == 'rgb':
        stack_depth = 1
    elif args.modality == 'flow':
        stack_depth = 5

    step = (frame_cnt - stack_depth) / (args.num_frame_per_video-1)
    if step > 0:
        frame_ticks = range(1, min((2 + step * (args.num_frame_per_video-1)), frame_cnt+1), step)
    else:
        frame_ticks = [1] * args.num_frame_per_video

    assert(len(frame_ticks) == args.num_frame_per_video)

    frame_scores = []
    for tick in frame_ticks:
        if args.modality == 'rgb':
            name = '{}{:05d}.jpg'.format(args.rgb_prefix, tick)
            frame = cv2.imread(os.path.join(video_frame_path, name), cv2.IMREAD_COLOR)
            scores = net.predict_single_frame([frame,], score_name, frame_size=(340, 256))
            frame_scores.append(scores)
        if args.modality == 'flow':
            frame_idx = [min(frame_cnt, tick+offset) for offset in xrange(stack_depth)]
            flow_stack = []
            for idx in frame_idx:
                x_name = '{}{:05d}.jpg'.format(args.flow_x_prefix, idx)
                y_name = '{}{:05d}.jpg'.format(args.flow_y_prefix, idx)
                flow_stack.append(cv2.imread(os.path.join(video_frame_path, x_name), cv2.IMREAD_GRAYSCALE))
                flow_stack.append(cv2.imread(os.path.join(video_frame_path, y_name), cv2.IMREAD_GRAYSCALE))
            scores = net.predict_single_flow_stack(flow_stack, score_name, frame_size=(340, 256))
            frame_scores.append(scores)

    print 'video {} done'.format(vid)
    sys.stdin.flush()
    return np.array(frame_scores), label

if args.num_worker > 1:
    pool = multiprocessing.Pool(args.num_worker, initializer=build_net)
    video_scores = pool.map(eval_video, eval_video_list)
else:
    build_net()
    video_scores = map(eval_video, eval_video_list)

video_pred = [np.argmax(default_aggregation_func(x[0])) for x in video_scores]
video_labels = [x[1] for x in video_scores]

cf = confusion_matrix(video_labels, video_pred).astype(float)

cls_cnt = cf.sum(axis=1)
cls_hit = np.diag(cf)

cls_acc = cls_hit/cls_cnt

print cls_acc

print 'Accuracy {:.02f}%'.format(np.mean(cls_acc)*100)

if args.save_scores is not None:
    np.savez(args.save_scores, scores=video_scores, labels=video_labels)