import os, re
from stanfordcorenlp import StanfordCoreNLP
from jpype import *

root_path = "D:\\change\\code_ner\\grammer\\hanlp"
djclass_path = "-Djava.class.path=" + root_path + os.sep + "hanlp-1.6.2.jar;" + root_path
startJVM(getDefaultJVMPath(), djclass_path, "-Xms1g", "-Xmx1g")

stanford_nlp = StanfordCoreNLP(r'E:\stanfordnlp', lang='zh')

drop_pos_set = set(
    ['xu', 'xx', 'y', 'yg', 'wh', 'wky', 'wkz', 'wp', 'ws', 'wyy', 'wyz', 'wb', 'u', 'ud', 'ude1', 'ude2', 'ude3',
     'udeng', 'udh'])
han_pattern = re.compile(r'[^\dA-Za-z\u3007\u4E00-\u9FCB\uE815-\uE864]+')

NLPTokenizer = JClass('com.hankcs.hanlp.tokenizer.NLPTokenizer')

Tokenizer = JClass('com.hankcs.hanlp.tokenizer.StandardTokenizer')

HanLP = JClass('com.hankcs.hanlp.HanLP')


def to_string(sentence, return_generator=False):
    if return_generator:
        return (word_pos_item.toString().split('/') for word_pos_item in Tokenizer.segment(sentence))
    else:
        # res=[(word_pos_item.toString().split('/')[0],word_pos_item.toString().split('/')[1]) for word_pos_item in Tokenizer.segment(sentence)]
        return [(word_pos_item.toString().split('/')[0], word_pos_item.toString().split('/')[1]) for word_pos_item in
                Tokenizer.segment(sentence)]




def to_string_hanlp(sentence, return_generator=False):
    if return_generator:
        return (word_pos_item.toString().split('/') for word_pos_item in HanLP.segment(sentence))
    else:
        # res=[(word_pos_item.toString().split('/')[0],word_pos_item.toString().split('/')[1]) for word_pos_item in Tokenizer.segment(sentence)]
        return [(word_pos_item.toString().split('/')[0], word_pos_item.toString().split('/')[1]) for word_pos_item in
                Tokenizer.segment(sentence)]


def seg_sentences(sentence, with_filter=True, return_generator=False):
    segs = to_string(sentence, return_generator=return_generator)
    # print(segs)
    # g=[]
    if with_filter:
        g = [word_pos_pair[0] for word_pos_pair in segs if
             len(word_pos_pair) == 2 and word_pos_pair[0] != ' ' and word_pos_pair[1] not in drop_pos_set]
    else:
        g = [word_pos_pair[0] for word_pos_pair in segs if len(word_pos_pair) == 2 and word_pos_pair[0] != ' ']
    return iter(g) if return_generator else g


def ner_stanford(raw_sentence, return_list=True):
    if len(raw_sentence.strip()) > 0:
        return stanford_nlp.ner(raw_sentence) if return_list else iter(stanford_nlp.ner(raw_sentence))


def ner_hanlp(raw_sentence, return_list=True):
    if len(raw_sentence.strip()) > 0:
        return NLPTokenizer.segment(raw_sentence) if return_list else iter(NLPTokenizer.segment(raw_sentence))


def cut_stanford(raw_sentence, return_list=True):
    if len(raw_sentence.strip()) > 0:
        return stanford_nlp.pos_tag(raw_sentence) if return_list else iter(stanford_nlp.pos_tag(raw_sentence))
