package org.studiorailgun.conversation.categorization;

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

import org.studiorailgun.conversation.tracking.Quote;
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 sentences based on function
 */
public class SentenceFunctionCategorizor {

    /**
     * The function of the sentence
     */
    public static enum SentenceFunction {
        /**
         * A sentence whose purpose is to perform some non-informational duty in a conversation.
         * An example of this could be a greeting or farewell.
         */
        UTILITY,

        /**
         * Transfers a piece of information to another participant. For instance, describing
         * the details of an object.
         */
        TRANSFER,
        
        /**
         * Queries a piece of information from another participant.
         */
        QUERY,

        /**
         * Commands a participant to do something.
         */
        IMPERATIVE,
    }

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

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

        //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 quote
     * @return The function of the sentence
     */
    public static void categorize(Quote input){
        //construct input
        TString inputTensor = TString.scalarOf(input.getRaw());
        inputTensor.shape().append(1);
        NdArray<String> stringArr = NdArrays.ofObjects(String.class, Shape.of(1, 1));
        stringArr.setObject(input.getRaw());
        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[]{
            floatArr.getFloat(0),
            floatArr.getFloat(1),
            floatArr.getFloat(2),
            floatArr.getFloat(3),
        };

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

        switch(maxAt){
            case 0: {
                input.setFunction(SentenceFunction.UTILITY);
            } break;
            case 1: {
                input.setFunction(SentenceFunction.TRANSFER);
            } break;
            case 2: {
                input.setFunction(SentenceFunction.QUERY);
            } break;
            case 3: {
                input.setFunction(SentenceFunction.IMPERATIVE);
            } break;
        }
    }
    
}
