import tensorflow as tf

import numpy as np

from sklearn.metrics import f1_score

import os

import sys

import math

import plot

from tfrecord_stft_util import input_pipeline_har

os.environ["CUDA_VISIBLE_DEVICES"]="0"

layers = tf.contrib.layers

BATCH_SIZE = 64

GEN_FFT_N = [16, 32, 64, 128]

GEN_FFT_STEP = [FFT_N_ELEM for FFT_N_ELEM in GEN_FFT_N]

FILTER_LEN = [3, 3, 3, 3]

DILATION_LEN = [1, 2, 4, 8]

GEN_FFT_N2 = [12, 24, 48, 96]

SERIES_SIZE2 = 384

GEN_FFT_STEP2 = [FFT_N_ELEM for FFT_N_ELEM in GEN_FFT_N2]

GEN_C_OUT = 64 #72

KEEP_PROB = 0.8

select = 'wifi' # {'hhar', 'wifi'}

if len(sys.argv) > 1:

select = sys.argv[1]

if select != 'wifi' and select != 'hhar':

print 'select wifi or hhar'

sys.exit("select wifi or hhar")

print 'select', select

if select == 'wifi':

SERIES_SIZE = 512

SENSOR_AXIS = 30

SENSOR_NUM = 2

OUT_DIM = 6

if select == 'hhar':

SERIES_SIZE = 512

SENSOR_AXIS = 3

SENSOR_NUM = 2

OUT_DIM = 6

print 'GEN_FFT_N', GEN_FFT_N

print 'GEN_FFT_N2', GEN_FFT_N2

print 'FILTER_LEN', FILTER_LEN

print 'DILATION_LEN', DILATION_LEN

print 'KEEP_PROB', KEEP_PROB

print 'GEN_C_OUT', GEN_C_OUT

FILTER_EXP_SEL = 'linear_interp' #{linear_interp', 'time_zeropadding'}

print 'FILTER_EXP_SEL', FILTER_EXP_SEL

FILTER_INIT = 'real' #{'real', 'complex'}

print 'FILTER_INIT', FILTER_INIT

GLOBAL_KERNEL_SIZE = 32

print 'GLOBAL_KERNEL_SIZE', GLOBAL_KERNEL_SIZE

CONV_KERNEL_INIT = 'freq' #{'time', 'freq'}

print 'CONV_KERNEL_INIT', CONV_KERNEL_INIT

MERGE_INIT = 'zero'

print 'MERGE_INIT', MERGE_INIT

ADAM_LR = 1e-4

ADAM_B1 = 0.9

ADAM_B2 = 0.99

if select == 'wifi':

ACT_DOMIAN = 'freq'

FILTER_FLAG = False

FREQ_CONV_FLAG = True

if select == 'hhar':

ACT_DOMIAN = 'time'

FILTER_FLAG = True

FREQ_CONV_FLAG = False

ADAM_B1 = 0.5

ADAM_B2 = 0.9

DROP_FLAG = True

INPUT_COMPLEX_NORM_FLAG = True

CLIP_FLAG = False

print 'ACT_DOMIAN', ACT_DOMIAN

print 'DROP_FLAG', DROP_FLAG

print 'INPUT_COMPLEX_NORM_FLAG', INPUT_COMPLEX_NORM_FLAG

print 'FILTER_FLAG', FILTER_FLAG

print 'FREQ_CONV_FLAG', FREQ_CONV_FLAG

print 'CLIP_FLAG', CLIP_FLAG

print 'ADAM_LR', ADAM_LR

print 'ADAM_B1', ADAM_B1

print 'ADAM_B2', ADAM_B2

metaDict = {'hhar':[13544, 1765],

'wifi':[11100, 900]}

TRAIN_SIZE = metaDict[select.split('_')[-1]][0]

EVAL_DATA_SIZE = metaDict[select.split('_')[-1]][1]

EVAL_ITER_NUM = int(math.ceil(EVAL_DATA_SIZE / BATCH_SIZE))

TOTAL_ITER_NUM = 10000000

CLIP_VAL = 0.3

if CLIP_FLAG:

print 'CLIP_VAL', CLIP_VAL

def complex_glorot_uniform(c_in, c_out_total, fft_list, fft_n, use_bias=True, name='complex_mat'):

return kernel_complex_dict

def spectral_filter_gen(c_in, c_out_total, basic_len, len_list, use_bias, name='spectral_filter'):

return kernel_dict

def complex_layerNorm(in_r, in_i, reuse=False, name='complex_layerNorm'):

return out_r, out_i

def zero_interp(in_patch, ratio, seg_num, in_fft_n, out_fft_n, f_dim):

return in_patch[:,:,:out_fft_n,:]

def complex_merge(merge_ratio, name='time_merge'):

return kernel_complex, bias_complex

def atten_merge(patch, kernel, bias):

return tf.reduce_sum(patch*patch_atten, 4)

def STFLayer(inputs, fft_list, f_step_list, kenel_len_list, dilation_len_list, c_in, c_out, reuse, out_fft_list=[0], ser_size=SERIES_SIZE, pooling=False, name='STFLayer'):

return patch_time_final

def STFNet(inputs, train, reuse=False, name='STFNet'):

return logits

global_step = tf.Variable(0, trainable=False)

batch_feature, batch_label = input_pipeline_har(os.path.join(select, 'train.tfrecord'), BATCH_SIZE, SERIES_SIZE, SENSOR_AXIS*SENSOR_NUM, OUT_DIM)

batch_eval_feature, batch_eval_label = input_pipeline_har(os.path.join(select, 'eval.tfrecord'), BATCH_SIZE, SERIES_SIZE, SENSOR_AXIS*SENSOR_NUM, OUT_DIM, shuffle_sample=False)

logits = STFNet(batch_feature, True, name='STFNet')

predict = tf.argmax(logits, axis=1)

batchLoss = tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=batch_label)

loss = tf.reduce_mean(batchLoss)

logits_eval = STFNet(batch_eval_feature, False, reuse=True, name='STFNet')

predict_eval = tf.argmax(logits_eval, axis=1)

loss_eval = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=logits_eval, labels=batch_eval_label))

t_vars = tf.trainable_variables()

regularizers = 0.

for var in t_vars:

print var.name

if 'angle' in var.name:

continue

regularizers += tf.nn.l2_loss(var)

loss += 5e-4 * regularizers

if CLIP_FLAG:

discOpt = tf.train.AdamOptimizer(

learning_rate = ADAM_LR,

beta1 = ADAM_B1,

beta2 = ADAM_B2

)

gvs = discOpt.compute_gradients(loss, var_list=t_vars)

capped_gvs = [(tf.clip_by_value(grad, -CLIP_VAL, CLIP_VAL), var) for grad, var in gvs]

discOptimizer = discOpt.apply_gradients(capped_gvs)

else:

discOptimizer = tf.train.AdamOptimizer(

learning_rate = ADAM_LR,

beta1 = ADAM_B1,

beta2 = ADAM_B2

).minimize(loss, var_list=t_vars)

with tf.Session() as sess:

tf.global_variables_initializer().run()

coord = tf.train.Coordinator()

threads = tf.train.start_queue_runners(coord=coord)

for iteration in xrange(TOTAL_ITER_NUM):

_, lossV, _trainY, _predict = sess.run([discOptimizer, loss, batch_label, predict])

_label = np.argmax(_trainY, axis=1)

_accuracy = np.mean(_label == _predict)

plot.plot('train cross entropy', lossV)

plot.plot('train accuracy', _accuracy)

if iteration % 50 == 49:

dev_accuracy = []

dev_cross_entropy = []

total_label = []

total_predt = []

for eval_idx in xrange(EVAL_ITER_NUM):

eval_loss_v, _trainY, _predict = sess.run([loss, batch_eval_label, predict_eval])

_label = np.argmax(_trainY, axis=1)

_accuracy = np.mean(_label == _predict)

total_label += _label.tolist()

total_predt += _predict.tolist()

dev_accuracy.append(_accuracy)

dev_cross_entropy.append(eval_loss_v)

plot.plot('dev accuracy', np.mean(dev_accuracy))

plot.plot('dev cross entropy', np.mean(dev_cross_entropy))

plot.plot('dev macro f1', f1_score(total_label, total_predt, average='macro'))

if (iteration < 5) or (iteration % 50 == 49):

plot.flush()

plot.tick()