import numpy as np
import re

with open('语料.txt') as f:
    X_raw = f.readlines()


X = [list(sentence) for sentence in X_raw]  # 拆词
regex = re.compile(u'[\u4e00-\u9fa5]')
X = [[word for word in sentence if regex.match(word) is not None] for sentence in X]
# X=[[word for word in sentence if word not in stops] for sentence in X] # 停词
X = [' '.join(sentence) for sentence in X]

from sklearn.feature_extraction import text

max_n_gram = 4
count_vectorizer = text.CountVectorizer(ngram_range=[1, max_n_gram],
                                        token_pattern=r'\S',
                                        max_df=0.8,
                                        min_df=5)

count_vectorizer.fit(X)

X_cnt_vec = count_vectorizer.transform(X)

X_cnt_vec.toarray().sum(axis=0)

# 每个 feature 是什么（文本）
feature_names = count_vectorizer.get_feature_names()

# 每个feature含多少词（n-gram）
feature_len = np.array([feature_name.count(' ') + 1 for feature_name in feature_names])

# 全部语料中，每个feature出现多少次
feature_cnt = X_cnt_vec.toarray().sum(axis=0)

# 每种长度的词，在全部语料中出现多少次
feature_toal = {i: feature_cnt[feature_len == i].sum() for i in range(1, max_n_gram + 1)}

res = []
for idx in range(len(feature_names)):
    if feature_len[idx] == 1:
        continue
    feature_name_split = feature_names[idx].split(' ')
    p_x_y = feature_cnt[idx] / feature_toal[feature_len[idx]]
    px_multiply_py = 0
    for i in range(1, feature_len[idx]):
        part_left = ' '.join(feature_name_split[:i])
        part_right = ' '.join(feature_name_split[i:])
        idx_left = feature_names.index(part_left)
        idx_right = feature_names.index(part_right)
        if part_left in feature_names and part_right in feature_names:
            px_multiply_py_ = feature_cnt[idx_left] * feature_cnt[idx_right] / feature_toal[feature_len[idx_left]] / \
                              feature_toal[feature_len[idx_right]]
            if px_multiply_py_ > px_multiply_py:
                px_multiply_py = px_multiply_py_
    res.append([feature_names[idx], p_x_y / px_multiply_py])

res

sorted(res, key=lambda x: x[1], reverse=True)[:20]

import collections

# 取上一步的结果，这里应当卡域值，但作为演示就取top20
possible_new_words = sorted(res, key=lambda x: x[1], reverse=True)[:20]


def cal_entropy(cnt):
    # 计算熵
    freq = np.array(list(cnt.values()))
    prob = freq / freq.sum()
    return -(prob * np.log2(prob)).sum()


new_word = []
for word in possible_new_words:
    regex_left = re.compile(r'[\S]* {word}'.format(word=word[0]))
    regex_right = re.compile(r'{word} [\S]*'.format(word=word[0]))

    # 遍历语料，找到左字和右字（用空格分割的一个基础单元，字或词或英语单词）
    cnt_left, cnt_right = collections.Counter(), collections.Counter()
    for x in X:
        cnt_left.update(collections.Counter(regex_left.findall(x)))
        cnt_right.update(collections.Counter(regex_right.findall(x)))

    left_entropy, right_entropy = cal_entropy(cnt_left), cal_entropy(cnt_right)
    if left_entropy > 1 and right_entropy > 1:
        new_word.append(word)

new_word