# -*- coding: utf-8 -*-
"""
Created on Mon May 01 17:26:45 2017

@author: medialab
"""
import h5py
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from scipy.special import expit
import tensorflow as tf


def conv2d(x, W):
    return tf.nn.conv2d(x, W, strides=[1, 1, 1, 1], padding = 'SAME')

def deconv2d(x, W, output_shape):
    return tf.nn.conv2d_transpose(x, W, output_shape, strides = [1, 1, 1, 1], padding = 'SAME')
idx = 99
start = 200
#KandM.h5 is generated by concatenate.py
src_path = './Concatenate_Data/CNN/'
date_ext = '_CNN_0504'
data_ext = '_K2M_rel'

NK = h5py.File(src_path + 'limb_KandM_ex4'+date_ext+'.h5','r')['N_Kinect'][:]
K  = h5py.File(src_path + 'limb_KandM_ex4'+date_ext+'.h5','r')['Kinect'][:]
M  = h5py.File(src_path + 'limb_KandM_ex4'+date_ext+'.h5','r')['Mcam'][:]   
[MIN,MAX]  = h5py.File(src_path + 'NLdata'+date_ext+'.h5','r')['minmax'][:]  
  
    
We1   = h5py.File('./data/CNN/We1'+date_ext+data_ext+'.h5','r')['w_e_conv1_'+str(idx)][:]
We2   = h5py.File('./data/CNN/We2'+date_ext+data_ext+'.h5','r')['w_e_conv2_'+str(idx)][:]
Wd1   = h5py.File('./data/CNN/Wd1'+date_ext+data_ext+'.h5','r')['w_d_conv1_'+str(idx)][:]
Wd2   = h5py.File('./data/CNN/Wd2'+date_ext+data_ext+'.h5','r')['w_d_conv2_'+str(idx)][:]
be1   = h5py.File('./data/CNN/be1'+date_ext+data_ext+'.h5','r')['b_e_conv1_'+str(idx)][:]
be2   = h5py.File('./data/CNN/be2'+date_ext+data_ext+'.h5','r')['b_e_conv2_'+str(idx)][:]
bd1   = h5py.File('./data/CNN/bd1'+date_ext+data_ext+'.h5','r')['b_d_conv1_'+str(idx)][:]
bd2   = h5py.File('./data/CNN/bd2'+date_ext+data_ext+'.h5','r')['b_d_conv2_'+str(idx)][:]
[joints_num,group_size,batch_size,conv_ker_L1,conv_ker_L2] = h5py.File('./data/CNN/model'+date_ext+data_ext+'.h5','r')['parm'][:]

x = tf.placeholder(tf.float32, shape = [None, group_size,joints_num*3])
x_origin = tf.reshape(x, [-1, group_size,joints_num*3, 1])
he1 = tf.nn.relu(tf.add(conv2d(x_origin, We1), be1))
he2 = tf.nn.relu(tf.add(conv2d(he1, We2), be2))

output_shape_d1 = tf.pack([batch_size, group_size, joints_num*3, conv_ker_L1])
output_shape_d2 = tf.pack([batch_size, group_size, joints_num*3, 1])
hd1 = tf.nn.relu(deconv2d(he2, Wd1,output_shape_d1)+bd1)
hd2 = tf.nn.relu(deconv2d(hd1, Wd2,output_shape_d2)+bd2)

sess = tf.InteractiveSession()
#        batch_size = 16
loop = 10


tmp = np.zeros([joints_num*3,loop*batch_size])



for i in range(loop):
    batch_raw = np.zeros([batch_size,group_size,joints_num*3])
    batch_raw[:] = NK[:,:,start+i*batch_size:start+(i+1)*batch_size].T        
           

    a = sess.run(hd2,feed_dict={x:batch_raw}).T[0,:,:,:] 
              
#    if i == 0:
#        tmp[:,:group_size] = a[:,:,0]
#        for j in range(1,batch_size):
#            tmp[:,(group_size-1)+j] = a[:,:,j][:,-1]
#    else:
#    for j in range(batch_size):
#        tmp[:,i*batch_size+j] = a[:,:,j][:,-1]

    tmp[:,i*batch_size:(i+1)*batch_size] = a[:,-1,:]


KJ_mod = tmp*(MAX-MIN)+MIN



fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.set_xlabel('Z axis')
ax.set_ylabel('X axis')
ax.set_zlabel('Y axis')    

for i in range(loop*batch_size):
    plt.cla()
    kxm = KJ_mod.T[i,::3]
    kym = KJ_mod.T[i,1::3]
    kzm = KJ_mod.T[i,2::3]


    mx = M.T[start+group_size-1+i,-1,::3]
    my = M.T[start+group_size-1+i,-1,1::3]
    mz = M.T[start+group_size-1+i,-1,2::3]    
#
    kx = K.T[start+group_size-1+i,-1,::3]
    ky = K.T[start+group_size-1+i,-1,1::3]
    kz = K.T[start+group_size-1+i,-1,2::3]  

#    if taro:
#        kx_taro = K_taro.T[start+i,::3]
#        ky_taro = K_taro.T[start+i,1::3]
#        kz_taro = K_taro.T[start+i,2::3]
#
#        mx_taro = M_taro.T[start+i,::3]
#        my_taro = M_taro.T[start+i,1::3]
#        mz_taro = M_taro.T[start+i,2::3]
#            
#        ax.scatter(kz_taro, kx_taro, ky_taro, c = 'blue', s = 100)    
#        ax.scatter(mz_taro, mx_taro, my_taro,c = 'green',s = 50,alpha=.4)
#        ax.scatter(kz_taro, kx_taro, ky_taro,c = 'red',s = 50,alpha=.4)
#        
    
    ax.scatter(kz, kx, ky, c = 'blue', s = 100,label='Kinect Joints')    
    ax.scatter(mz, mx, my,c = 'green',s = 50,alpha=.4,label='MoCam Joints')
#    ax.scatter(kzm, kxm, kym,c = 'red',s = 50,alpha=.4,label='modified Kjoints')
    ax.set_xlim(-300,300)
    ax.set_ylim(-200,400)
    ax.set_zlim(50,600)
    ax.set_title('frame no : '+str(i+start)+'  training '+ str(idx*1000) +' times ' )
    ax.set_xlabel('Z axis')
    ax.set_ylabel('X axis')
    ax.set_zlabel('Y axis')
    plt.legend( loc=1)
    plt.draw()
    plt.pause(1.0/60)
    
