# -*- coding: UTF-8 -*-
import numpy as np
import pandas as pd
import math
import jieba

def doc2tfidf_matrix():
    # 读取待编码的文件
    docs = ["少年救出溺水老人",
            "少年强中国强"]
    print(docs)
    # ['少年救出溺水老人', '少年强中国强']
    # 将文件每行分词，分词后的词语放入words中
    words=[]
    for i in range(len(docs)):
        docs[i]=jieba.lcut(docs[i].strip("\n"))
        words+=docs[i]
    print("words = %s" % words)
    # 找出分词后不重复的词语，作为词袋
    vocab=sorted(set(words),key=words.index)
    print("vocab = %s" % vocab)
    # ['少年', '救出', '溺水', '老人', '强', '中国']
    
    # 建立一个M行V列的全0矩阵，M问文档样本数，这里是行数，V为不重复词语数，即编码维度
    V=len(vocab)
    M=len(docs)
    onehot=np.zeros((M,V)) # 二维矩阵要使用双括号
    tf=np.zeros((M,V))
    print(tf)
    for i,doc in enumerate(docs):
        for word in doc:
            if word in vocab:
                pos=vocab.index(word)
                onehot[i][pos]=1
                tf[i][pos]+=1 # tf,统计某词语在一条样本中出现的次数
    #[[1. 1. 1. 1. 0. 0.]
    # [1. 0. 0. 0. 2. 1.]]
    row_sum=tf.sum(axis=1) # 行相加，得到每个样本出现的词语数
    print(row_sum)
    # [4. 4.]
    # 计算TF(t,d)
    tf=tf/row_sum[:,np.newaxis] #分母表示各样本出现的词语数，tf为单词在样本中出现的次数，[:,np.newaxis]作用类似于行列转置
    print(row_sum[:,np.newaxis])
    # [[4.]
    #  [4.]]
    print("tf=%s" % tf)
    # [[0.25 0.25 0.25 0.25 0.   0.  ]
    #  [0.25 0.   0.   0.   0.5  0.25]]
    # 计算DF(t,D)，IDF
    print("onehot = %s" % onehot)
    df=onehot.sum(axis=0) # 列相加，表示有多少样本包含词袋某词
    print("df = %s" % df)
    print("M = %s" % M)
    idf=list(map(lambda x:math.log10((M + 1)/(x+1)),df))
    print("idf = %s" % idf)
    # 计算TFIDF
    tfidf=tf*np.array(idf)
    tfidf=pd.DataFrame(tfidf,columns=vocab)
    return tfidf


print(doc2tfidf_matrix())