from PyPDF2 import PdfReader, PdfWriter
from unstructured.partition.pdf import partition_pdf

import os
from os import path

# Create the output directory if it doesn't exist
# os.makedirs('./output', exist_ok=True)
path = path.abspath(path.dirname(__file__))

# pdf_file = path + '/sample01.pdf'

filename =  path + "/sample02.pdf"

# Read the original PDF
input_pdf = PdfReader(f'{filename}')

batch_size = 2
num_batches = len(input_pdf.pages) // batch_size + 1

filename = path + "/output" 
# Extract batches of 100 pages from the PDF
for b in range(num_batches):
    writer = PdfWriter()

    # Get the start and end page numbers for this batch
    start_page = b * batch_size
    end_page = min((b+1) * batch_size, len(input_pdf.pages))

    # Add pages in this batch to the writer
    for i in range(start_page, end_page):
        writer.add_page(input_pdf.pages[i])

    # Save the batch to a separate PDF file
    batch_filename = f'{filename}-batch{b+1}.pdf'
    with open(batch_filename, 'wb') as output_file:
        writer.write(output_file)

    # Now you can use the `partition_pdf` function from Unstructured.io to analyze the batch
    elements = partition_pdf(filename=batch_filename)
    print(elements)
    # Do something with `elements`...
    
    # This will process without issue
    # 抽取表格數(shù)據(jù)
	elements = partition_pdf("copy-protected.pdf", strategy="hi_res")

import PyPDF2

def split_pdf(input_pdf, output_prefix, pages_per_file=5):
    with open(input_pdf, 'rb') as file:
        pdf_reader = PyPDF2.PdfFileReader(file)
        num_pages = pdf_reader.numPages

        for i in range(0, num_pages, pages_per_file):
            pdf_writer = PyPDF2.PdfFileWriter()
            output_pdf = f'{output_prefix}_{i // pages_per_file + 1}.pdf'

            for j in range(i, min(i + pages_per_file, num_pages)):
                page = pdf_reader.getPage(j)
                pdf_writer.addPage(page)

            with open(output_pdf, 'wb') as new_file:
                pdf_writer.write(new_file)

            print(f'已創(chuàng)建文件: {output_pdf}')

# 示例調(diào)用
split_pdf('large_file.pdf', 'output_split')

def main():
    directory = input("請輸入PDF文件所在目錄：")
    pdf_files = get_pdf_files(directory)
    split_rule = get_split_rule()
    output_directory = get_output_directory()

    for file in pdf_files:
        output_files = split_pdf(file, split_rule)
        save_output_files(output_files, output_directory)

    print("分割完成！")

if __name__ == "__main__":
    main()