#!/usr/bin/python3
# phasenn_test5.py
#
# David Rowe Oct 2019

# Sparse model like test4 but with non-linear layers (activation functions) as
# these will be required later when we add amplitude information.

import numpy as np
import sys
from keras.layers import Dense
from keras import models,layers
from keras import initializers
import keras.backend as K

import matplotlib.pyplot as plt

# constants

N                 = 80      # number of time domain samples in frame
nb_samples        = 100000
nb_batch          = 32
nb_epochs         = 50
width             = 256
pairs             = 2*width
fo_min            = 50
fo_max            = 400
Fs                = 8000

# Generate training data.  Given the phase at the start of the frame,
# and the frequency, determine the phase at the end of the frame

# phase encoded as cos,sin pairs
phase_start = np.zeros((nb_samples, pairs))
phase_end = np.zeros((nb_samples, pairs))
Wo = np.zeros(nb_samples)
L = np.zeros(nb_samples, dtype=int)

# Map 0...width-1 to 0...pi
for i in range(nb_samples):
    r = np.random.rand(1)
    fo = fo_min + (fo_max-fo_min)*r[0]
    Wo[i] = fo*2*np.pi/Fs
    L[i] = int(np.floor(np.pi/Wo[i]))
    for m in range(1, L[i]):
        wm = m*Wo[i]
        bin = int(np.round(wm*width/np.pi))
        # Quantise frequency to this bin, this step is important to
        # match results of test3.  Without it we are adding +/- 0.5 bin
        # uncertainty in the frequency, which will randomise phase shifts
        # across bins
        wm = bin*np.pi/width
        r = np.random.rand(1)
        phase_start_pol = -np.pi + r[0]*2*np.pi
        phase_start[i,2*bin]   = np.cos(phase_start_pol)
        phase_start[i,2*bin+1] = np.sin(phase_start_pol)
        phase_end_pol = phase_start_pol + N*wm
        phase_end[i,2*bin]   = np.cos(phase_end_pol)
        phase_end[i,2*bin+1] = np.sin(phase_end_pol)
        
print(phase_start.shape)                       
print(phase_end.shape)

# note most of these weights whould end up being 0, as we only need
# two wieghts to map each phase rotation
model = models.Sequential()
model.add(layers.Dense(pairs, activation='relu', input_dim=pairs))
model.add(layers.Dense(pairs, input_dim=pairs))
model.summary()

import keras.backend as K

# Compile and fit our model 

from keras import optimizers
sgd = optimizers.SGD(lr=0.02, decay=1e-6, momentum=0.9, nesterov=True)
model.compile(loss='mse', optimizer=sgd)
history = model.fit(phase_start, phase_end, batch_size=nb_batch, epochs=nb_epochs)

# measure error over non-zero samples

phase_end_est = model.predict(phase_start)
ind = np.nonzero(phase_start)
err = (phase_end_est[ind] - phase_end[ind])
var = np.var(err)
std = np.std(err)
print("rect var: %f std: %f" % (var,std))
#print(err[:5,:])

# approximation if error is small
err_angle = np.arctan2(err[1::2], 1)
#print(err[:5,:])
print(err_angle.shape)
var = np.var(err_angle)
std = np.std(err_angle)
print("angle var: %f std: %f" % (var,std))
print("angle var: %4.2f std: %4.2f degs" % (var*180/np.pi,std*180/np.pi))

plot_en = 1;
if plot_en:
    plt.figure(1)
    plt.plot(history.history['loss'])
    plt.title('model loss')
    plt.xlabel('epoch')
 
    plt.figure(2)
    plt.hist(err_angle, bins=20)
    plt.title('phase angle error')

    plt.figure(3)
    r = 0
    phase = np.zeros(width)
    phase_est = np.zeros(width)
    for m in range(1,L[r]):
        wm = m*Wo[r]
        bin = int(np.round(wm*width/np.pi))
        #print("bin: %d %f %f" % (bin, phase_end[r,2*bin], phase_end[r,2*bin+1]))
        phase[m] = np.arctan2(phase_end[r,2*bin+1], phase_end[r,2*bin])
        phase_est[m] = np.arctan2(phase_end_est[r,2*bin+1], phase_end_est[r,2*bin])
    
    plt.plot(phase[1:L[r]+1],'g')
    plt.plot(phase_est[1:L[r]+1],'r')
    
    plt.show()