package org.studiorailgun.sim.eval.parse;

import java.util.HashMap;
import java.util.Map;

import org.tensorflow.Result;
import org.tensorflow.SavedModelBundle;
import org.tensorflow.Tensor;
import org.tensorflow.ndarray.FloatNdArray;
import org.tensorflow.ndarray.NdArray;
import org.tensorflow.ndarray.NdArrays;
import org.tensorflow.ndarray.Shape;
import org.tensorflow.types.TFloat32;
import org.tensorflow.types.TString;

/**
 * Categorizes player commands
 */
public class CommandCategorizationModel {
    
    /**
     * The type of the command
     */
    public static enum CommandType {
        /**
         * Asking the system for help
         */
        HELP,

        /**
         * Looking at something
         */
        LOOK,

        /**
         * Walk towards something
         */
        WALK,

        /**
         * Speak at something
         */
        SPEAK,
    }

    /**
     * The number of different command types
     */
    static final int NUM_COMMAND_TYPES = 4;

    /**
     * Cutoff after which we say the command was not classified successfully
     */
    static final float CLASSIFICAITON_CUTOFF = 0.5f;

    /**
     * The model
     */
    static SavedModelBundle model;

    /**
     * Initializes the categorization model
     */
    public static void init(){
        model = SavedModelBundle.load("./data/model/sim/command_cat");

        //describe the model
        // System.out.println(model.functions());
        // System.out.println(model.functions().get(0).signature());
        // System.out.println(model.functions().get(1).signature());
        // System.out.println(model.functions().get(2).signature());
        // System.out.println(model.functions().get(0).signature().getInputs());
        // System.out.println(model.functions().get(0).signature().getInputs().values().iterator().next().name);
        // System.out.println(model.functions().get(0).signature().getInputs().values().iterator().next().shape);
    }

    /**
     * Categorizes the sentence by function
     * @param input The input string
     * @return The function of the sentence
     */
    public static CommandType categorize(String input){
        //construct input
        TString inputTensor = TString.scalarOf(input);
        inputTensor.shape().append(1);
        NdArray<String> stringArr = NdArrays.ofObjects(String.class, Shape.of(1, 1));
        stringArr.setObject(input);
        inputTensor = TString.tensorOf(stringArr);

        Map<String,Tensor> inputMap = new HashMap<String,Tensor>();
        inputMap.put("keras_tensor",inputTensor);

        //call
        Result result = model.function("serve").call(inputMap);

        //parse results
        TFloat32 resultFloats = (TFloat32)result.get(0);
        FloatNdArray floatArr = resultFloats.get(0);
        float[] classification = new float[NUM_COMMAND_TYPES];
        for(int i = 0; i < NUM_COMMAND_TYPES; i++){
            classification[i] = floatArr.getFloat(i);
        }

        //figure out highest value category
        int maxAt = 0;
        for(int i = 0; i < classification.length; i++){
            if(classification[i] > classification[maxAt]){
                maxAt = i;
            }
        }

        //throw an error if not confident in the classification
        if(classification[maxAt] < CLASSIFICAITON_CUTOFF){
            String message = "Failed to classify command definitively! \n";
            for(int i = 0; i < NUM_COMMAND_TYPES; i++){
                message = message + "class[" + i + "] " + classification[i] + "\n";
            }
            throw new Error(message);
        }

        //convert to enum of command types
        switch(maxAt){
            case 0: {
                return CommandType.HELP;
            }
            case 1: {
                return CommandType.LOOK;
            }
            case 2: {
                return CommandType.WALK;
            }
            case 3: {
                return CommandType.SPEAK;
            }
            default: {
                throw new Error("Failed to categorize!");
            }
        }
    }

}
