package hex.deepwater;

import hex.genmodel.algos.deepwater.caffe.DeepwaterCaffeModel;
import org.junit.Assume;
import org.junit.BeforeClass;
import org.junit.Ignore;
import org.junit.Test;

import java.io.DataInputStream;
import java.io.File;
import java.io.FileInputStream;
import java.util.Random;
import java.util.zip.GZIPInputStream;

@Ignore
public class DeepWaterCaffeIntegrationTest extends DeepWaterAbstractIntegrationTest {

  @Override
  DeepWaterParameters.Backend getBackend() { return DeepWaterParameters.Backend.caffe; }

  @BeforeClass
  public static void checkBackend() { Assume.assumeTrue(DeepWater.haveBackend(DeepWaterParameters.Backend.caffe)); };

  @Test
  public void run() throws Exception {
    /*
      MNIST demo. Get the data first in your home folder:
      cd
      wget http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz
      wget http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz
    */
    final int PIXELS = 28 * 28;
    String home = System.getProperty("user.home");
    DataInputStream pixels = new DataInputStream(new GZIPInputStream(new FileInputStream(new File(
        home + "/train-images-idx3-ubyte.gz"))));
    DataInputStream labels = new DataInputStream(new GZIPInputStream(new FileInputStream(new File(
        home + "/train-labels-idx1-ubyte.gz"))));

    pixels.readInt(); // Magic
    int count = pixels.readInt();
    pixels.readInt(); // Rows
    pixels.readInt(); // Cols

    labels.readInt(); // Magic
    labels.readInt(); // Count

    System.out.println("Read " + count + " samples");
    byte[][] rawI = new byte[count][PIXELS];
    byte[] rawL = new byte[count];
    for (int i = 0; i < count; i++) {
      pixels.readFully(rawI[i]);
      rawL[i] = labels.readByte();
    }

    System.out.println("Randomize");
    Random rand = new Random();
    for (int i = count - 1; i >= 0; i--) {
      int shuffle = rand.nextInt(i + 1);
      byte[] image = rawI[shuffle];
      rawI[shuffle] = rawI[i];
      rawI[i] = image;
      byte label = rawL[shuffle];
      rawL[shuffle] = rawL[i];
      rawL[i] = label;
    }

    System.out.println("Create model");
    final int batch = 256;
    DeepwaterCaffeModel model = new DeepwaterCaffeModel(
        batch,
        new int[] {PIXELS, 4024, 4024, 4048, 10},
        new String[] {"data", "relu", "relu", "relu", "loss"},
        new double[] {.9, .5, .5, .5, 0.},
        1234,
        true // GPU
    );
    model.learning_rate(.01f);
    model.momentum(.9f);

    model.start();
    System.out.println("Press enter");
//    System.in.read();

    System.out.println("Train");

    float[] ps = new float[batch * PIXELS];
    float[] ls = new float[batch];
    for (int iter = 0; iter < 10; iter++) {
      for (int b = 0; b < batch; b++) {
        for (int i = 0; i < PIXELS; i++)
          ps[b * PIXELS + i] = (rawI[b][i] & 0xff) * 0.00390625f;
        ls[b] = rawL[b];
      }
      model.train(ps, ls);
      model.predict(ps);
    }

    model.saveModel("graph");
    model.saveParam("params");
    model.loadParam("params");

//    // Monitor training
//    final long start = System.nanoTime();
//    double time = (System.nanoTime() - start) / 1e9;
//    long processed = 0;
//    int ps = (int) (processed / time);
//    String text = (int) time + "s, " + processed + " (" + (ps) + "/s) ";
  }
}

@@ -337,16 +337,8 @@
 m;%0A%0A
-@Ignore%0A
 publ
@@ -539,16 +539,18 @@
 Class%0A  
+//
 public s
@@ -656,16 +656,82 @@
 fe)); %7D;
+%0A  public static void checkBackend() %7B Assume.assumeTrue(false); %7D
 %0A%0A  @Tes

Disable Caffe tests differently.