# -*- coding: utf-8 -*-

import math
import jieba
from collections import Counter
from collections import defaultdict


def TFIDF(dataset):
    # 統(tǒng)計每個文章中每個詞語的tf
    j = 0
    doc_len = len(dataset)
    cntr = [0] * doc_len          # Counter每個文章
    total = [0] * doc_len          # 獲取每個文章中各自的詞語總數(shù)
    word_tf = [0] * doc_len          # 獲取每個文章中每個詞語的tf
    for sent in dataset:
        iters = jieba.cut(sent)
        cntr[j] = Counter(iters)    
        total[j] = sum(cntr[j].values())    
        word_tf[j] = defaultdict(int)
        for i in cntr[j].keys():        
            word_tf[j][i] = cntr[j][i] / total[j]
        j += 1
    
    
    # 構建詞匯表和總詞頻
    word_cnt = defaultdict(int)
    for word_list in word_tf:
        for word in word_list:
            word_cnt[word] += 1
    
    # 統(tǒng)計詞匯在文檔中的頻率
    doc_tf = defaultdict(int)
    for word in word_cnt:
        for article in dataset:
            if word in article:
                doc_tf[word] += 1
    
    # 構建詞匯的idf
    word_idf = {}
    for word in doc_tf:
        word_idf[word] = math.log(doc_len / (doc_tf[word] + 1))
        
    # 構建每篇文章中詞的tf-idf     
    t = 0    
    word_tf_idf = [0] * doc_len        
    for article in word_tf:
        word_tf_idf[t] = defaultdict(int)
        for word in article:
            word_tf_idf[t][word] = article[word] * word_idf[word]            
        t += 1
    
    # 對每篇文章每個詞的tf-idf進行由大到小排序
    t = 0
    for i in word_tf_idf:
        for j in i:  
            word_tf[t] = dict(sorted(word_tf[t].items(), key=lambda x:x[1], reverse=True))
            word_tf_idf[t] = dict(sorted(word_tf_idf[t].items(), key=lambda x:x[1], reverse=True))        
        t += 1
            
    return word_tf_idf, word_tf
    
    

if __name__ == '__main__':

    dataset = ["中國的北京是中國的首都","那里有一句話叫做“北京歡迎您”,北京啊北京",
               "北京歡迎您！相信夢想，夢，相信有夢就會有奇跡"]    
    word_tf_idf,word_tf = TFIDF(dataset)
    
    print("=============tf-idf=============")
    for i in word_tf_idf:
        print(i)
    print("=============tf=============")
    for i in word_tf:
        print(i)