#!/bin/env python
import sys, getopt
import math
from PIL import Image
import numpy as np

phenotypes = {
	"subnordid" : "pheno_regions/subnordid.gif",
	"atlanto_mediterranean" : "pheno_regions/atlanto_mediterranean.gif",
	"basic_veddoid" : "pheno_regions/veddoid.gif",
	"basic_turanid" : "pheno_regions/turanid.gif",
	"basic_tungid" : "pheno_regions/tungid.gif",
	"basic_sudanid" : "pheno_regions/sudanid.gif",
	"basic_south_mongolid" : "pheno_regions/southmongolid.gif",
	"basic_sinid" : "pheno_regions/sinid.gif",
	"basic_sibirid" : "pheno_regions/sibirid.gif",
	"basic_orientalid" : "pheno_regions/orientalid.gif",
	"basic_nordid" : "pheno_regions/nordid.gif",
	"basic_nilotid" : "pheno_regions/nilotid.gif",
	"basic_mediterranean" : "pheno_regions/mediterranid.gif",
	"basic_lappid" : "pheno_regions/lappid.gif",
	"basic_indo_melanid" : "pheno_regions/indo-melanid.gif",
	"basic_indid" : "pheno_regions/indid.gif",
	"basic_ethiopid" : "pheno_regions/ethiopid.gif",
	"basic_east_europid" : "pheno_regions/easteuropid.gif",
	"basic_dinarid" : "pheno_regions/dinarid.gif",
	"basic_congolid" : "pheno_regions/congolid.gif",
	"basic_armenoid" : "pheno_regions/armenoid.gif",
	"basic_alpine" : "pheno_regions/alpine.gif",
	"basic_ainuid" : "pheno_regions/ainuid.gif",
	"yakonin" : "pheno_regions/yakonin.gif",
	"west_alpinid" : "pheno_regions/west_alpinid.gif",
	"troender" : "pheno_regions/tronder.gif",
	"tonkinesid" : "pheno_regions/tonkinesid.gif",
	"tibetid" : "pheno_regions/tibetid.gif",
	"tavastid" : "pheno_regions/tavastid.gif",
	"targid" : "pheno_regions/targid.gif",
	"strandid" : "pheno_regions/strandid.gif",
	"shanid" : "pheno_regions/shanid.gif",
	"savolaxid" : "pheno_regions/savolaxid.gif",
	"satsuma" : "pheno_regions/satsuma.gif",
	"pre_slavic" : "pheno_regions/pre-slavic.gif",
	"pontid" : "pheno_regions/pontid.gif",
	"paleo_sardinian" : "pheno_regions/paleo_sardinian.gif",
	"paleo_atlantid" : "pheno_regions/paleo_atlantid.gif",
	"north_pontid" : "pheno_regions/north_pontid.gif",
	"north_indid" : "pheno_regions/north_indid.gif",
	"north_atlantid" : "pheno_regions/north_atlantid.gif",
	"norid" : "pheno_regions/norid.gif",
	"neo_danubian" : "pheno_regions/neo_danubian.gif",
	"mtebid" : "pheno_regions/mtebid.gif",
	"mountain_indid" : "pheno_regions/mountain_indid.gif",
	"moorish" : "pheno_regions/moorish.gif",
	"manchukorean" : "pheno_regions/manchukorean.gif",
	"litorid" : "pheno_regions/litorid.gif",
	"libyid" : "pheno_regions/libyid.gif",
	"ladogan" : "pheno_regions/ladogan.gif",
	"kham" : "pheno_regions/kham.gif",
	"ishikawa" : "pheno_regions/ishikawa.gif",
	"iranid" : "pheno_regions/iranid.gif",
	"indo_nordid" : "pheno_regions/indo_nordid.gif",
	"indo_iranid" : "pheno_regions/indo_iranid.gif",
	"indo_brachid" : "pheno_regions/indo_brachid.gif",
	"huanghoid" : "pheno_regions/huanghoid.gif",
	"halstatt" : "pheno_regions/halstatt.gif",
	"gracile_indid" : "pheno_regions/gracile_indid.gif",
	"gorid" : "pheno_regions/gorid.gif",
	"gobid" : "pheno_regions/gobid.gif",
	"fenno_nordid" : "pheno_regions/fenno_nordid.gif",
	"faelid" : "pheno_regions/faelid.gif",
	"eurafricanid" : "pheno_regions/eurafricanid.gif",
	"egyptid" : "pheno_regions/egyptid.gif",
	"east_brachid" : "pheno_regions/east_brachid.gif",
	"choshiu" : "pheno_regions/choshiu.gif",
	"chikuzen" : "pheno_regions/chikuzen.gif",
	"chukiangid" : "pheno_regions/chukiangid.gif",
	"changkiangid" : "pheno_regions/changkiangid.gif",
	"central_pamirid" : "pheno_regions/central_pamirid.gif",
	"central_brachid" : "pheno_regions/central_brachid.gif",
	"carpathid" : "pheno_regions/carpathid.gif",
	"borreby" : "pheno_regions/borreby.gif",
	"berberid" : "pheno_regions/berberid.gif",
	"baskid" : "pheno_regions/baskid.gif",
	"asian_alpine" : "pheno_regions/asian_alpine.gif",
	"assyroid" : "pheno_regions/assyroid.gif",
	"plains_pamirid" : "pheno_regions/plains_pamirid.gif",
	"aralid" : "pheno_regions/aralid.gif",
	"arabid" : "pheno_regions/arabid.gif",
	"annamid" : "pheno_regions/annamid.gif",
	"anglo_saxon" : "pheno_regions/anglo_saxon.gif",
	"andronovo_turanid" : "pheno_regions/andronovo-turanid.gif",
	"anatolid" : "pheno_regions/anatolid.gif",
	"alfoeld" : "pheno_regions/alfoeld.gif",
	"aisto_nordid" : "pheno_regions/aisto_nordid.gif",
	"african_alpine" : "pheno_regions/african_alpine.gif"}

def fittingPixels(maskImg, typeImg, rgb):
	yellow = np.asarray([255,255,0 ],dtype=np.uint8)
	darkYellow = np.asarray([127,127,0 ],dtype=np.uint8)
	if maskImg.shape[0] != typeImg.shape[0]:
		return 0
	if maskImg.shape[1] != typeImg.shape[1]:
		return 0
	toReturn = 0.0
	for x in range(maskImg.shape[0]):
		for y in range(maskImg.shape[1]):
			a = maskImg[x][y]
			b = typeImg[x][y]
			if np.array_equal(a,rgb) :
				if np.array_equal(b,yellow):
					toReturn = toReturn + 1.0
				elif np.array_equal(b,darkYellow):
					toReturn = toReturn + 0.1
	return toReturn

def compareWIth(maskImg,rgb):
	items= {}
	for x in phenotypes:
		typeImage = Image.open(phenotypes[x])
		if typeImage.mode != 'RGB' :
			typeImage = typeImage.convert('RGB')
		typeImage = np.array(typeImage)
		temp = fittingPixels(maskImg,typeImage,rgb)
		if temp > 0.0:
			items[x] = temp
	summa = 0.0
	for x in items:
		summa = summa + items[x]
	for x in items:
		items[x] = int(math.ceil((items[x]/summa)*100.0))
	sorted_items = dict(sorted(items.items(), key=lambda x: x[1], reverse=True))
	return sorted_items

def main(argv):
	maskfile = ''
	color = 'ffffff'
	try:
		opts, args = getopt.getopt(argv,"hm:c:",["maskfile=","color="])
	except getopt.GetoptError:
		print('test.py -m <inputfile> -c <color>')
		sys.exit(2)
	for opt, arg in opts:
		if opt == '-h':
			print('test.py -m <inputfile>')
			sys.exit()
		elif opt in ("-m", "--maskfile"):
			maskfile = arg
		elif opt in ("-c", "--color"):
			color = arg

	rgb = tuple(int(color[i:i+2], 16) for i in (0, 2, 4))
	rgb = np.asarray(rgb,dtype=np.uint8)
	maskImage = Image.open(maskfile)
	if maskImage.mode != 'RGB' :
		maskImage = maskImage.convert('RGB')
	maskImage = np.array(maskImage)
	items = compareWIth(maskImage,rgb)
	for x in items:
		print(str(items[x]) + " = " + x)

main(sys.argv[1:])


# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
#
# 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 argparse
import copy
import os

import cv2
import numpy as np
import paddle
from skimage.io import imread
from skimage.transform import rescale
from skimage.transform import resize

import paddlehub as hub
from .api import PRN
from .predictor import PosPrediction
from .util import base64_to_cv2
from .utils.render import render_texture
from paddlehub.module.module import moduleinfo
from paddlehub.module.module import runnable
from paddlehub.module.module import serving

@moduleinfo(name="prnet", type="CV/", author="paddlepaddle", author_email="", summary="", version="1.0.0")
class PRNet:
    def __init__(self):
        self.pretrained_model = os.path.join(self.directory, "pd_model/model.pdparams")
        self.network = PRN(is_dlib=True, prefix=self.directory)

    def face_swap(self,
                  images: list = None,
                  paths: list = None,
                  mode: int = 0,
                  output_dir: str = './swapping_result/',
                  use_gpu: bool = False,
                  visualization: bool = True):
        '''
        Denoise a raw image in the low-light scene.

        images (list[dict]): data of images, each element is a dict:
          - source (numpy.ndarray): input image，shape is \[H, W, C\]，BGR format；

          - ref (numpy.ndarray) : style image，shape is \[H, W, C\]，BGR format；

        paths (list[dict]): paths to images, eacg element is a dict:
          - source (str): path to input image；

          - ref (str) : path to reference image；

        mode (int): option, 0 for change part of texture, 1 for change whole face
        output_dir (str): the dir to save the results
        use_gpu (bool): if True, use gpu to perform the computation, otherwise cpu.
        visualization (bool): if True, save results in output_dir.
        '''
        results = []
        paddle.disable_static()
        place = 'gpu:0' if use_gpu else 'cpu'
        place = paddle.set_device(place)
        if images == None and paths == None:
            print('No image provided. Please input an image or a image path.')
            return

        if images != None:
            for image_dict in images:
                source_img = image_dict['source'][:, :, ::-1]
                ref_img = image_dict['ref'][:, :, ::-1]
                results.append(self.texture_editing(source_img, ref_img, mode))

        if paths != None:
            for path_dict in paths:
                source_img = cv2.imread(path_dict['source'])[:, :, ::-1]
                ref_img = cv2.imread(path_dict['ref'])[:, :, ::-1]
                results.append(self.texture_editing(source_img, ref_img, mode))

        if visualization == True:
            if not os.path.exists(output_dir):
                os.makedirs(output_dir, exist_ok=True)
            for i, out in enumerate(results):
                cv2.imwrite(os.path.join(output_dir, 'output_{}.png'.format(i)), out[:, :, ::-1])

        return results

    def texture_editing(self, source_img, ref_img, mode):
        # read image
        image = source_img
        [h, w, _] = image.shape
        prn = self.network
        #-- 1. 3d reconstruction -> get texture.
        pos = prn.process(image)
        vertices = prn.get_vertices(pos)
        image = image / 255.
        texture = cv2.remap(
            image,
            pos[:, :, :2].astype(np.float32),
            None,
            interpolation=cv2.INTER_NEAREST,
            borderMode=cv2.BORDER_CONSTANT,
            borderValue=(0))

        #-- 2. Texture Editing
        Mode = mode
        # change part of texture(for data augumentation/selfie editing. Here modify eyes for example)
        if Mode == 0:
            # load eye mask
            uv_face_eye = imread(os.path.join(self.directory, 'Data/uv-data/uv_face_eyes.png'), as_gray=True) / 255.
            uv_face = imread(os.path.join(self.directory, 'Data/uv-data/uv_face.png'), as_gray=True) / 255.
            eye_mask = (abs(uv_face_eye - uv_face) > 0).astype(np.float32)

            # texture from another image or a processed texture
            ref_image = ref_img
            ref_pos = prn.process(ref_image)
            ref_image = ref_image / 255.
            ref_texture = cv2.remap(
                ref_image,
                ref_pos[:, :, :2].astype(np.float32),
                None,
                interpolation=cv2.INTER_NEAREST,
                borderMode=cv2.BORDER_CONSTANT,
                borderValue=(0))

            # modify texture
            new_texture = texture * (1 - eye_mask[:, :, np.newaxis]) + ref_texture * eye_mask[:, :, np.newaxis]

        # change whole face(face swap)
        elif Mode == 1:
            # texture from another image or a processed texture
            ref_image = ref_img
            ref_pos = prn.process(ref_image)
            ref_image = ref_image / 255.
            ref_texture = cv2.remap(
                ref_image,
                ref_pos[:, :, :2].astype(np.float32),
                None,
                interpolation=cv2.INTER_NEAREST,
                borderMode=cv2.BORDER_CONSTANT,
                borderValue=(0))
            ref_vertices = prn.get_vertices(ref_pos)
            new_texture = ref_texture  #(texture + ref_texture)/2.

        else:
            print('Wrong Mode! Mode should be 0 or 1.')
            exit()

        #-- 3. remap to input image.(render)
        vis_colors = np.ones((vertices.shape[0], 1))
        face_mask = render_texture(vertices.T, vis_colors.T, prn.triangles.T, h, w, c=1)
        face_mask = np.squeeze(face_mask > 0).astype(np.float32)

        new_colors = prn.get_colors_from_texture(new_texture)
        new_image = render_texture(vertices.T, new_colors.T, prn.triangles.T, h, w, c=3)
        new_image = image * (1 - face_mask[:, :, np.newaxis]) + new_image * face_mask[:, :, np.newaxis]

        # Possion Editing for blending image
        vis_ind = np.argwhere(face_mask > 0)
        vis_min = np.min(vis_ind, 0)
        vis_max = np.max(vis_ind, 0)
        center = (int((vis_min[1] + vis_max[1]) / 2 + 0.5), int((vis_min[0] + vis_max[0]) / 2 + 0.5))
        output = cv2.seamlessClone((new_image * 255).astype(np.uint8), (image * 255).astype(np.uint8),
                                   (face_mask * 255).astype(np.uint8), center, cv2.NORMAL_CLONE)

        return output

    @runnable
    def run_cmd(self, argvs: list):
        """
        Run as a command.
        """
        self.parser = argparse.ArgumentParser(
            description="Run the {} module.".format(self.name),
            prog='hub run {}'.format(self.name),
            usage='%(prog)s',
            add_help=True)

        self.arg_input_group = self.parser.add_argument_group(title="Input options", description="Input data. Required")
        self.arg_config_group = self.parser.add_argument_group(
            title="Config options", description="Run configuration for controlling module behavior, not required.")
        self.add_module_config_arg()
        self.add_module_input_arg()
        self.args = self.parser.parse_args(argvs)

        self.face_swap(
            paths=[{
                'source': self.args.source,
                'ref': self.args.ref
            }],
            mode=self.args.mode,
            output_dir=self.args.output_dir,
            use_gpu=self.args.use_gpu,
            visualization=self.args.visualization)

    @serving
    def serving_method(self, images, **kwargs):
        """
        Run as a service.
        """
        images_decode = copy.deepcopy(images)
        for image in images_decode:
            image['source'] = base64_to_cv2(image['source'])
            image['ref'] = base64_to_cv2(image['ref'])
        results = self.face_swap(images_decode, **kwargs)
        tolist = [result.tolist() for result in results]
        return tolist

    def add_module_config_arg(self):
        """
        Add the command config options.
        """
        self.arg_config_group.add_argument(
            '--mode', type=int, default=0, help='process option, 0 for part texture, 1 for whole face.', choices=[0, 1])
        self.arg_config_group.add_argument('--use_gpu', action='store_true', help="use GPU or not")

        self.arg_config_group.add_argument(
            '--output_dir', type=str, default='swapping_result', help='output directory for saving result.')
        self.arg_config_group.add_argument('--visualization', type=bool, default=False, help='save results or not.')

    def add_module_input_arg(self):
        """
        Add the command input options.
        """
        self.arg_input_group.add_argument('--source', type=str, help="path to source image.")
        self.arg_input_group.add_argument('--ref', type=str, help="path to reference image.")


"""Train a caption model.

This script trains the captioning model.

Code is adapted from:
https://machinelearningmastery.com/develop-a-deep-learning-caption-generation-model-in-python/
"""

from keras.callbacks import ModelCheckpoint
from keras.preprocessing.text import Tokenizer
from keras.preprocessing.sequence import pad_sequences
from keras.utils import to_categorical
from model import caption_model
from numpy import array
import pickle
from sklearn.model_selection import train_test_split

from nltk.corpus import stopwords

TITLES = "../resources/titles.txt"
FEATURES = "../resources/features.pkl"
MODEL = "../resources/model.h5"
TOKENIZER = "../resources/tokenizer.pkl"

stopwords = stopwords.words("english")

def clean_titles(title):
    title = title.split()
    title = [word.lower() for word in title]
    title = [word for word in title if len(word) > 1]
    title = [word for word in title if word.isalpha()]
    title = [word for word in title if word not in stopwords]
    title = " ".join(title)
    return title

def load_titles(path):
    titles = dict()
    with open(path, "r") as f:
        for line in f:
            tokens = line.strip("\n").split("\t")
            video_id, title = tokens[0], tokens[1]
            title = "startseq " + title + " endseq"
            titles[video_id] = title

    return titles

def load_image_features(path, dataset):
    features = pickle.load(open(path, "rb"))
    return {k: features[k] for k in dataset}

def to_lines(titles):
    all_titles = []
    for key in titles.keys():
        for t in titles[key].split():
            all_titles.append(t)
    return all_titles

def create_tokenizer(titles):
    all_titles = []
    for k in titles.keys():
        all_titles.append(titles[k])
    tokenizer = Tokenizer()
    tokenizer.fit_on_texts(all_titles)
    return tokenizer

def max_length(titles):
    lengths = []
    for key in titles.keys():
        for t in titles[key].split():
            lengths.append(len(t))
    return max(lengths)

def make_title_set(titles, dataset):
    title_set = dict()
    for video_id in dataset:
        title_set[video_id] = titles[video_id]
    return title_set

def create_sequences(tokenizer, max_length, titles, images):
    X_img = []
    X_seq = []
    Y = []

    for key, title in titles.items():
        # title = clean_titles(title)
        seq = tokenizer.texts_to_sequences([title])[0]

        for i in range(1, len(seq)):
            in_seq = seq[:i]
            out_seq = seq[i]

            in_seq = pad_sequences([in_seq], maxlen=max_length)[0]
            out_seq = to_categorical([out_seq], num_classes=vocab_size)[0]

            X_img.append(images[key][0])
            X_seq.append(in_seq)
            Y.append(out_seq)

    return array(X_img), array(X_seq), array(Y)

if __name__ == "__main__":
    # Load titles
    titles = load_titles(TITLES)
    dataset = list(titles.keys())

    train_dataset, test_dataset = train_test_split(
        dataset, train_size=0.7, test_size=0.3
    )

    # Make image feature sets
    train_features = load_image_features(FEATURES, train_dataset)  # dict()
    test_features = load_image_features(FEATURES, test_dataset)  # dict()

    # Make title sets
    train_titles = make_title_set(titles, train_dataset)  # dict()
    test_titles = make_title_set(titles, test_dataset)  # dict()

    # Prepare tokenizer
    tokenizer = create_tokenizer(train_titles)
    vocab_size = len(tokenizer.word_index) + 1
    print("Vocab size: ", vocab_size)
    max_length = max_length(train_titles)
    print("Title length:", max_length)

    pickle.dump(tokenizer, open(TOKENIZER, "wb"))

    X1_train, X2_train, Ytrain = create_sequences(
        tokenizer, max_length, train_titles, train_features
    )
    X1_test, X2_test, Ytest = create_sequences(
        tokenizer, max_length, test_titles, test_features
    )

    model = caption_model(vocab_size, max_length)

    filepath = MODEL
    checkpoint = ModelCheckpoint(
        filepath, monitor="val_loss", verbose=1, save_best_only=True, mode="min"
    )

    model.fit(
        [X1_train, X2_train],
        Ytrain,
        epochs=50,
        verbose=2,
        callbacks=[checkpoint],
        validation_data=([X1_test, X2_test], Ytest),
    )


import numpy as np
from app import main, init_activity_detection, init_classificator, init_feature_extraction, init_pre_processing
from models import Waveform
from utils import load_groundtruth, read_csv, plot_metrics_classification_boxplot
from utils import plot_audio, plot_odf, plot_confusion_matrix, plot_evaluation_report
from evaluate import evaluate_system, evaluate_activity_detection, evaluate_classificator
import pickle
import pandas as pd

"""
The run script is used to run an off-line simulation of the system in order to evaluate some interfaces
for performance testing.
"""
# Import model
from utils.load_csv import get_prediction_time_instants, load_annotation

filename = './app/finalized_model_mfccs.sav'
knn_model = pickle.load(open(filename, 'rb'))

model_normalization = './app/model_normalization_mfccs.csv'

normalization_values = pd.read_csv(model_normalization)

path = '../../TrainDataset/LOD_1624992635178/Snare_LOD'

audio = Waveform(path=path + ".wav")
groundtruth = load_annotation(path + ".csv")

# Init system
init_pre_processing()
init_activity_detection(func_type=1)
init_feature_extraction(func_type="mfcc", by_pass=False, n_mfcc_arg=20, norm_file=normalization_values)
init_classificator(knn_model=knn_model, by_pass=False)
buffer_len = 512

# Call system
result = main(audio, buffer_len)

prediction = get_prediction_time_instants(result['ONSET_LOCATIONS'], result['PREDICTION'], audio.sample_rate)
# Plot results
plot_audio(audio.waveform, result['SIGNAL_PROCESSED'], audio.sample_rate)
plot_odf(audio.filename, audio.waveform, result['SIGNAL_PROCESSED'], audio.sample_rate, result['ONSET_LOCATIONS'],
         result['HFC'], result['THRESHOLD'])

groundtruth_activity = np.zeros(len(result['ONSET_LOCATIONS']))
# Transform annotation in the desired format (1 activity, 0 non-activity)
for i in range(0, len(groundtruth), 2):
    sample_instant_1 = int(float(groundtruth[i][0]) * audio.sample_rate)
    sample_instant_2 = int(float(groundtruth[i + 1][0]) * audio.sample_rate)
    groundtruth_activity[sample_instant_1:sample_instant_2] = 1

# Evaluation
precision_ad, recall_ad, fscore_ad, accuracy_ad = evaluate_activity_detection(groundtruth_activity,
                                                                              result['ONSET_LOCATIONS'])

report, cm = evaluate_classificator(groundtruth, prediction)
precision, recall, fscore = evaluate_system(groundtruth, prediction)

print('----------------------------------------------------------------------')
print('Prototype report')
print('----------------------------------------------------------------------')
print('Activity detection evaluation:')
print('\n Precision:' + str(precision_ad) +
      '\n Recall:' + str(recall_ad) +
      '\n F1-score:' + str(fscore_ad) +
      '\n Accuracy:' + str(accuracy_ad)
      )
print('----------------------------------------------------------------------')
print('Classification evaluation:')
plot_evaluation_report(report)
print(prediction)
if len(cm) == 3:
    plot_confusion_matrix(cm)
print('----------------------------------------------------------------------')
print('System evaluation:')
print('\n Precision:' + str(precision) +
      '\n Recall:' + str(recall) +
      '\n F1-score:' + str(fscore)
      )
print('----------------------------------------------------------------------')


import pytest
import numpy as np
import os.path as op

from .. import clf as CLF

def load_data(category, name, basepath=None):
    from scipy.io import loadmat
    from datalad.api import get

    if basepath is None:
        basepath = op.join(
        'remodnav', 'tests', 'data', 'anderson_etal', 'annotated_data',
        'data used in the article')

    fname = op.join(
        basepath,
        category, name + ('' if name.endswith('.mat') else '.mat'))
    get(fname)
    m = loadmat(fname)
    # viewing distance
    vdist = m['ETdata']['viewDist'][0][0][0][0]
    screen_width = m['ETdata']['screenDim'][0][0][0][0]
    screen_res = m['ETdata']['screenRes'][0][0][0][0]
    px2deg = CLF.deg_per_pixel(screen_width, vdist, screen_res)
    #testing
    #px2deg = 0.001
    sr = float(m['ETdata']['sampFreq'][0][0][0][0])
    data = np.rec.fromarrays([
        m['ETdata']['pos'][0][0][:, 3],
        m['ETdata']['pos'][0][0][:, 4]],
        names=('x', 'y'))
    data[np.logical_and(data['x'] == 0, data['y'] == 0)] = (np.nan, np.nan)

    labels = m['ETdata']['pos'][0][0][:, 5]

    label_remap = {
        1: 'FIXA',
        2: 'SACC',
        3: 'PSO',
        4: 'PURS',
    }
    events = []
    ev_type = None
    ev_start = None
    for i in range(len(labels)):
        s = labels[i]
        if ev_type is None and s in label_remap.keys():
            ev_type = s
            ev_start = i
        elif ev_type is not None and s != ev_type:
            events.append(dict(
                id=len(events),
                label=label_remap.get(ev_type),
                start_time=0.0 if ev_start is None else
                float(ev_start) / sr,
                start_index=0 if ev_start is None else
                ev_start,

                end_time=float(i) / sr,
                end_index=i,
            ))

            ev_type = s if s in label_remap.keys() else None
            ev_start = i

    if ev_type is not None:
        events.append(dict(
            id=len(events),
            label=label_remap.get(ev_type),
            start_time=0.0 if ev_start is None else
            float(ev_start) / sr,
            start_index=0 if ev_start is None else
            ev_start,

            end_time=float(i) / sr,
            end_index=i,

        ))
    return data, labels, events, px2deg, sr

@pytest.mark.parametrize(
    'name',
    [
        ('dots', 'TH20_trial1_labelled_MN.mat'),
        ('dots', 'TH20_trial1_labelled_RA.mat'),
        ('dots', 'TH38_trial1_labelled_MN.mat'),
        ('dots', 'TH38_trial1_labelled_RA.mat'),
        ('dots', 'TL22_trial17_labelled_MN.mat'),
        ('dots', 'TL22_trial17_labelled_RA.mat'),
        ('dots', 'TL24_trial17_labelled_MN.mat'),
        ('dots', 'TL24_trial17_labelled_RA.mat'),
        ('dots', 'UH21_trial17_labelled_MN.mat'),
        ('dots', 'UH21_trial17_labelled_RA.mat'),
        ('dots', 'UH21_trial1_labelled_MN.mat'),
        ('dots', 'UH21_trial1_labelled_RA.mat'),
        ('dots', 'UH25_trial1_labelled_MN.mat'),
        ('dots', 'UH25_trial1_labelled_RA.mat'),
        ('dots', 'UH33_trial17_labelled_MN.mat'),
        ('dots', 'UH33_trial17_labelled_RA.mat'),
        ('dots', 'UL27_trial17_labelled_MN.mat'),
        ('dots', 'UL27_trial17_labelled_RA.mat'),
        ('dots', 'UL31_trial1_labelled_MN.mat'),
        ('dots', 'UL31_trial1_labelled_RA.mat'),
        ('dots', 'UL39_trial1_labelled_MN.mat'),
        ('dots', 'UL39_trial1_labelled_RA.mat'),
        ('img', 'TH34_img_Europe_labelled_MN.mat'),
        ('img', 'TH34_img_Europe_labelled_RA.mat'),
        ('img', 'TH34_img_vy_labelled_MN.mat'),
        ('img', 'TH34_img_vy_labelled_RA.mat'),
        ('img', 'TL20_img_konijntjes_labelled_MN.mat'),
        ('img', 'TL20_img_konijntjes_labelled_RA.mat'),
        ('img', 'TL28_img_konijntjes_labelled_MN.mat'),
        ('img', 'TL28_img_konijntjes_labelled_RA.mat'),
        ('img', 'UH21_img_Rome_labelled_MN.mat'),
        ('img', 'UH21_img_Rome_labelled_RA.mat'),
        ('img', 'UH27_img_vy_labelled_MN.mat'),
        ('img', 'UH27_img_vy_labelled_RA.mat'),
        ('img', 'UH29_img_Europe_labelled_MN.mat'),
        ('img', 'UH29_img_Europe_labelled_RA.mat'),
        ('img', 'UH33_img_vy_labelled_MN.mat'),
        ('img', 'UH33_img_vy_labelled_RA.mat'),
        ('img', 'UH47_img_Europe_labelled_MN.mat'),
        ('img', 'UH47_img_Europe_labelled_RA.mat'),
        ('img', 'UL23_img_Europe_labelled_MN.mat'),
        ('img', 'UL23_img_Europe_labelled_RA.mat'),
        ('img', 'UL31_img_konijntjes_labelled_MN.mat'),
        ('img', 'UL31_img_konijntjes_labelled_RA.mat'),
        ('img', 'UL39_img_konijntjes_labelled_MN.mat'),
        ('img', 'UL39_img_konijntjes_labelled_RA.mat'),
        ('img', 'UL43_img_Rome_labelled_MN.mat'),
        ('img', 'UL43_img_Rome_labelled_RA.mat'),
        ('img', 'UL47_img_konijntjes_labelled_MN.mat'),
        ('img', 'UL47_img_konijntjes_labelled_RA.mat'),
        ('video', 'TH34_video_BergoDalbana_labelled_MN.mat'),
        ('video', 'TH34_video_BergoDalbana_labelled_RA.mat'),
        ('video', 'TH38_video_dolphin_fov_labelled_MN.mat'),
        ('video', 'TH38_video_dolphin_fov_labelled_RA.mat'),
        ('video', 'TL30_video_triple_jump_labelled_MN.mat'),
        ('video', 'TL30_video_triple_jump_labelled_RA.mat'),
        ('video', 'UH21_video_BergoDalbana_labelled_MN.mat'),
        ('video', 'UH21_video_BergoDalbana_labelled_RA.mat'),
        ('video', 'UH29_video_dolphin_fov_labelled_MN.mat'),
        ('video', 'UH29_video_dolphin_fov_labelled_RA.mat'),
        ('video', 'UH47_video_BergoDalbana_labelled_MN.mat'),
        ('video', 'UH47_video_BergoDalbana_labelled_RA.mat'),
        ('video', 'UL23_video_triple_jump_labelled_MN.mat'),
        ('video', 'UL23_video_triple_jump_labelled_RA.mat'),
        ('video', 'UL27_video_triple_jump_labelled_MN.mat'),
        ('video', 'UL27_video_triple_jump_labelled_RA.mat'),
        ('video', 'UL31_video_triple_jump_labelled_MN.mat'),
        ('video', 'UL31_video_triple_jump_labelled_RA.mat'),
    ])
def test_labeled(name):
    data, target_labels, target_events, px2deg, sr = load_data(name[0], name[1])

    clf = CLF.EyegazeClassifier(
        px2deg=px2deg,
        sampling_rate=sr,
        pursuit_velthresh=5.,
        noise_factor=3.0,
        lowpass_cutoff_freq=10.0,
    )
    p = clf.preproc(
        data,
    )

    events = clf(p)
    for i, ev in enumerate(events[1:]):
        spacing = ev['start_time'] - events[i]['end_time']
        # events are either touching or have a small gap between
        # two saccade-type events (up to min_fixation_duration),
        # or we have missing data
        assert spacing == 0 or \
            (('PS' in events[i]['label'] or 'SAC' in events[i]['label']) and \
             'SAC' in ev['label'] and spacing < 0.04) or \
            np.sum(np.isnan(
                data[int(events[i]['end_time'] * sr):
                     int(ev['start_time'] * sr)]['x']))
        # PSOs are always following a saccade
        if 'PS' in ev['label']:
            assert 'SAC' in events[i]['label']

#    import pylab as pl
#    pl.subplot(211)
#    ut.show_gaze(pp=p, events=events, sampling_rate=sr)
#    pl.subplot(212)
#    ut.show_gaze(pp=p, events=target_events, sampling_rate=sr)
#    pl.show()


import copy
import os

import matplotlib.pyplot as plt
import numpy as np
from sklearn import metrics
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from dgl.dataloading import GraphDataLoader

from reactome_gnn import model, dataset, hyperparameters

def draw_loss_plot(train_loss, valid_loss, save_path):
    """Draw and save plot of train and validation loss over epochs.

    Parameters
    ----------
    train_loss : list
        List of training loss for each epoch
    valid_loss : list
        List of validation loss for each epoch
    save_path : str
        A path where the plot will be saved

    Returns
    -------
    None
    """
    plt.figure()
    plt.plot(train_loss, label='train')
    plt.plot(valid_loss, label='validation')
    plt.title('Loss over epochs')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.legend()
    plt.savefig(f'{save_path}.png')
    plt.show()

def draw_f1_plot(train_f1, valid_f1, save_path):
    """Draw and save plot of F1 score calculated during validation.

    Parameters
    ----------
    f1 scores : list
        List of f1 scores during validation
    save_path : str
        A path where the plot will be saved

    Returns
    -------
    None
    """
    plt.figure()
    plt.plot(train_f1, label='train')
    plt.plot(valid_f1, label='validation')
    plt.title('F1-score over epochs')
    plt.xlabel('Epoch')
    plt.ylabel('F1-score')
    plt.legend()
    plt.savefig(f'{save_path}.png')
    plt.show()

def train(hyperparams=None, data_path='demo/data/example', plot=True):
    """Train a network which will produce the embeddings.

    Parameters
    ----------
    hyperparams : dict, optional
        Dictionary with hyperparameters
    data_path : str, optional
        Relative path to where the data is stored
    plot : bool, optional
        Whether to plot loss over epochs

    Returns
    -------
    None
    """
    if hyperparams is None:
        hyperparams = hyperparameters.get_hyperparameters()

    # Unpack hyperparameters
    num_epochs = hyperparams['num_epochs']
    dim_latent = hyperparams['dim_latent']
    num_layers = hyperparams['num_layers']
    learning_rate = hyperparams['lr']
    batch_size = hyperparams['batch_size']
    device = hyperparams['device']

    model_path = os.path.join(data_path, 'models')
    model_path = os.path.join(model_path, f'model_dim{dim_latent}_lay{num_layers}.pth')

    # Create datasets
    ds = dataset.PathwayDataset(data_path)
    ds_train = [ds[0], ds[1], ds[2], ds[5], ds[6], ds[7], ds[8], ds[9]]
    ds_valid = [ds[3], ds[4], ds[10], ds[11]]
    dl_train = GraphDataLoader(ds_train, batch_size=batch_size, shuffle=True)
    dl_valid = GraphDataLoader(ds_valid, batch_size=batch_size, shuffle=False)

    # Initialize networks and optimizer
    net = model.GCNModel(dim_latent, num_layers, do_train=True).to(device)
    optimizer = optim.Adam(net.parameters(), lr=learning_rate)

    best_net = model.GCNModel(dim_latent, num_layers, do_train=True)
    best_net.load_state_dict(copy.deepcopy(net.state_dict()))

    loss_per_epoch_train, loss_per_epoch_valid = [], []
    f1_per_epoch_train, f1_per_epoch_valid = [], []

    criterion = nn.BCEWithLogitsLoss(reduction='none')
    weight = torch.tensor([0.00001, 0.99999]).to(device)

    # Start training
    for epoch in range(num_epochs):

        # Training iteration
        loss_per_graph = []
        f1_per_graph = []
        net.train()
        for data in dl_train:
            graph, name = data
            name = name[0]
            logits = net(graph)
            labels = graph.ndata['significance'].unsqueeze(-1)
            weight_ = weight[labels.data.view(-1).long()].view_as(labels)

            # Get weighted loss
            loss = criterion(logits, labels)
            loss_weighted = loss * weight_
            loss_weighted = loss_weighted.mean()

            # Update parameters
            optimizer.zero_grad()
            loss_weighted.backward()
            optimizer.step()

            # Append output metrics
            loss_per_graph.append(loss_weighted.item())
            preds = logits.sigmoid().squeeze(1).int()
            labels = labels.squeeze(1).int()
            f1 = metrics.f1_score(labels, preds)
            f1_per_graph.append(f1)

        # Output loss and f1
        running_loss = np.array(loss_per_graph).mean()
        running_f1 = np.array(f1_per_graph).mean()
        loss_per_epoch_train.append(running_loss)
        f1_per_epoch_train.append(running_f1)

        if (epoch + 1) % 20 == 0:
            print(f'Epoch: {epoch+1}\tTraining loss:\t\t{running_loss}')

        # Validation iteration
        with torch.no_grad():
            loss_per_graph = []
            f1_per_graph = []
            net.eval()
            for data in dl_valid:
                graph, name = data
                name = name[0]
                logits = net(graph)
                labels = graph.ndata['significance'].unsqueeze(-1)

                # Get weighted loss
                weight_ = weight[labels.data.view(-1).long()].view_as(labels)
                loss = criterion(logits, labels)
                loss_weighted = loss * weight_
                loss_weighted = loss_weighted.mean()

                # Append output metrics
                loss_per_graph.append(loss_weighted.item())
                preds = logits.sigmoid().squeeze(1).int()
                labels = labels.squeeze(1).int()
                f1 = metrics.f1_score(labels, preds)
                f1_per_graph.append(f1)

        # Output loss and f1
        running_loss = np.array(loss_per_graph).mean()
        running_f1 = np.array(f1_per_graph).mean()
        loss_per_epoch_valid.append(running_loss)
        f1_per_epoch_valid.append(running_f1)

        # Save the best model
        if len(loss_per_epoch_valid) > 0 and running_loss < min(loss_per_epoch_valid):
            best_net.load_state_dict(copy.deepcopy(net.state_dict()))

        if (epoch+1) % 20 == 0:
            print(f'Epoch: {epoch+1}\tValidation loss:\t{running_loss}')
            print(f'Epoch: {epoch+1}\tF1 score:\t\t{running_f1}\n')

    # Plot loss
    if plot:
        plot_path = os.path.join(data_path, 'figures')
        loss_path = os.path.join(plot_path, f'loss_dim{dim_latent}_lay{num_layers}.png')
        f1_path = os.path.join(plot_path, f'f1_dim{dim_latent}_lay{num_layers}.png')
        draw_loss_plot(loss_per_epoch_train, loss_per_epoch_valid, loss_path)
        draw_f1_plot(f1_per_epoch_train, f1_per_epoch_valid, f1_path)

    # Save the best model
    torch.save(best_net.state_dict(), model_path)

    return model_path

if __name__ == '__main__':
    train()


from typing import Any, Callable
import unittest
from os.path import join
from tempfile import TemporaryDirectory
from uuid import uuid4
import logging

import numpy as np
import torch

from rastervision.core.data import (
    ClassConfig, DatasetConfig, RasterioSourceConfig, MultiRasterSourceConfig,
    SubRasterSourceConfig, ReclassTransformerConfig, SceneConfig,
    LabelSourceConfig)
from rastervision.pytorch_learner import (
    RegressionModelConfig, SolverConfig, RegressionGeoDataConfig,
    GeoDataWindowConfig, RegressionLearnerConfig, RegressionPlotOptions)
from tests import data_file_path

class MockRegressionabelSourceConfig(LabelSourceConfig):
    def build(self, *args, **kwargs):
        pass

def make_scene(num_channels: int, num_classes: int,
               tmp_dir: str) -> SceneConfig:
    path = data_file_path('multi_raster_source/const_100_600x600.tiff')
    rs_cfgs_img = []
    for _ in range(num_channels):
        rs_cfg = RasterioSourceConfig(
            uris=[path],
            channel_order=[0],
            transformers=[
                ReclassTransformerConfig(
                    mapping={100: np.random.randint(0, 256)})
            ])
        rs_cfgs_img.append(rs_cfg)
    rs_cfg_img = MultiRasterSourceConfig(
        raster_sources=[
            SubRasterSourceConfig(raster_source=rs_cfg, target_channels=[i])
            for i, rs_cfg in enumerate(rs_cfgs_img)
        ],
        channel_order=list(range(num_channels)))

    scene_cfg = SceneConfig(
        id=str(uuid4()),
        raster_source=rs_cfg_img,
        label_source=MockRegressionabelSourceConfig())
    return scene_cfg

class TestClassificationLearner(unittest.TestCase):
    def assertNoError(self, fn: Callable, msg: str = ''):
        try:
            fn()
        except Exception:
            self.fail(msg)

    def test_learner(self):
        self.assertNoError(lambda: self._test_learner(3, None))
        self.assertNoError(
            lambda: self._test_learner(6, [(0, 1, 2), (3, 4, 5)]))

    def _test_learner(self,
                      num_channels: int,
                      channel_display_groups: Any,
                      num_classes: int = 5):
        """Tests whether the learner can be instantiated correctly and
        produce plots."""
        logging.disable(logging.CRITICAL)

        with TemporaryDirectory() as tmp_dir:
            class_config = ClassConfig(
                names=[f'class_{i}' for i in range(num_classes)])
            dataset_cfg = DatasetConfig(
                class_config=class_config,
                train_scenes=[
                    make_scene(num_channels, num_classes, tmp_dir)
                    for _ in range(2)
                ],
                validation_scenes=[
                    make_scene(num_channels, num_classes, tmp_dir)
                    for _ in range(2)
                ],
                test_scenes=[])
            data_cfg = RegressionGeoDataConfig(
                scene_dataset=dataset_cfg,
                img_channels=num_channels,
                window_opts=GeoDataWindowConfig(size=20, stride=20),
                class_names=class_config.names,
                class_colors=class_config.colors,
                plot_options=RegressionPlotOptions(
                    channel_display_groups=channel_display_groups),
                num_workers=0)

            learner_cfg = RegressionLearnerConfig(
                output_uri=tmp_dir,
                data=data_cfg,
                model=RegressionModelConfig(pretrained=False),
                solver=SolverConfig(),
                log_tensorboard=False)
            data_cfg.update(learner_cfg)
            learner = learner_cfg.build(tmp_dir, training=True)
            x = torch.rand((4, num_channels, 100, 100))
            y = torch.rand((4, num_classes))
            z = torch.rand((4, num_classes))
            learner.plot_batch(x, y, join(tmp_dir, '1.png'))
            learner.plot_batch(x, y, join(tmp_dir, '2.png'), z=z)

if __name__ == '__main__':
    unittest.main()


import argparse
import network_utils
import helper_utils
import torch
import json
import numpy as np

def get_args():
    parser = argparse.ArgumentParser(description="Predict flower classification with DNN")
    parser.add_argument('input', default='./flowers/test/17/image_03911.jpg', type=str, help="input flower image to predict")
    parser.add_argument('checkpoint', type=str, help='pre-trained model path')
    parser.add_argument('--top_k', default=3, type=int, help='default top_k results')
    parser.add_argument('--category_names', default='./cat_to_name.json', type=str, help='default category file')
    parser.add_argument('--gpu', default='False',type=str, help='If GPU should be enabled')
    return parser.parse_args()

def predict(image, model, use_gpu, topk):
    #Predict the class (or classes) of an image using a trained deep learning model.

    if helper_utils.str_to_bool(use_gpu) and torch.cuda.is_available():
        image = torch.from_numpy(image).type(torch.cuda.FloatTensor)  # Convert numpy to tensor
        model.cuda()
        print("GPU active")
    else:
        image = torch.from_numpy(image).type(torch.FloatTensor)  # Convert numpy to tensor
        model.cpu()
        print("CPU active: Either cuda is not available or gpu option has been turn off")
    model.eval()  # set model to evaluation mode

    image = torch.unsqueeze(image, dim=0)                    # form a column tensor
    with torch.no_grad ():                                   # Turn off gradient
        output = model.forward(image)
    preds, classes = torch.exp(output).topk(topk)            # Get prediction and classes of top 5
    probs = preds.cpu().numpy().tolist()[0]
    classes = classes.cpu().numpy().tolist()[0]

    idx_to_class = {model.class_to_idx[k]: k for k in model.class_to_idx}
    topk_classes = [idx_to_class[i] for i in classes]

    return probs, topk_classes

def main():
    args = get_args()
    processed_image = helper_utils.process_image(args.input) # Process the image to numpy array
    model = network_utils.loading_model(args.checkpoint)

    probs, topk_classes = predict(processed_image, model, args.gpu, args.top_k)
    cat_to_name = helper_utils.load_cat_to_name(args.category_names)
    class_names = [cat_to_name [item] for item in topk_classes]
    max_prob_idx = np.argmax(probs)
    max_class_nb = topk_classes[max_prob_idx]
    predicted_class = cat_to_name[max_class_nb]
    #helper_utils.display_image(args.input,predicted_class)

    helper_utils.display_result(class_names, probs)

if __name__ == '__main__':
    main()


#!/usr/bin/env python

__version__ = "0.0.0-dev"

import random
import io
import itertools
from collections import defaultdict, namedtuple
import hashlib
import os
import shutil
import glob
import math
from functools import partial

import numpy as np
import pandas as pd

import ipywidgets
from ipywidgets import interactive, fixed, IntSlider
from IPython.display import display
from scipy.optimize import minimize_scalar

import skbio
from skbio.diversity import alpha_diversity, beta_diversity
from skbio.stats.ordination import pcoa
from skbio.stats import subsample_counts
from skbio.util import safe_md5

from q2d2.wui import metadata_controls
from q2d2.registry import get_workflows

data_type_to_study_filename = {'sample_metadata': '.sample-md',
                               'otu_metadata': '.otu-md',
                               'unrarefied_biom': '.biom',
                               'rarefied_biom': '.rarefied-biom',
                               'tree': '.tree'}

workflows = {w.id: w for w in get_workflows()}

def get_data_info(study_id='.'):
    existing_data_types = get_existing_data_types(study_id)
    data_info = []
    for data_type in data_type_to_study_filename:
        filename = data_type_to_study_filename[data_type]
        exists = data_type in existing_data_types
        data_info.append((data_type, filename, exists))
    return data_info

def get_workflow_info(workflow_id):
    workflow = workflows[workflow_id]
    return {
        'workflow-id': workflow_id,
        'title': workflow.title,
        'inputs': list(workflow.inputs),
        'outputs': list(workflow.outputs)
    }

def get_study_state(study_id='.'):
    existing_data_types = get_existing_data_types(study_id)

    state = {
        'study-id': study_id,
        'workflow': {'exe': [], 'nexe': []},
        'data': {}
    }

    for workflow_id in workflows:
        workflow = workflows[workflow_id]
        if workflow.inputs.issubset(existing_data_types):
            state['workflow']['exe'].append(workflow_id)
        else:
            state['workflow']['nexe'].append(workflow_id)

    for data_type in existing_data_types:
        data_filepath = get_data_filepath(data_type, study_id)
        with open(data_filepath, 'rb') as data_file:
            # should we be using sha256 instead?
            md5 = safe_md5(data_file).hexdigest()
        state['data'][data_filepath] = md5

    return state

def get_system_info():
    # what other info goes here? dependencies?
    return {'version': __version__}

def get_existing_data_types(study_id='.'):
    data_types = set()
    for data_type in data_type_to_study_filename:
        try:
            get_data_filepath(data_type, study_id)
        except FileNotFoundError:
            pass
        else:
            data_types.add(data_type)
    return data_types

def create_index(study_id, command):
    markdown_s = get_index_markdown(study_id, command)
    output_filepath = os.path.join(study_id, 'index.md')
    open(output_filepath, 'w').write(markdown_s)

def get_data_filepath(data_type, study_id='.'):
    data_filepath = os.path.join(study_id, data_type_to_study_filename[data_type])
    if not os.path.exists(data_filepath):
        raise FileNotFoundError(data_filepath)
    return data_filepath

def create_input_files(study_id='.', **kwargs):
    for input_type, input_filepath in kwargs.items():
        study_filepath = data_type_to_study_filename[input_type]
        study_filepath = os.path.join(study_id, study_filepath)
        shutil.copy(input_filepath, study_filepath)

def load_table(rarefied=False):
    if rarefied:
        table_path = data_type_to_study_filename['rarefied_biom']
    else:
        table_path = data_type_to_study_filename['unrarefied_biom']
    result = pd.read_csv(table_path, sep='\t', skiprows=1, index_col=0)
    result.index = result.index.astype(str)
    if 'taxonomy' in result:
        result.drop('taxonomy', axis=1, inplace=True)
    return result

def store_table(table, rarefied=False):
    if rarefied:
        table_path = data_type_to_study_filename['rarefied_biom']
    else:
        table_path = data_type_to_study_filename['unrarefied_biom']
    with open(table_path, 'w') as table_file:
        table_file.write('# Constructed by [q2d2](github.com/gregcaporaso/q2d2)\n')
        table.to_csv(table_file, index_label="#OTU ID", sep='\t')

load_rarefied_table = partial(load_table, rarefied=True)
store_rarefied_table = partial(store_table, rarefied=True)

def load_tree():
    return skbio.TreeNode.read(data_type_to_study_filename['tree'], format='newick')

def load_sample_metadata():
    return pd.read_csv(data_type_to_study_filename['sample_metadata'], sep='\t', index_col=0, dtype=object)

def load_otu_metadata():
    return pd.read_csv(data_type_to_study_filename['otu_metadata'], sep='\t', names=['OTU ID', 'taxonomy'],
                       index_col=0, usecols=[0, 1], dtype=object)

def compute_alphas(biom, tree=None,
                   metrics=['chao1',
                            'faith_pd',
                            'observed_otus']):
    alphas = {}
    for metric in metrics:
        if metric == 'faith_pd':
            alphas[metric] = alpha_diversity(metric, counts=np.asarray(biom.T),
                                             ids=biom.columns,
                                             otu_ids=biom.index, tree=tree)
        else:
            alphas[metric] = alpha_diversity(metric, counts=np.asarray(biom.T),
                                                 ids=biom.columns)
    return alphas

# def biom_to_dm(metric, biom, tree=None):
#     if metric in ['unweighted_unifrac', 'weighted_unifrac']:
#         return beta_diversity(metric, counts=np.asarray(biom.T),
#                               ids=biom.columns, otu_ids=biom.index, tree=tree)
#     else:
#         return beta_diversity(metric, counts=np.asarray(biom.T),
#                               ids=biom.columns)

def dm_to_pcoa(dm, sample_md, category):
    title = "Samples colored by %s." % category
    pcoa_results = pcoa(dm)
    _ = pcoa_results.plot(df=sample_md,
                          column=category,
                          axis_labels=['PC 1', 'PC 2', 'PC 3'],
                          title=title,
                          s=35)

def table_summary(df):
    print("Samples: ", len(df.columns))
    print("Observations: ", len(df.index))
    print("Sequence/sample count detail:")
    print(df.sum().describe())

def get_workflow_template_filepath(workflow_id):
    base_dir = os.path.abspath(os.path.split(__file__)[0])
    return os.path.join(base_dir, "markdown", "%s.md" % workflow_id)

def create_workflow(workflow_id, study_id='.'):
    return workflows[workflow_id].create_workflow(study_id)

def delete_workflow(workflow_id, study_id='.'):
    workflows[workflow_id].delete_workflow(study_id)

def get_index_markdown(study_id, command):
    index_md_template = open(get_workflow_template_filepath('index')).read()
    md_fps = glob.glob(os.path.join(study_id, '*.md'))
    md_fps.sort()
    toc = []
    for md_fp in md_fps:
        md_fn = os.path.split(md_fp)[1]
        title = os.path.splitext(md_fn)[0].replace('-', ' ').title()
        toc.append(' * [%s](%s)' % (title, md_fn))
    toc = '\n'.join(toc)
    result = index_md_template.format(toc, study_id, __version__, command)
    return result

def _summarize_even_sampling_depth(even_sampling_depth, counts):
    samples_retained = (counts >= even_sampling_depth)
    num_samples_retained = samples_retained.sum()
    num_sequences_retained = num_samples_retained * even_sampling_depth
    return samples_retained, num_samples_retained, num_sequences_retained

def _get_depth_for_max_sequence_count(counts):
    """Find the even sampling depth that retains the most sequences."""
    count_summary = counts.describe()
    def f(d):
        return -1 * _summarize_even_sampling_depth(d, counts)[2]

    res = minimize_scalar(f,
                          bounds=(count_summary['min'], count_summary['max']),
                          method='bounded')
    return int(np.floor(res.x))

def get_default_even_sampling_depth(biom):
    counts = biom.sum()
    return _get_depth_for_max_sequence_count(counts)

def explore_sampling_depth(biom):
    import seaborn as sns
    counts = biom.sum()
    count_summary = counts.describe()
    total_num_samples = len(counts)
    total_num_sequences = counts.sum()
    depth_for_max_sequence_count = _get_depth_for_max_sequence_count(counts)
    sampling_depth_slider = IntSlider(min=count_summary['min'],
                                      max=count_summary['max'],
                                      step=10 ** (math.log(count_summary['max'], 10) - 2),
                                      value=depth_for_max_sequence_count)
    default_samples_retained, default_num_samples_retained, default_num_sequences_retained = \
            _summarize_even_sampling_depth(depth_for_max_sequence_count, counts)

    default_percent_samples_retained = default_num_samples_retained * 100 / total_num_samples
    default_percent_sequences_retained = default_num_sequences_retained * 100 / total_num_sequences

    label_s = "Depth {0}: {1:.2f}% of sequences and {2:.2f}% of samples retained."

    def f(even_sampling_depth):
        samples_retained, num_samples_retained, num_sequences_retained = \
            _summarize_even_sampling_depth(even_sampling_depth, counts)
        percent_samples_retained = num_samples_retained * 100 / total_num_samples
        percent_sequences_retained = num_sequences_retained * 100 / total_num_sequences
        ax = sns.distplot(counts)
        ax.set_xlabel("Number of sequences per sample")
        ax.set_ylabel("Frequency")
        line_label = label_s.format(depth_for_max_sequence_count,
                                    default_percent_sequences_retained,
                                    default_percent_samples_retained)
        ax.plot([depth_for_max_sequence_count, depth_for_max_sequence_count], ax.get_ylim(),
                'k--', label=line_label)

        line_label = label_s.format(even_sampling_depth,
                                    percent_sequences_retained,
                                    percent_samples_retained)
        ax.plot([even_sampling_depth, even_sampling_depth], ax.get_ylim(),
                'k-', label=line_label)
        ax.legend(loc='center left', bbox_to_anchor=(1, 0.5))

    def reset_depth(_):
        sampling_depth_slider.value = depth_for_max_sequence_count

    reset = ipywidgets.Button(icon='fa-refresh')
    reset.on_click(reset_depth)

    w = interactive(f, even_sampling_depth=sampling_depth_slider)
    display(ipywidgets.HBox(children=[w, reset]))

def rarify(biom, even_sampling_depth):
    data = []
    sample_ids = []
    for e in biom.columns:
        count_vector = biom[e]
        if count_vector.sum() < even_sampling_depth:
            continue
        else:
            sample_ids.append(e)
            data.append(subsample_counts(count_vector.astype(int), even_sampling_depth))
    return pd.DataFrame(np.asarray(data).T, index=biom.index, columns=sample_ids)

def filter_dm_and_map(dm, map_df):
    ids_to_exclude = set(dm.ids) - set(map_df.index.values)
    ids_to_keep = set(dm.ids) - ids_to_exclude
    filtered_dm = dm.filter(ids_to_keep)
    filtered_map = map_df.loc[ids_to_keep]

    return filtered_dm, filtered_map

def get_within_between_distances(map_df, dm, col):
    filtered_dm, filtered_map = filter_dm_and_map(dm, map_df)
    groups = []
    distances = []
    map_dict = filtered_map[col].to_dict()
    for id_1, id_2 in itertools.combinations(filtered_map.index.tolist(), 2):
        row = []
        if map_dict[id_1] == map_dict[id_2]:
            groups.append('Within')
        else:
            groups.append('Between')
        distances.append(filtered_dm[(id_1, id_2)])
    groups = zip(groups, distances)
    distances_df = pd.DataFrame(data=list(groups), columns=['Groups', 'Distance'])

    return distances_df

def distance_histogram(dm, category, metadata, metric='Distance', order=['Within', 'Between']):
    import seaborn as sns
    within_bw_distances = get_within_between_distances(metadata, dm, category)
    ax = sns.violinplot(x='Groups', y='Distance', data=within_bw_distances, order=order, orient='v')
    ax.set_xlabel(category)
    ax.set_ylabel(metric)

def interactive_distance_histograms(dm, sample_metadata):
    def on_update(category, metadata, check_within, check_between):
        order = []
        if check_within:
            order.append('Within')
        if check_between:
            order.append('Between')
        distance_histogram(dm, category, metadata, order=order)
    check_within = ipywidgets.Checkbox(description='Show within category', value=True)
    check_between = ipywidgets.Checkbox(description='Show between category', value=True)
    extras = ipywidgets.VBox(children=[check_within, check_between])
    return metadata_controls(sample_metadata, on_update, extras)

def distance_violinplots(dm, category, metadata, metric=None, order=['Within', 'Between']):
    import seaborn as sns
    within_bw_distances = get_within_between_distances(metadata, dm, category)
    ax = sns.violinplot(x='Groups', y='Distance', data=within_bw_distances, order=order, orient='v')
    ax.set_xlabel(category)
    ax.set_ylabel(metric)
    return ax

def interactive_distance_violinplots(dms, sample_metadata):

    def on_update(category, metadata, metric, check_within, check_between):
        order = []
        if check_within:
            order.append('Within')
        if check_between:
            order.append('Between')

        dm = dms[metric]
        distance_violinplots(dm, category, metadata, metric, order=order)

    check_within = ipywidgets.Checkbox(description='Show within category', value=True)
    check_between = ipywidgets.Checkbox(description='Show between category', value=True)
    metric_but = ipywidgets.Dropdown(options=list(dms.keys()), description='Metrics')

    extras = ipywidgets.VBox(children=[metric_but, check_within, check_between])
    return metadata_controls(sample_metadata, on_update, extras)

def compute_distance_matrices(
               biom,
               tree=None,
               metrics=['weighted_unifrac', 'unweighted_unifrac', 'braycurtis', 'jaccard']):
    dms = {}
    for metric in metrics:
        if metric in ['unweighted_unifrac', 'weighted_unifrac']:
            dms[metric] = beta_diversity(metric, counts=np.asarray(biom.T),
                                         ids=biom.columns, otu_ids=biom.index,
                                         tree=tree)
        else:
            dms[metric] = beta_diversity(metric, counts=np.asarray(biom.T),
                                         ids=biom.columns)
    return dms

def interactive_plot_pcoa(metadata, dms):

    def on_update(category, metadata, metric):
        dm = dms[metric]
        filtered_dm, _ = filter_dm_and_map(dm, metadata)
        pc = pcoa(filtered_dm)
        pc.plot(df=metadata,
        column=category,
        axis_labels=['PC 1', 'PC 2', 'PC 3'],
        s=35).set_size_inches(12, 9)

    metric_but = ipywidgets.Dropdown(options=list(dms.keys()), description='Metrics')
    extras = ipywidgets.VBox(children=[metric_but])

    return metadata_controls(metadata, on_update, extras)


import torch
import sys, os
sys.path.insert(0, os.path.dirname(__file__))
from LaFan import LaFan1
from torch.utils.data import Dataset, DataLoader
from model import StateEncoder, \
                  OffsetEncoder, \
                  TargetEncoder, \
                  LSTM, \
                  Decoder, \
                  ShortMotionDiscriminator, \
                  LongMotionDiscriminator
from skeleton import Skeleton
import torch.optim as optim
from tensorboardX import SummaryWriter
import numpy as np
from tqdm import tqdm
from functions import gen_ztta, write_to_bvhfile
import yaml
import time
import shutil
import imageio
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import axes3d, Axes3D
from remove_fs import remove_fs, save_bvh_from_network_output
from foot_sliding.animation_data import y_rotation_from_positions
from PIL import Image
from py_utils import mkd, count_param
import argparse

"""

"""

def plot_pose(pose, cur_frame, prefix):

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    parents = [-1, 0, 1, 2, 3, 0, 5, 6, 7, 0, 9, 10, 11, 12, 11, 14, 15, 16, 11, 18, 19, 20]
    ax.cla()
    num_joint = pose.shape[0] // 3
    for i, p in enumerate(parents):
        if i > 0:
            ax.plot([pose[i, 0], pose[p, 0]],\
                    [pose[i, 2], pose[p, 2]],\
                    [pose[i, 1], pose[p, 1]], c='r')
            ax.plot([pose[i+num_joint, 0], pose[p+num_joint, 0]],\
                    [pose[i+num_joint, 2], pose[p+num_joint, 2]],\
                    [pose[i+num_joint, 1], pose[p+num_joint, 1]], c='b')
            ax.plot([pose[i+num_joint*2, 0], pose[p+num_joint*2, 0]],\
                    [pose[i+num_joint*2, 2], pose[p+num_joint*2, 2]],\
                    [pose[i+num_joint*2, 1], pose[p+num_joint*2, 1]], c='g')
    # ax.scatter(pose[:num_joint, 0], pose[:num_joint, 2], pose[:num_joint, 1],c='b')
    # ax.scatter(pose[num_joint:num_joint*2, 0], pose[num_joint:num_joint*2, 2], pose[num_joint:num_joint*2, 1],c='b')
    # ax.scatter(pose[num_joint*2:num_joint*3, 0], pose[num_joint*2:num_joint*3, 2], pose[num_joint*2:num_joint*3, 1],c='g')
    xmin = np.min(pose[:, 0])
    ymin = np.min(pose[:, 2])
    zmin = np.min(pose[:, 1])
    xmax = np.max(pose[:, 0])
    ymax = np.max(pose[:, 2])
    zmax = np.max(pose[:, 1])
    scale = np.max([xmax - xmin, ymax - ymin, zmax - zmin])
    xmid = (xmax + xmin) // 2
    ymid = (ymax + ymin) // 2
    zmid = (zmax + zmin) // 2
    ax.set_xlim(xmid - scale // 2, xmid + scale // 2)
    ax.set_ylim(ymid - scale // 2, ymid + scale // 2)
    ax.set_zlim(zmid - scale // 2, zmid + scale // 2)

    plt.draw()
    plt.savefig(prefix + '_' + str(cur_frame)+'.png', dpi=200, bbox_inches='tight')
    plt.close()

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--cfg', type=str, default='test-base.yaml')
    args = parser.parse_args()

    opt = yaml.load(open('./config/' + args.cfg, 'r').read())
    model_dir =opt['test']['model_dir']

    ## initilize the skeleton ##
    skeleton_mocap = Skeleton(offsets=opt['data']['offsets'], parents=opt['data']['parents'])
    skeleton_mocap.cuda()
    skeleton_mocap.remove_joints(opt['data']['joints_to_remove'])

    ## load train data ##
    lafan_data_test = LaFan1(opt['data']['data_dir'], \
                              seq_len = opt['model']['seq_length'], \
                              offset = opt['data']['offset'],\
                              train = opt['test']['save_pose'], debug=opt['test']['debug'])
    lafan_data_test.cur_seq_length = opt['model']['seq_length']
    x_mean = lafan_data_test.x_mean.cuda()
    x_std = lafan_data_test.x_std.cuda().view(1, 1, opt['model']['num_joints'], 3)
    lafan_loader_test = DataLoader(lafan_data_test, \
                                    batch_size=opt['test']['batch_size'], \
                                    shuffle=False, num_workers=opt['data']['num_workers'])

    ## initialize model and load parameters ##
    state_encoder = StateEncoder(in_dim=opt['model']['state_input_dim'])
    state_encoder = state_encoder.cuda()
    state_encoder.load_state_dict(torch.load(os.path.join(opt['test']['model_dir'], 'state_encoder.pkl')))
    offset_encoder = OffsetEncoder(in_dim=opt['model']['offset_input_dim'])
    offset_encoder = offset_encoder.cuda()
    offset_encoder.load_state_dict(torch.load(os.path.join(opt['test']['model_dir'], 'offset_encoder.pkl')))
    target_encoder = TargetEncoder(in_dim=opt['model']['target_input_dim'])
    target_encoder = target_encoder.cuda()
    target_encoder.load_state_dict(torch.load(os.path.join(opt['test']['model_dir'], 'target_encoder.pkl')))
    lstm = LSTM(in_dim=opt['model']['lstm_dim'], hidden_dim = opt['model']['lstm_dim'] * 2)
    lstm = lstm.cuda()
    lstm.load_state_dict(torch.load(os.path.join(opt['test']['model_dir'], 'lstm.pkl')))
    decoder = Decoder(in_dim=opt['model']['lstm_dim'] * 2, out_dim=opt['model']['state_input_dim'])
    decoder = decoder.cuda()
    decoder.load_state_dict(torch.load(os.path.join(opt['test']['model_dir'], 'decoder.pkl')))
    print('model loaded')

    count_param(state_encoder, 'state_encoder')
    count_param(offset_encoder, 'offset_encoder')
    count_param(target_encoder, 'target_encoder')
    count_param(lstm, 'lstm')
    count_param(decoder, 'decoder')

    ## get positional code ##
    if opt['test']['use_ztta']:
        ztta = gen_ztta().cuda()
    # print('ztta:', ztta.size())
    # assert 0
    version = opt['test']['version']

    # writer = SummaryWriter(log_dir)
    loss_total_min = 10000000.0
    for epoch in range(opt['test']['num_epoch']):
        state_encoder.eval()
        offset_encoder.eval()
        target_encoder.eval()
        lstm.eval()
        decoder.eval()
        loss_total_list = []

        for i_batch, sampled_batch in enumerate(lafan_loader_test):
            # if i_batch != 33:
            #     continue
            pred_img_list = []
            gt_img_list = []
            img_list = []

            # print(i_batch, sample_batched['local_q'].size())

            # loss_pos = 0
            # loss_quat = 0
            # loss_contact = 0
            # loss_root = 0
            with torch.no_grad():
                # if True:
                # state input
                local_q = sampled_batch['local_q'].cuda()  # 32 x 50 x 22 x 4
                root_v = sampled_batch['root_v'].cuda()
                contact = sampled_batch['contact'].cuda()
                # offset input
                root_p_offset = sampled_batch['root_p_offset'].cuda()  # 32 x 3
                local_q_offset = sampled_batch['local_q_offset'].cuda()  # 32 x 88
                local_q_offset = local_q_offset.view(local_q_offset.size(0), -1)
                # target input
                target = sampled_batch['target'].cuda()
                target = target.view(target.size(0), -1)
                # root pos
                root_p = sampled_batch['root_p'].cuda()  # 32 x 50 x 3
                # X
                X = sampled_batch['X'].cuda()
                bs = 6#np.random.choice(X.size(0), 1)[0]
                # if False:
                #     print('local_q:', local_q.size(), \
                #         'root_v:', root_v.size(), \
                #         'contact:', contact.size(), \
                #         'root_p_offset:', root_p_offset.size(), \
                #         'local_q_offset:', local_q_offset.size(), \
                #         'target:', target.size())
                #     assert 0

                lstm.init_hidden(local_q.size(0))
                h_list = []
                quat_list = []
                quat_list.append(local_q[:,0,].view(local_q.size(0), -1, 4))
                pred_list = []
                pred_list.append(X[:,0])
                bvh_list = []
                bvh_list.append(torch.cat([X[:,0,0], local_q[:,0,].view(local_q.size(0), -1)], -1))
                contact_list = []
                contact_list.append(contact[:,0])
                root_list = []
                root_list.append(X[:,0,0])
                # print(X.size())
                # for t in range(opt['model']['seq_length'] - 1):
                for txn in range((opt['model']['seq_length'] - 1) * 10):
                    if txn < (opt['model']['seq_length'] - 1):
                        t = txn
                        # root pos
                        if t  == 0:
                            root_p_t = root_p[:,t]
                            local_q_t = local_q[:,t]
                            local_q_t = local_q_t.view(local_q_t.size(0), -1)
                            contact_t = contact[:,t]
                            root_v_t = root_v[:,t]
                        else:
                            root_p_t = root_pred[0]
                            local_q_t = local_q_pred[0]
                            contact_t = contact_pred[0]
                            root_v_t = root_v_pred[0]
                    else:
                    # if True:
                        # reset t
                        t = txn % (opt['model']['seq_length'] - 1)

                        if t == 0:
                            # reset target pose,
                            sz = local_q.shape  # 32 x 50 x 22 x 4
                            random_idx = np.random.randint(low=[0, 1], high=[sz[0], sz[1]], size=(sz[0], 2))

                            # offset input
                            # root_p_offset = sampled_batch['root_p_offset'].cuda()
                            tmp_root_pred = root_pred * 1.   # 1x32x3
                            tmp_root_p_offset = root_p[random_idx[:, 0], random_idx[:, 1]] * 1.
                            random_mv = (torch.rand_like(root_p_offset) - 0.5) * 2.
                            random_mv[:, 0] = random_mv[:, 0] * 100. + tmp_root_pred[0, :, 0] * 1.
                            random_mv[:, 1] = random_mv[:, 0] * 0.   + tmp_root_p_offset[:, 1] * 1.  # assume height no change
                            random_mv[:, 2] = random_mv[:, 2] * 100. + tmp_root_pred[0, :, 2] * 1.
                            root_p_offset = random_mv

                            # local_q_offset = sampled_batch['local_q_offset'].cuda()
                            local_q_offset = local_q[random_idx[:, 0], random_idx[:, 1]] * 1.
                            local_q_offset = local_q_offset.view(local_q_offset.size(0), -1)
                            # target input
                            # target = sampled_batch['target'].cuda()
                            # target = target.view(target.size(0), -1)
                            target = local_q_offset * 1.

                            # reset lstm
                            lstm.init_hidden(local_q.size(0))

                            # target pose modify
                            pos_target = skeleton_mocap.forward_kinematics(local_q_offset.view(1, -1, 22, 4), root_p_offset.view(1, -1, 3))
                            X[:, 0] = pos_pred[0, :] * 1.
                            X[:,-1] = pos_target[0, :] * 1.

                        # continue on pose seq
                        root_p_t = root_pred[0]
                        local_q_t = local_q_pred[0]
                        contact_t = contact_pred[0]
                        root_v_t = root_v_pred[0]

                    # state input
                    state_input = torch.cat([local_q_t, root_v_t, contact_t], -1)
                    # offset input
                    root_p_offset_t = root_p_offset - root_p_t
                    local_q_offset_t = local_q_offset - local_q_t
                    # print('root_p_offset_t:', root_p_offset_t.size(), 'local_q_offset_t:', local_q_offset_t.size())
                    offset_input = torch.cat([root_p_offset_t, local_q_offset_t], -1)
                    # target input
                    target_input = target

                    # print('state_input:',state_input.size())
                    h_state = state_encoder(state_input)
                    h_offset = offset_encoder(offset_input)
                    h_target = target_encoder(target_input)

                    if opt['test']['use_ztta']:
                        h_state += ztta[:, t]
                        h_offset += ztta[:, t]
                        h_target += ztta[:, t]

                    if opt['test']['use_adv']:
                        tta = opt['model']['seq_length'] - 2 - t
                        if tta < 5:
                            lambda_target = 0.0
                        elif tta >=5 and tta < 30:
                            lambda_target = (tta - 5) / 25.0
                        else:
                            lambda_target = 1.0
                        h_offset += 0.5 * lambda_target * torch.cuda.FloatTensor(h_offset.size()).normal_()
                        h_target += 0.5 * lambda_target * torch.cuda.FloatTensor(h_target.size()).normal_()

                    h_in = torch.cat([h_state, h_offset, h_target], -1).unsqueeze(0)
                    h_out = lstm(h_in)
                    # print('h_out:', h_out.size())

                    h_pred, contact_pred = decoder(h_out)
                    local_q_v_pred = h_pred[:,:,:opt['model']['target_input_dim']]
                    local_q_pred = local_q_v_pred + local_q_t
                    # print('q_pred:', q_pred.size())
                    local_q_pred_ = local_q_pred.view(local_q_pred.size(0), local_q_pred.size(1), -1, 4)
                    local_q_pred_ = local_q_pred_ / torch.norm(local_q_pred_, dim = -1, keepdim = True)
                    # print("local_q_pred_:", local_q_pred_.size())
                    quat_list.append(local_q_pred_[0])
                    root_v_pred = h_pred[:,:,opt['model']['target_input_dim']:]  # 1x32x3
                    root_pred = root_v_pred + root_p_t  # 1x32x3
                    root_list.append(root_pred[0])
                    # print(''contact:'', contact_pred.size())
                    # print('root_pred:', root_pred.size())
                    bvh_list.append(torch.cat([root_pred[0], local_q_pred_[0].view(local_q_pred_.size(1), -1)], -1))
                    pos_pred = skeleton_mocap.forward_kinematics(local_q_pred_, root_pred)

                    # pos_next = X[:,t+1]
                    # local_q_next = local_q[:,t+1]
                    # local_q_next = local_q_next.view(local_q_next.size(0), -1)
                    # root_p_next = root_p[:,t+1]
                    # contact_next = contact[:,t+1]
                    # # print(pos_pred.size(), x_std.size())
                    # loss_pos += torch.mean(torch.abs(pos_pred[0] - pos_next) / x_std) / opt['model']['seq_length']
                    # loss_quat += torch.mean(torch.abs(local_q_pred[0] - local_q_next)) / opt['model']['seq_length']
                    # loss_root += torch.mean(torch.abs(root_pred[0] - root_p_next) / x_std[:,:,0]) / opt['model']['seq_length']
                    # loss_contact += torch.mean(torch.abs(contact_pred[0] - contact_next)) / opt['model']['seq_length']
                    pred_list.append(pos_pred[0])
                    contact_list.append(contact_pred[0])

                    # if i_batch < 49:
                        # print("pos_pred:", pos_pred.size())
                    if opt['test']['save_img']:
                        mkd('../results' + version)
                        plot_pose(np.concatenate([X[bs,0].view(22, 3).detach().cpu().numpy(),\
                                                pos_pred[0, bs].view(22, 3).detach().cpu().numpy(),\
                                                X[bs,-1].view(22, 3).detach().cpu().numpy()], 0),\
                                                t, '../results'+version+'/pred')
                        # plot_pose(np.concatenate([X[bs,0].view(22, 3).detach().cpu().numpy(),\
                        #                         X[bs,t+1].view(22, 3).detach().cpu().numpy(),\
                        #                         X[bs,-1].view(22, 3).detach().cpu().numpy()], 0),\
                        #                         t, '../results'+version+'/gt')
                        pred_img = Image.open('../results'+version+'/pred_'+str(t)+'.png', 'r')
                        # gt_img = Image.open('../results'+version+'/gt_'+str(t)+'.png', 'r')
                        pred_img_list.append(pred_img)
                        # gt_img_list.append(gt_img)
                        # img_list.append(np.concatenate([pred_img, gt_img.resize(pred_img.size)], 1))
                        # img_list.append(np.concatenate([pred_img, pred_img], 1))
                        img_list.append(np.array(pred_img))
                        # img_list.append(np.array(pred_img))

                # print('pivots:', pivots.shape)
                # print('rot_data.size:', rot_data.shape)
                if opt['test']['save_bvh']:
                    assert False, 'not work currently, very slow'
                    # print("bs:", bs)
                    bvh_data = torch.cat([x[bs].unsqueeze(0) for x in bvh_list], 0).detach().cpu().numpy()
                    # print('bvh_data:', bvh_data.shape)
                    # print('bvh_data:', bvh_data[0,3:7])
                    # assert 0
                    mkd('../bvh_seq')
                    write_to_bvhfile(bvh_data, ('../bvh_seq/test_%03d.bvh' % i_batch), opt['data']['joints_to_remove'])
                    # assert 0
                    contact_data = torch.cat([x[bs].unsqueeze(0) for x in contact_list], 0).detach().cpu().numpy()
                    # rot_data = torch.cat([x[bs].unsqueeze(0) for x in quat_list], 0).detach().cpu().numpy()
                    # root_data = torch.cat([x[bs].unsqueeze(0) for x in root_list], 0).detach().cpu().numpy()
                    # pred_pose = torch.cat([x[bs].unsqueeze(0) for x in pred_list], 0).detach().cpu().numpy()
                    # quaters, pivots = y_rotation_from_positions(pred_pose, hips = (1,5), sdrs = (14,18))
                    # motion = np.concatenate([rot_data.reshape(rot_data.shape[0], -1),\
                    #                          root_data,\
                    #                          pivots], -1)
                    # motion = motion.transpose(1,0)
                    foot = contact_data.transpose(1,0)
                    foot[foot > 0.5] = 1.0
                    foot[foot <= 0.5] = 0.0
                    # print('foot[0]:',foot[0])
                    mkd("../bvh_seq_after"+version)
                    glb = remove_fs(('../bvh_seq/test_%03d.bvh' % i_batch), \
                              foot, \
                              fid_l=(3, 4), \
                              fid_r=(7, 8),\
                              output_path=("../bvh_seq_after"+version+"/test_%03d.bvh" % i_batch))
                    fix_img_list = []
                    # for t in range(opt['model']['seq_length']):
                    for t in range((opt['model']['seq_length']-1)*10):
                        plot_pose(np.concatenate([X[bs,0].view(22, 3).detach().cpu().numpy(),\
                                                glb[t],\
                                                X[bs,-1].view(22, 3).detach().cpu().numpy()], 0),\
                                                t, '../results'+version+'/fixed')
                        # plot_pose(np.concatenate([X[bs,0].view(22, 3).detach().cpu().numpy(),\
                        #                         X[bs,t].view(22, 3).detach().cpu().numpy(),\
                        #                         X[bs,-1].view(22, 3).detach().cpu().numpy()], 0),\
                        #                         t, '../results'+version+'/gt')
                        fix_img = Image.open('../results'+version+'/fixed_'+str(t)+'.png', 'r')
                        # gt_img = Image.open('../results'+version+'/gt_'+str(t)+'.png', 'r')
                        # fix_img_list.append(np.concatenate([fix_img, gt_img.resize(fix_img.size)], 1))
                        fix_img_list.append(np.array(fix_img))
                    mkd('../gif_fixed'+version)
                    imageio.mimsave(('../gif_fixed'+version+'/img_fix_%03d.gif' % i_batch), fix_img_list, duration=0.1)
                    # save_bvh_from_network_output(motion, output_path=("../bvh_seq_after/test_%03d.bvh" % i_batch))

                # if i_batch < 49:
                if opt['test']['save_img'] and opt['test']['save_gif']:
                    mkd('../gif'+version)
                    imageio.mimsave(('../gif'+version+'/img_%03d.gif' % i_batch), img_list, duration=0.1)
                if opt['test']['save_pose']:
                    # gt_pose = X[bs,:].view(opt['model']['seq_length'], 22, 3).detach().cpu().numpy()
                    # pred_pose = torch.cat([x[bs].unsqueeze(0) for x in pred_list], 0).detach().cpu().numpy()
                    # plt.clf()
                    # joint_idx = 13
                    # plt.plot(range(opt['model']['seq_length']), gt_pose[:,joint_idx,0])
                    # plt.plot(range(opt['model']['seq_length']), pred_pose[:,joint_idx,0])
                    # plt.legend(['gt', 'pred'])
                    # plt.savefig('../results'+version+'/pose_%03d.png' % i_batch)
                    # plt.close()

                    pred_pose = torch.stack(pred_list, 1).detach().cpu().numpy()
                    npy_name = '../traj_pose'+version+'/batch_%03d.npy' % i_batch
                    mkd('../traj_pose'+version)
                    np.save(npy_name, pred_pose)

                # if opt['test']['save_img'] and i_batch > 49:
                #     break

                if opt['test']['save_pose'] and i_batch > 49:
                    break

        # print("train epoch: %03d, cur total loss:%.3f, cur best loss:%.3f" % (epoch, loss_total_cur, loss_total_min))


"""
MIT License

Copyright (c) 2021 Overcomer

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
"""

import os
import pickle
import logging
import cv2
import numpy as np
from imutils import resize
from sklearn.svm import SVC
from sklearn.preprocessing import LabelEncoder

from new_tools import check_image
from new_face.tools.config import root_dir
from new_face.tools.download import download_models

class OpenFace(object):
    def __init__(self):
        self.label_encoder = None
        self.embedder_model = None
        self.classifier = None

    def extract_embeddings(self,
                           images_path=list(),
                           vision=False):
        """
        Use OpenFace nn4.small2.v1.t7 network model to extract face feature and output feature vectors.

        Args
        ----
        embedding_network_path: OpenFace nn4.small2.v1.t7 network model path.

        images_path: Images path.

        vision: Show image.

        Return
        ------
        embedding_vectors: Face feature vectors.
        """

        logging.info("Extracting images...")

        total = 0
        embedding_vectors = list()

        for num, image_path in enumerate(images_path, start=1):
            logging.info("Processing {}/{} images...".format(num, len(images_path)))

            state, bgr_image = check_image(image_path)
            if state == 0:
                if vision:
                    cv2.imshow("Extracting image...", resize(bgr_image, width=250))
                    cv2.waitKey(10)

                face_blob = cv2.dnn.blobFromImage(bgr_image, 1.0/255, (96, 96), (0, 0, 0), swapRB=True, crop=False)

                self.embedder_model.setInput(face_blob)
                vec = self.embedder_model.forward()

                embedding_vectors.append(vec.flatten())
                total += 1

        if vision:
            cv2.destroyAllWindows()

        logging.info("Extracted {} count face feature vectors.".format(total))

        return embedding_vectors

    def train_model(self,
                    C=1,
                    kernel="linear",
                    gamma="scale",
                    embedding_vectors=list(),
                    labels=list(),
                    classifier_path=str(),
                    label_encoder_path=str()):
        """
        Train svm classifier by face feature vectors.

        Args
        ----
        embedding_vectors: Extract face feature vector by OpenFace nn4.small2.v1.t7 network model.

        labels: People name.

        classifier_path: Classifier saved path.

        label_encoder_path: Label encoder saved path.
        """

        logging.info("Encoding labels...")
        label_encoder = LabelEncoder()
        encoded_labels = label_encoder.fit_transform(labels)

        logging.info("Training SVM classifier...")
        self.classifier = SVC(C=C, kernel=kernel, gamma=gamma, probability=True)
        self.classifier.fit(embedding_vectors, encoded_labels)

        logging.info("Saving Label encoder...")
        with open(label_encoder_path, "wb") as f:
            f.write(pickle.dumps(label_encoder))
        if os.path.exists(label_encoder_path):
            logging.info("Saved label encoder successfully !")

        logging.info("Saving classifier...")
        with open(classifier_path, "wb") as f:
            f.write(pickle.dumps(self.classifier))
        if os.path.exists(classifier_path):
            logging.info("Saved classifier successfully !")
            logging.info("Saved classifier to {}.".format(os.path.dirname(classifier_path)))

        logging.info("Trained OpenFace classifier finish !")

    def load_model(self,
                   label_encoder_path=str(),
                   classifier_path=str()):
        """
        Load OpenFace label encoder、classifier and embedding network model.

        Args
        ----
        label_encoder_path: Label encoder saved path.

        classifier_path: Classifier saved path.

        Return
        ------
        label_encoder_path: Label encoder saved path.

        classifier_path: Classifier saved path.
        """

        logging.info("Loading label encoder...")
        if os.path.exists(label_encoder_path):
            with open(label_encoder_path, "rb") as lab:
                self.label_encoder = pickle.load(lab)

        logging.info("Loading classifier...")
        if os.path.exists(classifier_path):
            with open(classifier_path, "rb") as classifier:
                self.classifier = pickle.load(classifier)

        logging.info("Loading OpenFace nn4.small2.v1.t7 model...")
        model_name = "nn4.small2.v1.t7"
        model_path = os.path.join(root_dir, model_name)
        if not os.path.exists(model_path):
            download_models(model_name, root_dir)
        self.embedder_model = cv2.dnn.readNetFromTorch(model_path)

        if label_encoder_path == str() or classifier_path == str():
            logging.critical("Please input label encoder and classifier path !")
            raise ValueError

        logging.info("Loaded OpenFace model finish !")

    def predict(self,
                bgr_image=np.ndarray):
        """
        Use OpenFace nn4.small2.v1.t7 network model predict image.

        Args
        ----
        classifier: Classifier saved path.

        embedder_network: Embedder network model saved path.

        bgr_image: BGR format image.

        Return
        ------
        predict_id:

        probability:
        """

        predict_id = None
        probability = None

        state, bgr_image = check_image(bgr_image)
        if state == 0:
            face_blob = cv2.dnn.blobFromImage(bgr_image, 1.0/255, (96, 96), (0, 0, 0), swapRB=True, crop=False)
            self.embedder_model.setInput(face_blob)
            vec = self.embedder_model.forward()

            predicts = self.classifier.predict_proba(vec)[0]
            predict_id = np.argmax(predicts)
            probability = predicts[predict_id]
            logging.debug("Prediction result: {}:{:.6f}".format(predict_id, probability))

        return predict_id, probability