import pandas as pd
import os
# import jieba
from snownlp import SnowNLP

CONST_COLUMN_WORDS_COUNT = '词数'


def read_keywords(keywords_xlsx_file):
    df = pd.read_excel(keywords_xlsx_file)
    return list(df['Phrase'])


def build_df(lst_cols, rows):
    df = pd.DataFrame(columns=lst_cols, index=rows.keys())
    for k, v in rows.items():
        df.loc[k] = pd.Series(v)
    # print(df)
    return df


def read_file(file):
    with open(file, 'r', encoding='utf8', errors='ignore') as f:
        return f.read()


def build_row_counts(txt_file, keywords):
    content = read_file(txt_file)
    counts = {}
    for keyword in keywords:
        if keyword == CONST_COLUMN_WORDS_COUNT:
            sn = SnowNLP(content)
            # seg_list = jieba.cut(content)
            counts[keyword] = len(sn.words)
        else:
            counts[keyword] = content.count(keyword)
    file_name = os.path.basename(txt_file).replace('.txt', '')
    return {file_name: counts}


def build_rows(txt_files_dir, keywords):
    rows = {}
    for f in os.listdir(txt_files_dir):
        full_name = os.path.join(txt_files_dir, f)
        rows.update(build_row_counts(full_name, keywords))
    # print(rows)
    return rows


def count_and_export(keywords_xlsx_file, txt_files_dir, export_xlsx_file):
    keywords = [CONST_COLUMN_WORDS_COUNT]
    keywords.extend(read_keywords(keywords_xlsx_file))
    # print(keywords)
    # build_df(['a', 'b'], {'x': {'a': 11, 'b': 22}, 'y': {'a': 111, 'b': 222}})
    df = build_df(keywords, build_rows(txt_files_dir, keywords))
    df.to_excel(export_xlsx_file)


if __name__ == '__main__':
    count_and_export(r'C:\Users\YANG.LEI\Downloads\test\Keyword_Phrases.xlsx',
                     r'C:\Users\YANG.LEI\Downloads\test\txt实验文件',
                     r'C:\Users\YANG.LEI\Downloads\test\导出.xlsx'
                     )
    print('main end')



#!/usr/bin/python
# -*- coding: utf-8 -*-

import textract, re
import PyPDF2
from snownlp import SnowNLP
import pandas as pd
import numpy as np
import os
import jieba.posseg as pseg
import jieba
from tika import parser


os.chdir('annual reports')
#first, read the pdfs as text, and save them as txt, which will make processing faster in the future
for filename in os.listdir('annual reports'):
   if filename.endswith('.pdf'):
    try:
        text = parser.from_file('pdf/' + filename)
        filename = filename.replace(".pdf", "")
        text_file = open('text_doc/' + filename + ".txt", "w+")
        n = text_file.write(text['content'])
        text_file.close()
        print("finished writing" + filename)
    except:
        print("can't write")