# 共起行列を作成する関数

## テキストから辞書（単語とIDのリスト）、コーパス（単語を単語IDに置き換えたもの）
def preprocess(text):
    #入力されたテキストをすべて小文字にして、単語ごとに分割
    words = word_tokenize(text.lower())

    # keyが単語でvalueがIDのリストとkeyがIDでvalueが単語の辞書を作る
    word_to_id = {}
    id_to_word = {}

    # 単語分割したテキスト（入力）をひとつひとつチェックして、辞書になければ新たなIDを振って保存する
    for word in words:
         if word not in word_to_id:
                new_id = len(word_to_id)
                word_to_id[word] = new_id
                id_to_word[new_id] = word

    #単語分割したテキスト（入力）を単語IDに置き換える（→corpus)
    corpus = [word_to_id[w] for w in words]
    corpus = np.array(corpus)
    
    return corpus,word_to_id,id_to_word

# 単語ベクトルの列（行列）を作成する
def create_co_matrix(corpus, vocab_size, window_size=1):

    # 総語数の取得
    corpus_size = len(corpus)
    
    # 行数、列数を単語の種類数とする行列(co_matrix）を作成
    # vocab_sizeは単語の種類数
    co_matrix = np.zeros((vocab_size, vocab_size), dtype=np.int32)

    # idxは単語の出現順序
    # word_idは単語ID
    for idx, word_id in enumerate(corpus):

        # 単語ベクトルの取得
        for i in range(1, window_size + 1):
            # コンテクストサイズの取得
            left_idx = idx - i
            right_idx = idx + i

            # 対象とする単語の左に単語がある場合処理を行う
            if left_idx >= 0:
                # 対象とする単語の左の単語の列に1を足す
                left_word_id = corpus[left_idx]
                co_matrix[word_id, left_word_id] += 1

            if right_idx < corpus_size:
                # 対象とする単語の右の単語の列に1を足す
                right_word_id = corpus[right_idx]
                co_matrix[word_id, right_word_id] += 1

    return co_matrix

# 対象テキストの取得
f = open("../samples/alice_all.txt","r")
raw_text = f.read()
f.close()

# preprocessとcreat_co_matrix()を使って共起行列Vを作成
corpus,word_to_id,id_to_word = preprocess(raw_text)
vocab_size = len(word_to_id)
V = create_co_matrix(corpus,vocab_size)

# 対象テキストの取得
f = open("../samples/alice_all.txt","r")
raw_text1 = f.read()
f.close()

f = open("../samples/tom_sawyer.txt","r")
raw_text2 = f.read()
f.close()

# テキストの結合
raw_text0 = raw_text1 + "\n" + raw_text2

# preprocessとcreat_co_matrix()を使って共起行列Vを作成
corpus0,word_to_id0,id_to_word0 = preprocess(raw_text0)
vocab_size0 = len(word_to_id0)
V = create_co_matrix(corpus0,vocab_size0)