# Run this script and training as:
# KERAS_BACKEND=tensorflow ipython --pylab
from pylab import *
from mutils import rms
from time import time
from keras.utils import np_utils
from keras.models import Sequential
from keras.layers.core import Dense, Dropout, Activation
from keras.optimizers import RMSprop
from keras import backend as K

seed(1337)

# generate noisy data
data=[]
for deg in range(360*100):
    ang_rng = (r_[-1:0] + deg)*pi/180.
    Pi = cos(ang_rng) + randn(ang_rng.shape[0])*.01
    Pq = sin(ang_rng) + randn(ang_rng.shape[0])*.01
    if deg % 100 == 0:
        Pi[:] = 0.
        Pq[:] = 0.
    data.append( (Pi, Pq, ((deg % 360.)/360.) ) )
data=array(data)
shuffle(data)
# want X = N x 10*2 and Y = N
row = int(.9*data.shape[0])
X_train = hstack(data[:row,:-1].flat).reshape(row,2*1)
y_train = data[:row,-1]
X_test = hstack(data[row:,:-1].flat).reshape(len(data)-row,2*1)
y_test = data[row:,-1]

print('compiling model')
t1 = time()
model = Sequential()
model.add(Dense(2,input_dim=2))
#model.add(Activation('linear')) # 20=>val_loss=0.0332
#model.add(Activation('relu')) # 20=>val_loss=.0044
#model.add(Activation('softplus')) # 20=>val_loss=0.0132
model.add(Activation('tanh')) # 20=>val_loss=.0017
if False: # Adding a second layer helps performance
    model.add(Dense(10))
    model.add(Activation('tanh'))
model.add(Dense(1))
model.compile(loss='mean_squared_error', optimizer=RMSprop())
print('compile time %.1f' % (time()-t1))

print(model.summary())

print('running model')
t1 = time()
hist = model.fit(X_train, y_train, nb_epoch=2, validation_split=0.1, verbose=1)
score = model.evaluate(X_test, y_test)
print('run time %.1f' % (time()-t1))
print('')
print('loss',score)

print('running predict on test data')
t1 = time()
y_pred_raw = model.predict(X_test)
y_pred = y_pred_raw[:,0]
print('test predict time %.1f' % (time()-t1))
y_arctan = arctan2(X_test[:,1],X_test[:,0])*180/pi/360.
y_arctan[find(y_arctan<0.)] += 1.0
t=r_[0:len(y_pred)]
plot( t, y_test, label='truth' )
plot( t, y_arctan, label='arctan' )
plot( t, y_pred, label='net' )
print('arctan',rms(y_test-y_arctan))
print('net',rms(y_test-y_pred))
legend(loc='best')
show()

x = model.get_input_at(0)
y = model.get_output_at(0)

sess = K.get_session()
#K.set_learning_phase(0)

import tfconvert as tf2e

m = tf2e.Model(sess)
m.build( x, y )
#m.print_data()

import pprint
test_input = X_test[0:2]
result = y_pred_raw[0:2]

m.output_numpy('test2_out_gen')
m.output_numpy_validation('test2_verify.py',test_input,result)

m.output_cpp('test2_out_gen','test')
m.output_cpp_validation('test2_verify.cpp','test',test_input,result)