package org.algorithm.core.cnn.model;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.alibaba.fastjson.TypeReference;
import org.algorithm.core.cnn.AlgorithmCNNExecutorPacs;
import org.algorithm.core.cnn.dataset.RelationExtractionDataSet;
import org.algorithm.core.cnn.entity.LemmaInfo;
import org.algorithm.core.cnn.entity.Triad;
import org.tensorflow.SavedModelBundle;
import org.tensorflow.Session;
import org.tensorflow.Tensor;

import java.nio.IntBuffer;
import java.util.ArrayList;
import java.util.List;

/**
 * @Author: bijl
 * @Date: 2019/1/22 10:21
 * @Decription:
 */
public class RelationExtractionModel extends AlgorithmCNNExecutorPacs {
//    self.X = tf.placeholder(tf.int32, shape=[None, self.max_length], name='X')
//    self.pos1 = tf.placeholder(tf.int32, shape=[None, self.max_length], name='pos1')
//    self.pos2 = tf.placeholder(tf.int32, shape=[None, self.max_length], name='pos2')
//
//    self.y = tf.placeholder(tf.float32, shape=[None, self.n_classes], name='y')
//    self.keep_prob = tf.placeholder(dtype=tf.float32, name='keep_prob')
    private final String X_PLACEHOLDER = "X";
    private final String pos1_PLACEHOLDER = "pos1";
    private final String pos2_PLACEHOLDER = "pos2";
    private final String SOFT_MAX = "softmax/Softmax";
    private final int NUM_LABEL = 2;
    private SavedModelBundle bundle; // 模型捆绑
    private Session session;  // 会话
    protected RelationExtractionDataSet dataSet;

    /**
     * 初始化：加载模型，获取会话。
     * @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();
    }

    @Override
    public List<List<String>> execute(String content, List<Triad> triads) {
//        List<Lemma[]> combinations = new ArrayList<>();
//        // 组合
//        for(int i=0; i < lemmas.size() - 1; i++){  // 两两组合成实体对
//            for (int j = i + 1; j< lemmas.size(); j++){
//                Lemma[] pair = new Lemma[2];
//                pair[0] = lemmas.get(i);
//                pair[1] = lemmas.get(j);
//                combinations.add(pair);
//            }
//        }
//
//        List<Triad> triads = new ArrayList<>();

//        // 遍历组合
//        for (Lemma[] LemmaPair: combinations) {
//            int[][] example = dataSet.getExample(content, LemmaPair[0], LemmaPair[1]);
//            // 调用模型
//            float[][] relation = this.run(example, 1);
//
//            if (relation[0][0] < relation[0][1] ){  // 返回有关系的样本
//                Triad triad = new Triad();
//                triad.setL_1(LemmaPair[0]);
//                triad.setL_2(LemmaPair[1]);
//                triad.setRelation("有");
//                triads.add(triad);
//            }
//
//        }
        return null;
    }

    /**
     * json数组字符串装对象list
     * @param json_content
     * @return
     */
    public List<LemmaInfo> stringToObjects(String json_content) {
        List<LemmaInfo> lemmaInfos = new ArrayList<>();
        JSONArray jsonArray = JSONArray.parseArray(json_content);
        for (int i = 0; i < jsonArray.size(); i++) {
            JSONObject job = jsonArray.getJSONObject(i);
            LemmaInfo info = JSON.parseObject(job.toJSONString(), new TypeReference<LemmaInfo>() {
            });
            lemmaInfos.add(info);
        }
        return lemmaInfos;
    }


    /**
     *
     * @param inputValues 字符id，相对于实体1位置，相对于实体2位置
     * @param batchSize 批量大小
     * @return
     */
    private float[][] run(int[][] inputValues, int batchSize){
        long[] shape = {1, dataSet.MAX_LEN};  // 老模型
        Tensor<Integer> charId = Tensor.create(
                shape,
                IntBuffer.wrap(inputValues[0])
        );
        Tensor<Integer> pos1 = Tensor.create(
                shape,
                IntBuffer.wrap(inputValues[1])
        );
        Tensor<Integer> pos2 = Tensor.create(
                shape,
                IntBuffer.wrap(inputValues[2])
        );

        return this.session.runner()
                .feed(this.X_PLACEHOLDER, charId)
                .feed(this.pos1_PLACEHOLDER, pos1)
                .feed(this.pos2_PLACEHOLDER, pos2)
                .feed("keep_prob", Tensor.create(1.0f, Float.class))  // dropout保留率
                .fetch(this.SOFT_MAX).run().get(0)
                .copyTo(new float[1][this.NUM_LABEL]);
    }

    /**
     * 关闭会话,释放资源
     */
    public void close() {
        this.session.close();
        this.bundle.close();
    }
}
