package org.algorithm.core.cnn.dataset;

import java.io.BufferedReader;
import java.io.FileReader;
import java.io.IOException;
import java.util.*;

import org.algorithm.core.cnn.entity.Lemma;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;

/**
 * @Author: bijl
 * @Date: 2019年1月21日-下午2:43:44
 * @Description:
 */
public class RelationExtractionDataSet {

    private Map<String, Integer> char2id = new HashMap<>();
    public final int MAX_LEN = 256;


    public RelationExtractionDataSet(String dir) {
        this.loadDictionary(dir);
    }

    /**
     * 加载字典
     *
     * @param dir
     */
    private void loadDictionary(String dir) {
        BufferedReader br = null;
        try {
            br = new BufferedReader(new FileReader(dir));// 读取原始json文件
            String s = null;
            while ((s = br.readLine()) != null) {
                JSONObject jsonObject = (JSONObject) JSON.parse(s);
                Set<Map.Entry<String, Object>> entries = jsonObject.entrySet();
                for (Map.Entry<String, Object> entry : entries)
                    this.char2id.put(entry.getKey(), (Integer) entry.getValue());
            }
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            try {
                br.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    /**
     * 句子转字符ids
     *
     * @param sentence 句子
     * @return ids
     */
    public int[] sentence2ids(String sentence) {
        int[] ids = new int[this.MAX_LEN];
        char ch = '1';
        Integer id = null;
        for (int i = 0; i < sentence.length(); i++) {
            ch = sentence.charAt(i);
            id = this.char2id.get(String.valueOf(ch));
            if (id == null) {
                id = this.char2id.get("<UKC>");
            }
            ids[i] = id.intValue();
        }
        for (int i = sentence.length(); i < this.MAX_LEN; i++)  // padding
            ids[i] = this.char2id.get("<PAD>");

        return ids;
    }

    /**
     * @param sentence 句子
     * @param position 一个实体的位置
     * @return 句子中各个汉子相对于实体的位置
     */
    public int[] getRelativePositions(String sentence, String position) {
        int[] relativePositions = new int[this.MAX_LEN];
        String[] positionPair = position.split(",");
        int startPos = Integer.parseInt(positionPair[0]);
        int endtPos = Integer.parseInt(positionPair[1]);

        for (int i = 0; i < sentence.length(); i++) {
            if (i < startPos)
                relativePositions[i] = startPos - i;
            else if (i >= startPos && i <= endtPos)
                relativePositions[i] = 0;
            else
                relativePositions[i] = i - endtPos;
        }

        for (int i = sentence.length(); i < this.MAX_LEN; i++)
            relativePositions[i] = this.MAX_LEN - 1;

        return relativePositions;
    }

    /**
     * @param sentence 输入句子
     * @param entity1  实体1信息
     * @param entity2  实体2信息
     * @return
     */
    public int[][] getExample(String sentence, Lemma entity1, Lemma entity2) {
        int[][] example = new int[3][this.MAX_LEN];

        example[0] = this.sentence2ids(sentence);
        example[1] = this.getRelativePositions(sentence, entity1.getPosition());

        example[2] = this.getRelativePositions(sentence, entity2.getPosition());

        return example;
    }


}
