# 读取PDF
import PyPDF2
# 分析PDF的layout，提取文本
from pdfminer.high_level import extract_pages
from pdfminer.layout import LTTextContainer, LTRect
import pdfplumber

import os

'''
已知发现问题：
1.表格和文本提取错误，表格和文本内容在同一页，文本在前表格在后的，文本数据提取不出来
2.大模型抽取错，抽取2023年营业收入：主营业务收入、分产品的营业收入、变动比例被错误抽取
'''

# 查找PDF路径
pdf_path = '/Users/zhengfei/Desktop/科润智控.pdf'
page_obj = []

# 我们从PDF中提取页面
for pagenum, page in enumerate(extract_pages(pdf_path)):
    page_text = ''
    text_obj = {}
    # 初始化检查表的数量
    table_num = 0
    first_element= True
    table_extraction_flag= False
    # # 打开pdf文件
    pdf = pdfplumber.open(pdf_path)
    # 查找已检查的页面
    page_tables = pdf.pages[pagenum]
    # 找出本页上的表格数目
    tables = page_tables.find_tables()

    # 找到所有的元素
    page_elements = [(element.y1, element) for element in page._objs]
    # 对页面中出现的所有元素进行排序
    page_elements.sort(key=lambda a: a[0], reverse=True)

    # 查找组成页面的元素
    for i,component in enumerate(page_elements):
        # 提取PDF中元素顶部的位置
        pos= component[0]
        # 提取页面布局的元素
        element = component[1]
        
        # 检查该元素是否为文本元素
        if isinstance(element, LTTextContainer):
            # 检查文本是否出现在表中
            
            if table_extraction_flag == False:
                # 使用该函数提取每个文本元素的文本和格式
                line_text = element.get_text().replace('\s+', '').replace('\n', '').replace('\r', '')
                # 将每行的文本追加到页文本
                if len(line_text) > 5:
                    page_text += line_text
                
                # 附加每一行包含文本的格式
            else:
                # 省略表中出现的文本
                pass

        # 检查表的元素
        if isinstance(element, LTRect):
            # 如果第一个矩形元素
            if first_element == True and (table_num+1) <= len(tables):
                # 找到表格的边界框
                lower_side = page.bbox[3] - tables[table_num].bbox[3]
                upper_side = element.y1 
                # 将标志设置为True以再次避免该内容
                table_extraction_flag = True
                # 让它成为另一个元素
                first_element = False

                # 检查我们是否已经从页面中提取了表
                if element.y0 >= lower_side and element.y1 <= upper_side:
                    pass
                elif not isinstance(page_elements[i+1][1], LTRect):
                    table_extraction_flag = False
                    first_element = True
                    table_num+=1

    text_obj['page_num'] = pagenum
    text_obj['text'] = page_text
    print("pagenum:",pagenum,"  text:",page_text)

# 打印提取的文本
# print(page_obj)