/*
 * Copyright (c) 2019, Oracle and/or its affiliates. All rights reserved.
 * Licensed under the MIT License.
 */
package sample;

import ai.onnxruntime.NodeInfo;
import ai.onnxruntime.OnnxTensor;
import ai.onnxruntime.OrtEnvironment;
import ai.onnxruntime.OrtException;
import ai.onnxruntime.OrtSession;
import ai.onnxruntime.OrtSession.Result;
import ai.onnxruntime.OrtSession.SessionOptions;
import ai.onnxruntime.OrtSession.SessionOptions.OptLevel;
import java.io.BufferedReader;
import java.io.FileReader;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.logging.Level;
import java.util.logging.Logger;
import java.util.regex.Pattern;

/**
 * Demo code, supporting both a pytorch CNN trained on MNIST and a scikit-learn model trained on
 * MNIST.
 */
public class ScoreMNIST {

  private static final Logger logger = Logger.getLogger(ScoreMNIST.class.getName());
  /** Pattern for splitting libsvm format files. */
  private static final Pattern splitPattern = Pattern.compile("\\s+");

  /** A named tuple for sparse classification data. */
  private static class SparseData {
    public final int[] labels;
    public final List<int[]> indices;
    public final List<float[]> values;

    public SparseData(int[] labels, List<int[]> indices, List<float[]> values) {
      this.labels = labels;
      this.indices = Collections.unmodifiableList(indices);
      this.values = Collections.unmodifiableList(values);
    }
  }

  /**
   * Converts a List of Integer into an int array.
   *
   * @param list The list to convert.
   * @return The int array.
   */
  private static int[] convertInts(List<Integer> list) {
    int[] output = new int[list.size()];
    for (int i = 0; i < list.size(); i++) {
      output[i] = list.get(i);
    }
    return output;
  }

  /**
   * Converts a List of Float into a float array.
   *
   * @param list The list to convert.
   * @return The float array.
   */
  private static float[] convertFloats(List<Float> list) {
    float[] output = new float[list.size()];
    for (int i = 0; i < list.size(); i++) {
      output[i] = list.get(i);
    }
    return output;
  }

  /**
   * Loads data from a libsvm format file.
   *
   * @param path The path to load the data from.
   * @return A named tuple containing the data.
   * @throws IOException If it failed to read the file.
   */
  private static SparseData load(String path) throws IOException {
    int pos = 0;
    List<int[]> indices = new ArrayList<>();
    List<float[]> values = new ArrayList<>();
    List<Integer> labels = new ArrayList<>();
    String line;
    int maxFeatureID = Integer.MIN_VALUE;
    try (BufferedReader reader = new BufferedReader(new FileReader(path))) {
      for (; ; ) {
        line = reader.readLine();
        if (line == null) {
          break;
        }
        pos++;
        String[] fields = splitPattern.split(line);
        int lastID = -1;
        try {
          boolean valid = true;
          List<Integer> curIndices = new ArrayList<>();
          List<Float> curValues = new ArrayList<>();
          for (int i = 1; i < fields.length && valid; i++) {
            int ind = fields[i].indexOf(':');
            if (ind < 0) {
              logger.warning(String.format("Weird line at %d", pos));
              valid = false;
            }
            String ids = fields[i].substring(0, ind);
            int id = Integer.parseInt(ids);
            curIndices.add(id);
            if (maxFeatureID < id) {
              maxFeatureID = id;
            }
            float val = Float.parseFloat(fields[i].substring(ind + 1));
            curValues.add(val);
            if (id <= lastID) {
              logger.warning(String.format("Repeated features at line %d", pos));
              valid = false;
            } else {
              lastID = id;
            }
          }
          if (valid) {
            // Store the label
            labels.add(Integer.parseInt(fields[0]));
            // Store the features
            indices.add(convertInts(curIndices));
            values.add(convertFloats(curValues));
          } else {
            throw new IOException("Invalid LibSVM format file at line " + pos);
          }
        } catch (NumberFormatException ex) {
          logger.warning(String.format("Weird line at %d", pos));
          throw new IOException("Invalid LibSVM format file", ex);
        }
      }
    }

    logger.info(
        "Loaded "
            + maxFeatureID
            + " features, "
            + labels.size()
            + " samples, from + '"
            + path
            + "'.");
    return new SparseData(convertInts(labels), indices, values);
  }

  /**
   * Naively takes the softmax of the input.
   *
   * @param input The input array.
   * @return The softmax of the input.
   */
  public static float[] softmax(float[] input) {
    double[] tmp = new double[input.length];
    double sum = 0.0;
    for (int i = 0; i < input.length; i++) {
      double val = Math.exp(input[i]);
      sum += val;
      tmp[i] = val;
    }

    float[] output = new float[input.length];
    for (int i = 0; i < output.length; i++) {
      output[i] = (float) (tmp[i] / sum);
    }

    return output;
  }

  /**
   * Zeros the supplied array.
   *
   * @param data The array to zero.
   */
  public static void zeroData(float[][][][] data) {
    // Zero the array
    for (int i = 0; i < data.length; i++) {
      for (int j = 0; j < data[i].length; j++) {
        for (int k = 0; k < data[i][j].length; k++) {
          Arrays.fill(data[i][j][k], 0.0f);
        }
      }
    }
  }

  /**
   * Writes out sparse data into the last two dimensions of the supplied 4d array.
   *
   * @param data The 4d array to write to.
   * @param indices The indices of the sparse data.
   * @param values The values of the sparse data.
   */
  public static void writeData(float[][][][] data, int[] indices, float[] values) {
    zeroData(data);

    for (int m = 0; m < indices.length; m++) {
      int i = (indices[m]) / 28;
      int j = (indices[m]) % 28;
      data[0][0][i][j] = values[m] / 255;
    }

    for (int i = 0; i < 28; i++) {
      for (int j = 0; j < 28; j++) {
        data[0][0][i][j] = (data[0][0][i][j] - 0.1307f) / 0.3081f;
      }
    }
  }

  /**
   * Zeros the array used by the scikit-learn model.
   *
   * @param data The array to zero.
   */
  public static void zeroDataSKL(float[][] data) {
    // Zero the array
    for (int i = 0; i < data.length; i++) {
      Arrays.fill(data[i], 0.0f);
    }
  }

  /**
   * Writes out sparse data to the last dimension of the supplied 2d array.
   *
   * @param data The 2d array to write to.
   * @param indices The indices of the sparse data.
   * @param values THe values of the sparse data.
   */
  public static void writeDataSKL(float[][] data, int[] indices, float[] values) {
    zeroDataSKL(data);

    for (int m = 0; m < indices.length; m++) {
      data[0][indices[m]] = values[m];
    }
  }

  /**
   * Find the maximum probability and return it's index.
   *
   * @param probabilities The probabilites.
   * @return The index of the max.
   */
  public static int pred(float[] probabilities) {
    float maxVal = Float.NEGATIVE_INFINITY;
    int idx = 0;
    for (int i = 0; i < probabilities.length; i++) {
      if (probabilities[i] > maxVal) {
        maxVal = probabilities[i];
        idx = i;
      }
    }
    return idx;
  }

  public static void main(String[] args) throws OrtException, IOException {
    if (args.length < 2 || args.length > 3) {
      System.out.println("Usage: ScoreMNIST <model-path> <test-data> <optional:scikit-learn-flag>");
      System.out.println("The test data input should be a libsvm format version of MNIST.");
      return;
    }

    try (OrtEnvironment env = OrtEnvironment.getEnvironment();
        OrtSession.SessionOptions opts = new SessionOptions()) {

      opts.setOptimizationLevel(OptLevel.BASIC_OPT);

      logger.info("Loading model from " + args[0]);
      try (OrtSession session = env.createSession(args[0], opts)) {

        logger.info("Inputs:");
        for (NodeInfo i : session.getInputInfo().values()) {
          logger.info(i.toString());
        }

        logger.info("Outputs:");
        for (NodeInfo i : session.getOutputInfo().values()) {
          logger.info(i.toString());
        }

        SparseData data = load(args[1]);

        float[][][][] testData = new float[1][1][28][28];
        float[][] testDataSKL = new float[1][780];

        int correctCount = 0;
        int[][] confusionMatrix = new int[10][10];

        String inputName = session.getInputNames().iterator().next();

        for (int i = 0; i < data.labels.length; i++) {
          if (args.length == 3) {
            writeDataSKL(testDataSKL, data.indices.get(i), data.values.get(i));
          } else {
            writeData(testData, data.indices.get(i), data.values.get(i));
          }

          try (OnnxTensor test =
                  OnnxTensor.createTensor(env, args.length == 3 ? testDataSKL : testData);
              Result output = session.run(Collections.singletonMap(inputName, test))) {

            int predLabel;

            if (args.length == 3) {
              long[] labels = (long[]) output.get(0).getValue();
              predLabel = (int) labels[0];
            } else {
              float[][] outputProbs = (float[][]) output.get(0).getValue();
              predLabel = pred(outputProbs[0]);
            }
            if (predLabel == data.labels[i]) {
              correctCount++;
            }

            confusionMatrix[data.labels[i]][predLabel]++;

            if (i % 2000 == 0) {
              logger.log(Level.INFO, "Cur accuracy = " + ((float) correctCount) / (i + 1));
              logger.log(Level.INFO, "Output type = " + output.get(0).toString());
              if (args.length == 3) {
                logger.log(Level.INFO, "Output type = " + output.get(1).toString());
                logger.log(Level.INFO, "Output value = " + output.get(1).getValue().toString());
              }
            }
          }
        }

        logger.info("Final accuracy = " + ((float) correctCount) / data.labels.length);

        StringBuilder sb = new StringBuilder();
        sb.append("Label");
        for (int i = 0; i < confusionMatrix.length; i++) {
          sb.append(String.format("%1$5s", "" + i));
        }
        sb.append("\n");

        for (int i = 0; i < confusionMatrix.length; i++) {
          sb.append(String.format("%1$5s", "" + i));
          for (int j = 0; j < confusionMatrix[i].length; j++) {
            sb.append(String.format("%1$5s", "" + confusionMatrix[i][j]));
          }
          sb.append("\n");
        }

        System.out.println(sb.toString());
      }
    }

    logger.info("Done!");
  }
}