package hex.genmodel.algos.deepwater.caffe;

import java.util.Arrays;

import deepwater.backends.BackendModel;
import deepwater.backends.BackendParams;
import deepwater.backends.BackendTrain;
import deepwater.backends.RuntimeOptions;
import deepwater.datasets.ImageDataSet;

/**
 * This backend forward requests to a docker images running the python
 * Caffe interface. C.f h20-docker/caffe for more information.
 */
public class DeepwaterCaffeBackend implements BackendTrain {
  @Override
  public void delete(BackendModel m) {
    ((DeepwaterCaffeModel) m).close();
  }

  @Override
  public BackendModel buildNet(ImageDataSet dataset, RuntimeOptions opts, BackendParams bparms, int num_classes, String name /*ignored for now*/) {
    if (name.equals("MLP")) {
      // TODO: add non-MLP Models such as lenet, inception_bn, etc.
      int[] hidden = (int[]) bparms.get("hidden");
      int[] sizes = new int[hidden.length + 2];
      sizes[0] = 0;  // ?
      System.arraycopy(hidden, 0, sizes, 1, hidden.length);
      sizes[sizes.length - 1] = num_classes;
      return new DeepwaterCaffeModel(
          (Integer) bparms.get("mini_batch_size"),
          sizes,
          (String[]) bparms.get("activations"),
          (double[]) bparms.get("hidden_dropout_ratios")
      );
    } else throw new UnsupportedOperationException("Only MLP is supported for now for Caffe.");
  }

  // graph (model definition) only
  @Override
  public void saveModel(BackendModel m, String model_path) {
    ((DeepwaterCaffeModel) m).saveModel(model_path);
  }

  // full state of everything but the graph to continue training
  @Override
  public void loadParam(BackendModel m, String param_path) {
    ((DeepwaterCaffeModel) m).loadParam(param_path);
  }

  // full state of everything but the graph to continue training
  @Override
  public void saveParam(BackendModel m, String param_path) {
    ((DeepwaterCaffeModel) m).saveParam(param_path);
  }

  @Override
  public float[] loadMeanImage(BackendModel m, String path) {
    throw new UnsupportedOperationException();
  }

  @Override
  public String toJson(BackendModel m) {
    throw new UnsupportedOperationException();
  }

  @Override
  public void setParameter(BackendModel m, String name, float value) {
    if (name.equals("learning_rate"))
      ((DeepwaterCaffeModel) m).learning_rate(value);
    else if (name.equals("momentum"))
      ((DeepwaterCaffeModel) m).momentum(value);
  }

  // given a mini-batch worth of data and labels, train
  @Override
  public float[]/*ignored*/ train(BackendModel m, float[/*mini_batch * input_neurons*/] data, float[/*mini_batch*/] label) {
    ((DeepwaterCaffeModel) m).train(data, label);
    return null; //return value is always ignored
  }

  @Override
  public float[] predict(BackendModel m, float[] data, float[] label) {
    throw new UnsupportedOperationException();
  }

  // return predictions (num_classes logits (softmax outputs) x mini_batch)
  @Override
  public float[/*mini_batch * num_classes*/] predict(BackendModel m, float[/*mini_batch * input_neurons*/] data) {
    // new float[cm.mini_batch_size * cm.num_classes];
    return ((DeepwaterCaffeModel) m).predict(data);
  }
}

@@ -824,16 +824,43 @@
 n, etc.%0A
+      // Neurons per layer%0A
       in
@@ -901,16 +901,53 @@
 idden%22);
+  // Hidden layers sizes are obtained
 %0A      i
@@ -990,34 +990,59 @@
  2%5D;
-%0A
      
- sizes%5B0%5D = 0;  // ?
+// Sizes array includes input and output sizes
 %0A   
@@ -1098,16 +1098,113 @@
 ength);%0A
+      sizes%5B0%5D = 0;                                 // Input size set to 0! see data information%0A
       si
@@ -1243,180 +1243,561 @@
 ses;
-%0A
+  
       
-return new DeepwaterCaffeModel(%0A          (Integer) bparms.get(%22mini_batch_size%22),%0A          sizes,%0A          (String%5B%5D) bparms.get(%22activ
+// Last size set to num_classes%0A      // Activation functions%0A      String%5B%5D act = new String%5Bhidden.length + 2%5D;%0A      System.arraycopy((String%5B%5D) bparms.get(%22activations%22), 0, act, 1, hidden.length);%0A      act%5B0%5D            = %22input%22;                   // Input type%0A      act%5Bact.length-1%5D = %22output%22;                  // Input type%0A      // Drop out ratios%0A      double%5B%5D dropout = new double%5Bhidden.length + 2%5D;%0A      System.arraycopy((double%5B%5D) bparms.get(%22hidden_dropout_r
 atio
-n
 s%22),
-%0A         
+ 0, dropout, 1, hidden.length);%0A      dropout%5B0%5D =
  (double
 %5B%5D) 
@@ -1788,26 +1788,24 @@
 0%5D = (double
-%5B%5D
 ) bparms.get
@@ -1798,38 +1798,37 @@
 le) bparms.get(%22
-hidden
+input
 _dropout_ratios%22
@@ -1825,19 +1825,259 @@
 ut_ratio
-s
 %22)
+;%0A      dropout%5Bact.length-1%5D = 0.0;                    // Output dropout%0A      return new DeepwaterCaffeModel(%0A          (Integer) bparms.get(%22mini_batch_size%22),    // Batch size information%0A          sizes,%0A          act,%0A          dropout
 %0A      )

converting Deepwater parameters into caffe