#-------------------------------------------------------------
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
#
#-------------------------------------------------------------

/*
 * Hyperparameter Tuning Script For LeNet-like CNN Model
 */
# Imports
source("breastcancer/convnet.dml") as clf

# Read data
# X = read("data/X_0.01_sample_binary")
# Y = read("data/Y_0.01_sample_binary")
# X_val = read("data/X_val_0.01_sample_binary")
# Y_val = read("data/Y_val_0.01_sample_binary")
X = read("data/X_really_small_sample_binary")
Y = read("data/Y_really_small_sample_binary")
X_val = read("data/X_val_really_small_sample_binary")
Y_val = read("data/Y_val_really_small_sample_binary")

# Smaller sample for now
# X = X[1:300000,]
# Y = Y[1:300000,]
# X = X[1:18818,]
# Y = Y[1:18818,]
# X_val = X_val[1:18818,]
# Y_val = Y_val[1:18818,]

# data shape
C = 1
Hin = 64
Win = 64

# Output directory
# dir = "models/systemml/lenet-1-sample/"
dir = "models/systemml/lenet-really-small-sample-hyperparam-2/"

# TODO: Get this working with `parfor`
#j = 1
#while(TRUE) {  # Ideally this would be a parfor loop with a bunch of tries
parfor(j in 1:10000) {
  # Hyperparameter Sampling & Settings
  lr = 10 ^ as.scalar(rand(rows=1, cols=1, min=-7, max=-1))  # learning rate
  mu = as.scalar(rand(rows=1, cols=1, min=0.5, max=0.9))  # momentum
  decay = as.scalar(rand(rows=1, cols=1, min=0.9, max=1))  # learning rate decay constant
  lambda = 10 ^ as.scalar(rand(rows=1, cols=1, min=-7, max=-1))  # regularization constant
  batch_size = 50
  epochs = 1
  log_interval = 10

  # Train
  [Wc1, bc1, Wc2, bc2, Wc3, bc3, Wa1, ba1, Wa2, ba2] =
      clf::train(X, Y, X_val, Y_val, C, Hin, Win, lr, mu, decay, lambda, batch_size, epochs,
                 log_interval, dir)

  # Eval
  #probs = clf::predict(X, C, Hin, Win, Wc1, bc1, Wc2, bc2, Wc3, bc3, Wa1, ba1, Wa2, ba2)
  #[loss, accuracy] = clf::eval(probs, Y)
  probs_val = clf::predict(X_val, C, Hin, Win, Wc1, bc1, Wc2, bc2, Wc3, bc3, Wa1, ba1, Wa2, ba2)
  [loss_val, accuracy_val] = clf::eval(probs_val, Y_val)

  # Save hyperparams with accuracy
  str = "lr: " + lr + ", mu: " + mu + ", decay: " + decay + ", lambda: " + lambda
        + ", batch_size: " + batch_size
  name = dir + accuracy_val + "," + j  #+","+accuracy+","+j
  write(str, name)
  #j = j + 1
}