#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Oct 29 13:11:37 2019

@author: dhzeng
"""

#!/usr/bin/env python
import numpy as np
import random
import matplotlib.patheffects as PathEffects

from keras.layers import Input, Conv2D, Lambda, Dense, Flatten,MaxPooling2D, concatenate, Activation, Dropout, Flatten
from keras.models import Model, Sequential, Model, load_model
from keras.regularizers import l2
from keras import backend as K
from keras.optimizers import SGD,Adam
from keras.losses import binary_crossentropy
import os
import pickle
import matplotlib.pyplot as plt

from itertools import permutations
import seaborn as sns

from keras.datasets import mnist
from sklearn.manifold import TSNE
import tensorflow as tf
from sklearn.svm import SVC
from sklearn.preprocessing import LabelBinarizer

import h5py
import sys
from utils_semihard import metric
## dataset
from keras.datasets import mnist

## for Model definition/training

from keras.utils import plot_model
from keras.callbacks import ModelCheckpoint

## required for semi-hard triplet loss:
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.framework import dtypes

## for visualizing 
import matplotlib.pyplot as plt, numpy as np
from sklearn.decomposition import PCA

path = "/home/dhzeng/AVIDEO/iEmbedding/ccca_vgg_inception10_02/"

def cross_fold_generate():
    files_list = os.listdir(path)
    for i in range(1):# len(files_list)
        new_files_list = files_list[:i]+files_list[i+1:]
        ## read out the file
        ff = h5py.File(path+files_list[i], 'r')
        test_audio, test_rgb = ff["featN1"], ff["featN2"]
        test_lab = ff["testLs"]

        ## train files
        train_audio, train_rgb, train_lab = np.empty((0, 10), np.float32), np.empty((0, 10), np.float32), np.empty((0,), np.float32)
        for ft in new_files_list:
            ftr = h5py.File(path+ft, 'r')
            ft_audio = ftr["featN1"]
            ft_rgb   = ftr["featN2"]
            ft_lab = ftr["testLs"]
            ft_lab = np.asarray([int(labb) for labb in ft_lab])
            train_audio = np.concatenate((train_audio, ft_audio))
            train_rgb   = np.concatenate((train_rgb, ft_rgb))
            train_lab   = np.concatenate((train_lab, ft_lab))
        yield train_audio, train_rgb, train_lab, test_audio, test_rgb, test_lab

def pairwise_distance(feature, squared=False):
    """Computes the pairwise distance matrix with numerical stability.

    output[i, j] = || feature[i, :] - feature[j, :] ||_2

    Args:
      feature: 2-D Tensor of size [number of data, feature dimension].
      squared: Boolean, whether or not to square the pairwise distances.

    Returns:
      pairwise_distances: 2-D Tensor of size [number of data, number of data].
    """
    pairwise_distances_squared = math_ops.add(
        math_ops.reduce_sum(math_ops.square(feature), axis=[1], keepdims=True),
        math_ops.reduce_sum(
            math_ops.square(array_ops.transpose(feature)),
            axis=[0],
            keepdims=True)) - 2.0 * math_ops.matmul(feature,
                                                    array_ops.transpose(feature))

    # Deal with numerical inaccuracies. Set small negatives to zero.
    pairwise_distances_squared = math_ops.maximum(pairwise_distances_squared, 0.0)
    # Get the mask where the zero distances are at.
    error_mask = math_ops.less_equal(pairwise_distances_squared, 0.0)

    # Optionally take the sqrt.
    if squared:
        pairwise_distances = pairwise_distances_squared
    else:
        pairwise_distances = math_ops.sqrt(
            pairwise_distances_squared + math_ops.to_float(error_mask) * 1e-16)

    # Undo conditionally adding 1e-16.
    pairwise_distances = math_ops.multiply(
        pairwise_distances, math_ops.to_float(math_ops.logical_not(error_mask)))

    num_data = array_ops.shape(feature)[0]
    # Explicitly set diagonals to zero.
    mask_offdiagonals = array_ops.ones_like(pairwise_distances) - array_ops.diag(
        array_ops.ones([num_data]))
    pairwise_distances = math_ops.multiply(pairwise_distances, mask_offdiagonals)
    return pairwise_distances

def masked_maximum(data, mask, dim=1):
    """Computes the axis wise maximum over chosen elements.

    Args:
      data: 2-D float `Tensor` of size [n, m].
      mask: 2-D Boolean `Tensor` of size [n, m].
      dim: The dimension over which to compute the maximum.

    Returns:
      masked_maximums: N-D `Tensor`.
        The maximized dimension is of size 1 after the operation.
    """
    axis_minimums = math_ops.reduce_min(data, dim, keepdims=True)
    masked_maximums = math_ops.reduce_max(
        math_ops.multiply(data - axis_minimums, mask), dim,
        keepdims=True) + axis_minimums
    return masked_maximums

def masked_minimum(data, mask, dim=1):
    """Computes the axis wise minimum over chosen elements.

    Args:
      data: 2-D float `Tensor` of size [n, m].
      mask: 2-D Boolean `Tensor` of size [n, m].
      dim: The dimension over which to compute the minimum.

    Returns:
      masked_minimums: N-D `Tensor`.
        The minimized dimension is of size 1 after the operation.
    """
    axis_maximums = math_ops.reduce_max(data, dim, keepdims=True)
    masked_minimums = math_ops.reduce_min(
        math_ops.multiply(data - axis_maximums, mask), dim,
        keepdims=True) + axis_maximums
    return masked_minimums

def triplet_loss_adapted_from_tf(y_true, y_pred):
    del y_true
    margin = 1.
    labels = y_pred[:, :1]

 
    labels = tf.cast(labels, dtype='int32')

    embeddings = y_pred[:, 1:]

    ### Code from Tensorflow function [tf.contrib.losses.metric_learning.triplet_semihard_loss] starts here:
    
    # Reshape [batch_size] label tensor to a [batch_size, 1] label tensor.
    # lshape=array_ops.shape(labels)
    # assert lshape.shape == 1
    # labels = array_ops.reshape(labels, [lshape[0], 1])

    # Build pairwise squared distance matrix.
    pdist_matrix = pairwise_distance(embeddings, squared=True)
    # Build pairwise binary adjacency matrix.
    adjacency = math_ops.equal(labels, array_ops.transpose(labels))
    # Invert so we can select negatives only.
    adjacency_not = math_ops.logical_not(adjacency)

    # global batch_size  
    batch_size = array_ops.size(labels) # was 'array_ops.size(labels)'

    # Compute the mask.
    pdist_matrix_tile = array_ops.tile(pdist_matrix, [batch_size, 1])
    mask = math_ops.logical_and(
        array_ops.tile(adjacency_not, [batch_size, 1]),
        math_ops.greater(
            pdist_matrix_tile, array_ops.reshape(
                array_ops.transpose(pdist_matrix), [-1, 1])))
    mask_final = array_ops.reshape(
        math_ops.greater(
            math_ops.reduce_sum(
                math_ops.cast(mask, dtype=dtypes.float32), 1, keepdims=True),
            0.0), [batch_size, batch_size])
    mask_final = array_ops.transpose(mask_final)

    adjacency_not = math_ops.cast(adjacency_not, dtype=dtypes.float32)
    mask = math_ops.cast(mask, dtype=dtypes.float32)

    # negatives_outside: smallest D_an where D_an > D_ap.
    negatives_outside = array_ops.reshape(
        masked_minimum(pdist_matrix_tile, mask), [batch_size, batch_size])
    negatives_outside = array_ops.transpose(negatives_outside)

    # negatives_inside: largest D_an.
    negatives_inside = array_ops.tile(
        masked_maximum(pdist_matrix, adjacency_not), [1, batch_size])
    semi_hard_negatives = array_ops.where(
        mask_final, negatives_outside, negatives_inside)

    loss_mat = math_ops.add(margin, pdist_matrix - semi_hard_negatives)

    mask_positives = math_ops.cast(
        adjacency, dtype=dtypes.float32) - array_ops.diag(
        array_ops.ones([batch_size]))

    # In lifted-struct, the authors multiply 0.5 for upper triangular
    #   in semihard, they take all positive pairs except the diagonal.
    num_positives = math_ops.reduce_sum(mask_positives)

    semi_hard_triplet_loss_distance = math_ops.truediv(
        math_ops.reduce_sum(
            math_ops.maximum(
                math_ops.multiply(loss_mat, mask_positives), 0.0)),
        num_positives,
        name='triplet_semihard_loss')
    
    ### Code from Tensorflow function semi-hard triplet loss ENDS here.
    return semi_hard_triplet_loss_distance

def create_base_networks(in_dims):
    """
    Base network to be shared.
    """
    model = Sequential()
    model.add(Dense(300, input_dim=in_dims))
    model.add(Activation("tanh"))
    model.add(Dropout(0.1))
    model.add(Dense(300, kernel_initializer='normal', activation='tanh'))
    model.add(Dropout(0.1))
    model.add(Dense(300, kernel_initializer='normal', activation='tanh'))
    model.add(Dropout(0.1))
    model.add(Dense(10, kernel_initializer='normal', activation='tanh'))
    model.add(Dropout(0.1))
    model.add(Activation("sigmoid"))
    # model.add(Dense(600))

    return model

def create_base_network(image_input_shape, embedding_size):
    """
    Base network to be shared (eq. to feature extraction).
    """
    input_image = Input(shape=image_input_shape)
    x = input_image 
    #x = Flatten()(input_image)
    x = Dense(128, activation='relu')(x)
    x = Dropout(0.1)(x)
    x = Dense(128, activation='relu')(x)
    x = Dropout(0.1)(x)
    x = Dense(embedding_size)(x)

    base_network = Model(inputs=input_image, outputs=x)
    #plot_model(base_network, to_file='base_network.png', show_shapes=True, show_layer_names=True)
    return base_network

embedding_size = 10
batch_size = 256
nepochs = 1

nepochs = 1
input_image_shape1=(10,)
input_image_shape2=(10,)
base_network1 = create_base_network(input_image_shape1, embedding_size)
base_network2 = create_base_network(input_image_shape2, embedding_size)
input_images1 = Input(shape=input_image_shape1, name='input_image1')
input_images2 = Input(shape=input_image_shape2, name='input_image2')
input_labels  = Input(shape=(1,), name='input_label') 
embeddings1   = base_network1([input_images1])
embeddings2   = base_network2([input_images2])
labels_plus_embeddings = concatenate([input_labels, embeddings1, embeddings2])
model = Model(inputs=[input_images1,input_images2, input_labels],
                  outputs=labels_plus_embeddings)
model.summary()
opt = Adam(lr=0.0001)  # choose optimiser. RMS is good too!
model.compile(loss=triplet_loss_adapted_from_tf,optimizer=opt)

for data_xy in cross_fold_generate():
    trainAudio, trainRGB, catTrain, testAudio, testRGB, catTest = data_xy
    dummy_gt_train = np.zeros((len(trainAudio), embedding_size + 1))
    dummy_gt_val = np.zeros((len(trainAudio), embedding_size + 1))

    trainAudio = np.reshape(trainAudio, (len(trainAudio),  trainAudio.shape[1]))
    trainRGB = np.reshape(trainRGB, (len(trainRGB),trainRGB.shape[1]))
    x_val1 =  trainAudio
    x_val2 = trainRGB
    print(x_val1.shape, x_val1.shape, catTrain.shape)
    H = model.fit(
            x=[trainAudio, trainRGB, catTrain],
            y=dummy_gt_train,
            batch_size=batch_size,
            epochs=nepochs,
            validation_data=([x_val1, x_val2, catTrain], dummy_gt_val))
    
    # Test the network
    from scipy.spatial.distance import cdist
    # creating an empty network
    testing_embeddings1 = create_base_network(input_image_shape1,
                                             embedding_size=embedding_size)
    testing_embeddings2 = create_base_network(input_image_shape2,
                                             embedding_size=embedding_size)
    x_embeddings_before_train1 = testing_embeddings1.predict(testAudio)
    x_embeddings_before_train2 = testing_embeddings2.predict(testRGB)
    views = x_embeddings_before_train1, x_embeddings_before_train2
    mAp_view1, mAp_view2,pre_final1,rec_final1, pre_final2, rec_final2 = metric(views, catTest)
    

