package spExtor;

/**------------------------------------------------------------------------
 *
 * Prepares the testing data to be used in the full process of the incremental learning procedure
 *
 * @author Hussein S. Al-Olimat github.com/halolimat
 *
 *------------------------------------------------------------------------
 *
 * This file is part of the Sparse Entity Extraction tool (SpExtor) which is the implementation of the proposed method
 * in our COLING 2018 publication titled: "A practical incremental learning framework for sparse entity extraction".
 * The full COLING 2018 paper can be found @ https://aclanthology.info/papers/C18-1059/c18-1059
 *
 * SpExtor is a free tool: you can redistribute it and/or modify it under the terms of the GNU General Public License as
 * published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version.
 *
 * SpExtor is distributed WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A
 * PARTICULAR PURPOSE.  See the GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License along with SpExtor.
 * If not, see <http://www.gnu.org/licenses/> and https://github.com/halolimat/SpExtor/blob/master/LICENSE.
 *
 *------------------------------------------------------------------------
 */

import java.io.BufferedReader;
import java.io.FileReader;
import java.util.*;
import java.util.stream.Collectors;

public class Dataset {

    private HashMap<Integer, String> conll_train_sentences;
    private HashMap<Integer, String> conll_test_sentences;

    private HashMap<Integer, String> raw_train_sentences;
    private HashMap<Integer, String> raw_test_sentences;

    private Set<String> train_entities;
    private HashMap<Integer, Set<String>> train_sentence_entities;

    // *****************************************************************************************************************
    // Private methods
    // *****************************************************************************************************************

    /**
     * Extract annotated entities from the CoNLL formatted annotated file.
     *
     * @param data_map: a CSV tuple of the semantics of each column in CoNLL formatted data.
     */
    private void extract_entities_from_train_conll(String data_map) {
        // Note: this is not correctly getting all sentences with entities from CoNLL-2003-LOC, getting only 977 from 1k

        HashMap<Integer, String> conll_train_sentences = getConll_train_sentences();
        HashMap<Integer, Set<String>> train_sentence_entities = new HashMap<>();
        Set<String> train_entities = new HashSet<>();

        for (int sentence_id : conll_train_sentences.keySet()) {
            String conll_sentence = conll_train_sentences.get(sentence_id);
            StringJoiner entity = new StringJoiner(" ");
            for (String line : conll_sentence.split("\n")) {
                String word = line.split(" ")[0];
                int eclass_index = Integer.parseInt(data_map.substring(data_map.length() - 1));
                String eclass = "";
                try{
                    eclass = line.split(" ")[eclass_index];
                } catch (Exception e){
                    System.out.println("####");
                    System.out.println(conll_sentence);
                    System.out.println(sentence_id);
                    System.out.println("####");
                    e.printStackTrace();
                }

                if (!eclass.equals("O")) {
                    entity.add(word);
                } else if (!entity.toString().equals("")) {
                    train_entities.add(entity.toString());
                    if (!train_sentence_entities.keySet().contains(sentence_id)) {
                        train_sentence_entities.put(sentence_id, new HashSet<>());
                    }
                    train_sentence_entities.get(sentence_id).add(entity.toString());
                    entity = new StringJoiner(" ");
                }
            }
        }

        setTrain_sentence_entities(train_sentence_entities);
        setTrain_entities(train_entities);
    }

    /**
     * Returns the raw text from the CoNLL formatted annotated file.
     *
     * @param file_text_lines: the path from which you want to read the dataset
     */
    private String get_raw_text_from_conll(List<String> file_text_lines){
        StringBuilder sb = new StringBuilder();
        for(String line: file_text_lines){
            sb.append(line.split(" ")[0]+ " ");
        }
        String raw_text = sb.toString().trim();
        // Remove space before punctuations
        //raw_text = raw_text.replaceAll("\\s+(?=\\p{Punct})", "");
        return raw_text;
    }

    /**
     * Returns the textual content of a CoNLL annotated dataset as raw or CoNLL formatted text
     *
     * @param absPath: the path from which you want to read the dataset
     * @param raw: boolean value to choose between returning a raw or CoNLL formatted string
     */
    private String get_file_text(String absPath, boolean raw) {
        try {
            BufferedReader reader = new BufferedReader(new FileReader(absPath));
            List<String> file_text_lines = reader.lines().collect(Collectors.toList());
            if (raw){
                return get_raw_text_from_conll(file_text_lines);
            } else{
                String file_text = String.join("\n", file_text_lines);
                return file_text;
            }
        } catch (Exception e) {
            System.out.println(e.getStackTrace());
            return "";
        }
    }

    // *****************************************************************************************************************
    // Public methods
    // *****************************************************************************************************************

    /**
     * Preparing test data
     *
     * @param abspath: the path from which you want to read the dataset
     * @param data_map: a CSV tuple of the semantics of each column in the CoNLL training data.
     */
    public Dataset(String abspath, String data_map) {

        String conll = get_file_text(abspath, false);
        String raw = get_file_text(abspath, true);

        // Tokenize to sentences
        String[] cs = conll.split("\n\n");
        String[] rs = raw.split("  ");

        assert cs.length == rs.length;

        HashMap<Integer, String> conll_train_sentences = new HashMap<>();
        HashMap<Integer, String> raw_train_sentences = new HashMap<>();
        for(int i = 0 ; i < cs.length ; i++){
            conll_train_sentences.put(i, cs[i]);
            raw_train_sentences.put(i, rs[i]);
        }
        setConll_train_sentences(conll_train_sentences);
        setRaw_train_sentences(raw_train_sentences);

        HashMap<Integer, String> conll_test_sentences = new HashMap<>();
        HashMap<Integer, String> raw_test_sentences = new HashMap<>();
        for(int i = 0 ; i < cs.length ; i++){
            conll_test_sentences.put(i, cs[i]);
            raw_test_sentences.put(i, rs[i]);
        }
        setConll_test_sentences(conll_test_sentences);
        setRaw_test_sentences(raw_test_sentences);

        extract_entities_from_train_conll(data_map);
    }

    /**
     * This is for testing to have a subset of the data only
     *
     * @param abspath: the path from which you want to read the dataset
     * @param data_map: a CSV tuple of the semantics of each column in the CoNLL training data.
     * @param size: the size of the pool of sentences
     */
    public Dataset(String abspath, String data_map, int size) {
        String conll = get_file_text(abspath, false);
        String raw = get_file_text(abspath, true);

        // Tokenize to sentences
        String[] cs = conll.split("\n\n");
        String[] rs = raw.split("  ");

        assert cs.length == rs.length;

        HashMap<Integer, String> conll_train_sentences = new HashMap<>();
        HashMap<Integer, String> raw_train_sentences = new HashMap<>();
        for(int i = 0 ; i < cs.length ; i++){
            conll_train_sentences.put(i, cs[i]);
            raw_train_sentences.put(i, rs[i]);

            if (i >= size-1)
                break;
        }
        setConll_train_sentences(conll_train_sentences);
        setRaw_train_sentences(raw_train_sentences);

        HashMap<Integer, String> conll_test_sentences = new HashMap<>();
        HashMap<Integer, String> raw_test_sentences = new HashMap<>();
        for(int i = 0 ; i < cs.length ; i++){
            conll_test_sentences.put(i, cs[i]);
            raw_test_sentences.put(i, rs[i]);

            if (i >= size-1)
                break;
        }
        setConll_test_sentences(conll_test_sentences);
        setRaw_test_sentences(raw_test_sentences);

        extract_entities_from_train_conll(data_map);
    }

    /**
     * Returns all train sentences from the train dataset CoNLL formatted
     *
     * @return map of sentence-id to sentence CoNLL text
     */
    public HashMap<Integer, String> getConll_train_sentences() {
        return conll_train_sentences;
    }

    public void setConll_train_sentences(HashMap<Integer, String> conll_train_sentences) {
        this.conll_train_sentences = conll_train_sentences;
    }

    /**
     * Returns all raw train sentences from the train dataset
     *
     * @return map of sentence-id to raw sentence text
     */
    public HashMap<Integer, String> getRaw_train_sentences() {
        return raw_train_sentences;
    }

    public void setRaw_train_sentences(HashMap<Integer, String> raw_train_sentences) {
        this.raw_train_sentences = raw_train_sentences;
    }

    /**
     * Returns all the annotated entities from the train sentences
     *
     * @return set of entities
     */
    public Set<String> getTrain_entities() {
        return train_entities;
    }

    public void setTrain_entities(Set<String> train_entities) {
        this.train_entities = train_entities;
    }

    /**
     * Returns all test sentences from the test dataset CoNLL formatted
     *
     * @return map of sentence-id to sentence CoNLL text
     */
    public HashMap<Integer, String> getConll_test_sentences() {
        return conll_test_sentences;
    }

    public void setConll_test_sentences(HashMap<Integer, String> conll_test_sentences) {
        this.conll_test_sentences = conll_test_sentences;
    }

    /**
     * Returns all raw test sentences from the test dataset
     *
     * @return map of sentence-id to raw sentence text
     */
    public HashMap<Integer, String> getRaw_test_sentences() {
        return raw_test_sentences;
    }

    public void setRaw_test_sentences(HashMap<Integer, String> raw_test_sentences) {
        this.raw_test_sentences = raw_test_sentences;
    }

    /**
     * Returns all the annotated entities for each sentence from the train dataset
     *
     * @return map of sentence-id to entities set in each sentence
     */
    public HashMap<Integer, Set<String>> getTrain_sentence_entities() {
        return train_sentence_entities;
    }

    public void setTrain_sentence_entities(HashMap<Integer, Set<String>> train_sentence_entities) {
        this.train_sentence_entities = train_sentence_entities;
    }
}