source("nnUtils.R")

f <- function (x)  0.2 + 0.05 * x + 0.4 * sin(x) + 0.05 * rnorm(length(x))
N <- 40
xmax <- 20
Xtrain <- matrix(seq(0,xmax,length=N),N,1)
Ttrain <- f(Xtrain)
Xtest <- Xtrain
Ttest <- f(Xtest)

nh <- 5
rhoh <- 0.1
rhoo <- 0.1

standardizeF <- makeStandardizeF(Xtrain)
XtrainS <- standardizeF(Xtrain)
XtestS <- standardizeF(Xtest)


X <- XtrainS
T <- Ttrain

ni <- ncol(X)
no <- ncol(T)
V <- matrix(0.1*(runif((ni+1)*nh)-0.5), ni+1,nh)
W <- matrix(0.1*(runif((nh+1)*no)-0.5), nh+1,no)


rh <- rhoh / (nrow(X)*ncol(T))
ro <- rhoo / (nrow(X)*ncol(T))
X1 <- cbind(1,X)
errorTrace <- NULL

x11(type="Xlib")
par(mfcol=c(2,2),mar=c(4,4,0,0),bty="n")

nReps <- 200000
for (reps in 1:nReps) {
  Z <- tanh(X1 %*% V)
  Y <- cbind(1,Z) %*% W
  error <- Y - T
  errorTrace <- c(errorTrace, sqrt(mean(error^2)))
  V <- V - ro * t(X1) %*% (error %*% t(W[-1,]) * (1-Z^2))
  W <- W - rh * t(cbind(1,Z)) %*% error

  if (reps %% (nReps/100) == 0) {
    plot(errorTrace,type="l",lwd=2,xlab="Epochs",ylab="Train RMSE")
    drawNNet(list(W=list(V,W)))
    matplot(Xtrain,Z,type="b",lwd=2,lty=1,xlab="x",ylab="Z")
    matplot(Xtrain,cbind(T,Y),lty=1,type="l",lwd=2,xlab="x",ylab="T and Y")
    legend("topleft",c("Target","Y"),lty=1,lwd=2,col=1:2)
  }
}