package info.magnolia.image.tagging;

import java.io.IOException;
import java.net.URL;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;

import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.factory.Nd4j;
import org.nd4j.shade.jackson.databind.ObjectMapper;

import com.google.common.collect.Lists;

/**
 * Helper class with a static method that returns the label description.
 *
 * @author susaneraly
 */
public class ImageNetLabels {

    private final static String jsonUrl = "http://blob.deeplearning4j.org/utils/imagenet_class_index.json";
    private static ArrayList<String> predictionLabels = null;

    public ImageNetLabels() {
        this.predictionLabels = getLabels();
    }

    private static ArrayList<String> getLabels() {
        if (predictionLabels == null) {
            HashMap<String, ArrayList<String>> jsonMap;
            try {
                jsonMap = new ObjectMapper().readValue(new URL(jsonUrl), HashMap.class);
                predictionLabels = new ArrayList<>(jsonMap.size());
                for (int i = 0; i < jsonMap.size(); i++) {
                    predictionLabels.add(jsonMap.get(String.valueOf(i)).get(1));
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        return predictionLabels;
    }

    /**
     * Returns the description of tne nth class in the 1000 classes of ImageNet.
     * @param n
     * @return
     */
    public String getLabel(int n) {
        return predictionLabels.get(n);
    }

    /**
     * Given predictions from the trained model this method will return a string
     * listing the top five matches and the respective probabilities
     * @param predictions
     * @return
     */
    public List<Result> decodePredictions(INDArray predictions) {
        List<Result> predictionDescription = Lists.newArrayList();
        int[] top5 = new int[5];
        float[] top5Prob = new float[5];

        //brute force collect top 5
        int i = 0;
        for (int batch = 0; batch < predictions.size(0); batch++) {
            INDArray currentBatch = predictions.getRow(batch).dup();
            while (i < 5) {
                top5[i] = Nd4j.argMax(currentBatch, 1).getInt(0, 0);
                top5Prob[i] = currentBatch.getFloat(batch, top5[i]);
                currentBatch.putScalar(0, top5[i], 0);
                predictionDescription.add(new Result(top5Prob[i] * 100, predictionLabels.get(top5[i])));
                i++;
            }
        }
        return predictionDescription;
    }

    public class Result {
        private final float probability;
        private final String prediction;

        public Result(float probability, String prediction) {
            this.prediction = prediction;
            this.probability = probability;
        }

        public String getPrediction() {
            return prediction;
        }

        public float getProbability() {
            return probability;
        }
    }

}


@@ -1,8 +1,1383 @@
+/**%0A * This file Copyright (c) 2011-2017 Magnolia International%0A * Ltd.  (http://www.magnolia-cms.com). All rights reserved.%0A *%0A *%0A * This file is dual-licensed under both the Magnolia%0A * Network Agreement and the GNU General Public License.%0A * You may elect to use one or the other of these licenses.%0A *%0A * This file is distributed in the hope that it will be%0A * useful, but AS-IS and WITHOUT ANY WARRANTY; without even the%0A * implied warranty of MERCHANTABILITY or FITNESS FOR A%0A * PARTICULAR PURPOSE, TITLE, or NONINFRINGEMENT.%0A * Redistribution, except as permitted by whichever of the GPL%0A * or MNA you select, is prohibited.%0A *%0A * 1. For the GPL license (GPL), you can redistribute and/or%0A * modify this file under the terms of the GNU General%0A * Public License, Version 3, as published by the Free Software%0A * Foundation.  You should have received a copy of the GNU%0A * General Public License, Version 3 along with this program;%0A * if not, write to the Free Software Foundation, Inc., 51%0A * Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.%0A *%0A * 2. For the Magnolia Network Agreement (MNA), this file%0A * and the accompanying materials are made available under the%0A * terms of the MNA which accompanies this distribution, and%0A * is available at http://www.magnolia-cms.com/mna.html%0A *%0A * Any modifications to this file must keep this entire header%0A * intact.%0A *%0A */%0A
 package 
@@ -3087,16 +3087,24 @@
 bilities
+.%0A     *
 %0A     * 
@@ -3957,24 +3957,62 @@
 ion;%0A    %7D%0A%0A
+    /**%0A     * TODO: JavaDoc.%0A     */%0A
     public c

Fix JavaDoc and license to avoid checkstyle errors