package org.algorithm.core.neural;

import org.algorithm.core.neural.dataset.NNDataSet;
import org.tensorflow.SavedModelBundle;
import org.tensorflow.Session;
import org.tensorflow.Tensor;

import java.nio.FloatBuffer;
import java.util.Map;

/**
 * tensorflow 模型类，要求单个样本是1维向量，而不是高维向量
 * @Author: bijl
 * @Date: 2018年7月19日-下午7:21:24
 * @Description:
 */
public class TensorflowModel {
    
    private final String INPUT_OPERATION_NAME;   // 输入op的名称
    private final String OUTPUT_OPERATION_NAME;  // 输出op的名称
    private final int NUM_FEATURE;  // 特征个数
    private final int NUM_LABEL;  //  标签（类别）个数
    private SavedModelBundle bundle; // 模型捆绑
    private Session session;  // 会话
    private NNDataSet dataSet;  // 数据集
    
    /**
     * 
     * @param exportDir  模型保存地址
     * @param inputOpName  输入op的名称
     * @param outputOpName  输出op的名称
     * @param dataSet  模型使用的数据集
     */
    public TensorflowModel(String exportDir, String inputOpName, String outputOpName, NNDataSet dataSet) {
        this.INPUT_OPERATION_NAME = inputOpName;
        this.OUTPUT_OPERATION_NAME = outputOpName;
        this.dataSet = dataSet;
        this.NUM_FEATURE = this.dataSet.getNumFeature();
        this.NUM_LABEL = this.dataSet.getNumLabel();
        this.init(exportDir);
                
    }
    
    /**
     * 初始化：加载模型，获取会话。
     * @param exportDir
     */
    public void init(String exportDir) {
        /* load the model Bundle */
        try {
            this.bundle = SavedModelBundle.load(exportDir, "serve");
        } catch (Exception e) {
            e.printStackTrace();
        }

        // create the session from the Bundle
        this.session = bundle.session(); 
    }
    
    /**
     * 运行模型
     * @param inputValue  输入值
     * @param numExamples  样本个数
     * @return 模型的输出
     */
    private float[][] run(float[] inputValues, int numExamples){
        // TODO: 修改的地方
//        long[] inputShape = {numExamples, this.NUM_FEATURE, 4, 1};  // 新模型
        long[] inputShape = {numExamples, this.NUM_FEATURE};  // 老模型
        Tensor<Float> inputTensor = Tensor.create(
                inputShape,  
                FloatBuffer.wrap(inputValues) 
        );
        return this.session.runner().feed(this.INPUT_OPERATION_NAME, inputTensor)
                .feed("keep_prob", Tensor.create(1.0f, Float.class))  // dropout保留率
                .fetch(this.OUTPUT_OPERATION_NAME).run().get(0)
                .copyTo(new float[numExamples][this.NUM_LABEL]);
    }
    
    
    /**
     * 运行模型，并将结果打包成目标格式
     */
    public Map<String, Float> execute(Map<String, Map<String, String>> inputs) {
        float[] inputValues = this.dataSet.toFeatureVector(inputs);
        float sum = 0;
        for (float f : inputValues)
            sum += f;
        if(sum == 0)  // 如果输入没有有效特征，则直接返回null
            return null;
        
        float[][] predict = this.run(inputValues, 1);  // 一次一个样本
        return this.dataSet.wrap(predict);  
    }
    
    
    /**
     * 关闭会话,释放资源
     */
    public void close() {
        this.session.close();
        this.bundle.close();
    }

}
