package DeepLearning;

import java.util.Random;

import java.util.List;

import java.util.ArrayList;

public class Dropout {

public int N;

public int n_in;

public int[] hidden_layer_sizes;

public int n_out;

public int n_layers;

public HiddenLayer[] hiddenLayers;

public LogisticRegression logisticLayer;

public Random rng;

public Dropout(int N, int n_in, int[] hidden_layer_sizes, int n_out, Random rng, String activation) {

this.N = N;

this.n_in = n_in;

this.hidden_layer_sizes = hidden_layer_sizes;

this.n_layers = hidden_layer_sizes.length;

this.n_out = n_out;

this.hiddenLayers = new HiddenLayer[n_layers];

if (rng == null) rng = new Random(1234);

this.rng = rng;

if (activation == null) activation = "ReLU";

// construct multi-layer

int input_size;

for(int i=0; i<this.n_layers; i++) {

// layer_size

if(i == 0) {

input_size = n_in;

} else {

input_size = hidden_layer_sizes[i-1];

}

// construct hiddenLayer

this.hiddenLayers[i] = new HiddenLayer(N, input_size, hidden_layer_sizes[i], null, null, rng, activation);

}

// construct logisticLayer

this.logisticLayer = new LogisticRegression(N, hidden_layer_sizes[this.n_layers-1], n_out);

}

public void train(int epochs, double[][] train_X, int[][] train_Y, boolean dropout, double p_dropout, double lr) {

List<int[]> dropout_masks;

List<double[]> layer_inputs;

double[] layer_input;

double[] layer_output = new double[0];

for(int epoch=0; epoch<epochs; epoch++) {

for(int n=0; n<N; n++) {

dropout_masks = new ArrayList<>(n_layers);

layer_inputs = new ArrayList<>(n_layers+1); // +1 for logistic layer

// forward hiddenLayers

for(int i=0; i<n_layers; i++) {

if(i == 0) layer_input = train_X[n];

else layer_input = layer_output.clone();

layer_inputs.add(layer_input.clone());

layer_output = new double[hidden_layer_sizes[i]];

hiddenLayers[i].forward(layer_input, layer_output);

if(dropout) {

int[] mask;

mask = hiddenLayers[i].dropout(layer_output.length, p_dropout, rng);

for(int j=0; j<layer_output.length; j++) layer_output[j] *= mask[j];

dropout_masks.add(mask.clone());

}

}

// forward & backward logisticLayer

double[] logistic_layer_dy; // = new double[n_out];

logistic_layer_dy = logisticLayer.train(layer_output, train_Y[n], lr); //, logistic_layer_dy);

layer_inputs.add(layer_output.clone());

// backward hiddenLayers

double[] prev_dy = logistic_layer_dy;

double[][] prev_W;

double[] dy = new double[0];

for(int i=n_layers-1; i>=0; i--) {

if(i == n_layers-1) {

prev_W = logisticLayer.W;

} else {

prev_dy = dy.clone();

prev_W = hiddenLayers[i+1].W;

}

if(dropout) {

for(int j=0; j<prev_dy.length; j++) {

prev_dy[j] *= dropout_masks.get(i)[j];

}

}

dy = new double[hidden_layer_sizes[i]];

hiddenLayers[i].backward(layer_inputs.get(i), dy, layer_inputs.get(i+1), prev_dy, prev_W, lr);

}

}

}

}

public void pretest(double p_dropout) {

for(int i=0; i<n_layers; i++) {

int in;

int out;

if (i == 0) in = n_in;

else in = hidden_layer_sizes[i];

if (i == n_layers - 1) out = n_out;

else out = hidden_layer_sizes[i+1];

for (int l = 0; l < out; l++) {

for (int m = 0; m < in; m++) {

hiddenLayers[i].W[l][m] *= 1 - p_dropout;

}

}

}

}

public void predict(double[] x, double[] y) {

double[] layer_input;

double[] layer_output = new double[0];

for(int i=0; i<n_layers; i++) {

if(i == 0) layer_input = x;

else layer_input = layer_output.clone();

layer_output = new double[hidden_layer_sizes[i]];

hiddenLayers[i].forward(layer_input, layer_output);

}

logisticLayer.predict(layer_output, y);

}

private static void test_dropout() {

Random rng = new Random(123);

double learning_rate = 0.1;

int n_epochs = 5000;

int train_N = 4;

int test_N = 4;

int n_in = 2;

int[] hidden_layer_sizes = {10, 10};

int n_out = 2;

boolean dropout = true;

double p_dropout = 0.5;

double[][] train_X = {

{0., 0.},

{0., 1.},

{1., 0.},

{1., 1.},

};

int[][] train_Y = {

{0, 1},

{1, 0},

{1, 0},

{0, 1},

};

// construct Dropout

Dropout classifier = new Dropout(train_N, n_in, hidden_layer_sizes, n_out, rng, "ReLU");

// train

classifier.train(n_epochs, train_X, train_Y, dropout, p_dropout, learning_rate);

// pretest

if(dropout) classifier.pretest(p_dropout);

// test data

double[][] test_X = {

{0., 0.},

{0., 1.},

{1., 0.},

{1., 1.},

};

double[][] test_Y = new double[test_N][n_out];

// test

for(int i=0; i<test_N; i++) {

classifier.predict(test_X[i], test_Y[i]);

for(int j=0; j<n_out; j++) {

System.out.print(test_Y[i][j] + " ");

}

System.out.println();

}

}

public static void main(String[] args) {

test_dropout();

}

}