app.py

import re
import fitz
from PIL import Image
import pytesseract
import gradio as gr
import pandas as pd
import os

config_val = "--psm 6 -c tessedit_char_whitelist=0123456789,.-+"

# Rectangles for Form 1040 Pages 1 & 2
page1_rects = [
    [(464, 399), (576, 399), (575, 409), (462, 410)],
    [(462, 519), (577, 518), (577, 531), (463, 529)],
    [(225, 517), (340, 518), (339, 530), (224, 530)],
    [(225, 530), (339, 532), (340, 541), (225, 542)],
    [(464, 531), (576, 531), (576, 542), (464, 542)],
    [(464, 589), (578, 589), (577, 602), (464, 602)],
    [(463, 624), (578, 626), (576, 639), (464, 637)],
    [(462, 652), (576, 651), (577, 661), (464, 663)],
    [(463, 661), (578, 664), (578, 676), (462, 674)],
    [(464, 699), (578, 684), (578, 699), (464, 699)]
]
page2_rects = [
    [(462, 15), (575, 15), (576, 26), (463, 26)],
    [(462, 62), (577, 63), (579, 75), (462, 73)],
    [(463, 98), (576, 98), (578, 110), (462, 110)],
    [(461, 111), (576, 111), (578, 123), (459, 122)]
]

schedule1_rects = [
    [(470, 204), (579, 203), (577, 216), (471, 216)],  # Schedule 1 Line 3
    [(470, 228), (577, 229), (576, 240), (470, 240)],  # Schedule 1 Line 5
    [(362, 274), (466, 274), (468, 288), (360, 288)]   # Schedule 1 Line 8
]

adjusted_page1_rects = [[(x, y + 23) for (x, y) in rect] for rect in page1_rects]
adjusted_page2_rects = [[(x, y + 23) for (x, y) in rect] for rect in page2_rects]

def get_bounding_rect(points):
    xs = [pt[0] for pt in points]
    ys = [pt[1] for pt in points]
    return fitz.Rect(min(xs), min(ys), max(xs), max(ys))

    
def extract_numeric_values(pdf_file, schedule1_file=None, client_name="", draw_bboxes=False):
    try:
        if not client_name or client_name.strip() == "":
            return "Error: Client name is required.", None

        if isinstance(pdf_file, str):
            doc = fitz.open(pdf_file)
        else:
            pdf_file.seek(0)
            doc = fitz.open(stream=pdf_file.read(), filetype="pdf")

        if len(doc) < 2:
            return "Error: Main PDF must have at least 2 pages.", None, None

        zoom = fitz.Matrix(2, 2)
        page1 = doc[0]
        page2 = doc[1]

        page1_values, page2_values = [], []

        for rect_points in adjusted_page1_rects:
            rect = get_bounding_rect(rect_points)
            pix = page1.get_pixmap(matrix=zoom, clip=rect)
            cropped_img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
            w, h = cropped_img.size
            val_img = cropped_img.crop((int(0.4 * w), 0, w, h))
            raw = pytesseract.image_to_string(val_img, config=config_val).strip()
            value_text = re.sub(r"[^\d,.\-+]", "", raw)
            page1_values.append(value_text)
            if draw_bboxes:
                page1.draw_rect(rect, color=(1, 0, 0), width=1.5)

        for rect_points in adjusted_page2_rects:
            rect = get_bounding_rect(rect_points)
            pix = page2.get_pixmap(matrix=zoom, clip=rect)
            cropped_img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
            w, h = cropped_img.size
            val_img = cropped_img.crop((int(0.4 * w), 0, w, h))
            raw = pytesseract.image_to_string(val_img, config=config_val).strip()
            value_text = re.sub(r"[^\d,.\-+]", "", raw)
            page2_values.append(value_text)
            if draw_bboxes:
                page2.draw_rect(rect, color=(0, 0, 1), width=1.5)

        output_pdf_path = None
        if draw_bboxes:
            output_pdf_path = "output_shifted.pdf"
            doc.save(output_pdf_path)

        doc.close()
        
        output = [f"1040 Value {i+1}: {val}" for i, val in enumerate(page1_values + page2_values)]
        all_extracted_values = page1_values + page2_values
        schedule1_values = []

        output_folder = "./temp_outputs"
        os.makedirs(output_folder, exist_ok=True)
        schedule1_pdf_path = None  # initialize
        
        if schedule1_file:
            if isinstance(schedule1_file, str):
                doc = fitz.open(schedule1_file)
            else:
                schedule1_file.seek(0)
                doc = fitz.open(stream=schedule1_file.read(), filetype="pdf")

            if len(doc) >= 1:
                page = doc[0]
                for idx, rect_points in enumerate(schedule1_rects):
                    rect = get_bounding_rect(rect_points)
                    pix = page.get_pixmap(matrix=zoom, clip=rect)
                    cropped_img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
                    w, h = cropped_img.size
                    val_img = cropped_img.crop((int(0.4 * w), 0, w, h))
                    raw = pytesseract.image_to_string(val_img, config=config_val).strip()
                    value_text = re.sub(r"[^\d,.\-+]", "", raw)
                    schedule1_values.append(value_text)
                    if draw_bboxes:
                        page.draw_rect(rect, color=(0, 1, 0), width=1.5)  # green boxes
                schedule1_pdf_path = None
                if draw_bboxes and schedule1_file:
                    schedule1_pdf_path = "schedule1_bbox.pdf"
                    doc.save(schedule1_pdf_path)
        
                output += [f"Schedule 1 Line {i*2+1 if i < 2 else 8}: {val}" for i, val in enumerate(schedule1_values)]
                doc.close()

        else:
            # Create a temporary empty file as a placeholder
            empty_path = os.path.join(output_folder, "empty_schedule1.pdf")
            if not os.path.exists(empty_path):
                with open(empty_path, "wb") as f:
                    f.write(b"%PDF-1.4\n%%EOF")
            schedule1_pdf_output = empty_path
    
        output_dir = "./Clients_Output_Data_Form_1040.csv"
        save_to_csv_flat(all_extracted_values, schedule1_values, client_name=client_name, csv_path=output_dir)
        
        return "\n".join(output), output_dir, output_pdf_path, schedule1_pdf_path


    except Exception as e:
        return f"Error occurred:\n{str(e)}", None, None


def save_to_csv_flat(all_extracted_values, schedule1_values, client_name="Unknown Client", csv_path=None):

    # Define the directory path explicitly
    if csv_path is None:
        csv_path = "./Clients_Output_Data_Form_1040.csv"

    # Header components    
    header_level_1 = [
        "Client Name","Gross Comp", "Taxable Wages", "Taxable Interest Income: Sch. B", "Tax- Exempt Interest",
        "Qualified Dividends", "Ordinary Dividends", "Long Term Capital Gain or Loss",
        "Other Adjustments (from Schedule 1)", "Business Income or Loss (Schedule C)",
        "Rent/ Royalty (Schedule E)", "Other Income", "Standard Deduction", "Qualified Business Income Deduction",
        "Taxable Income", "Tax", "", "", "Total Tax"
    ]
    header_level_2 = [
        "","W2 Box 5", "Line 1", "Line 2b", "Line 2a", "Line 3a", "Line 3b", "Line 7",
        "Line 10", "Schedule 1, Line 3", "Schedule 1, Line 5", "Schedule 1, Line 8",
        "Line 12", "Line 13", "Line 15", "Line 16", "Line 20, Schedule 3", "Line 23, Schedule 2", "Line 24"
    ]

    # Flatten headers for CSV
    flat_columns = [
        f"{h1.strip()} - {h2.strip()}" if h1.strip() and h2.strip()
        else (h1.strip() + h2.strip()) for h1, h2 in zip(header_level_1, header_level_2)
    ]

    # If file doesn't exist, create new DataFrame and write headers
    # if os.path.exists(csv_path):
    #     df = pd.read_csv(csv_path)
    # else:
    #     df = pd.DataFrame(columns=flat_columns)

    if os.path.exists(csv_path):
        print(f"Reading existing CSV file: {csv_path}")
        df = pd.read_csv(csv_path)
        print(f"CSV columns: {df.columns.tolist()}")
        print(f"CSV rows before append: {len(df)}")
    else:
        print("CSV does not exist. Creating new DataFrame.")
        df = pd.DataFrame(columns=flat_columns)

    # Create new row with None
    new_row = pd.Series([None] * len(flat_columns), index=flat_columns)
    new_row.iloc[0] = client_name
    # Map Page 1-2 values
    line_mapping = {
        "Taxable Wages - Line 1": 0,
        "Taxable Interest Income: Sch. B - Line 2b": 1,
        "Tax- Exempt Interest - Line 2a": 2,
        "Qualified Dividends - Line 3a": 3,
        "Ordinary Dividends - Line 3b": 4,
        "Long Term Capital Gain or Loss - Line 7": 5,
        "Other Adjustments (from Schedule 1) - Line 10": 6,
        "Standard Deduction - Line 12": 7,
        "Qualified Business Income Deduction - Line 13": 8,
        "Taxable Income - Line 15": 9,
        "Tax - Line 16": 10,
        "Line 20, Schedule 3": 11,
        "Line 23, Schedule 2": 12,
        "Total Tax - Line 24": 13
    }

    for key, idx in line_mapping.items():
        if idx < len(all_extracted_values):
            new_row[key] = all_extracted_values[idx] if all_extracted_values[idx] != '' else '0'

    # Add Schedule 1 values
    if schedule1_values:
        new_row["Business Income or Loss (Schedule C) - Schedule 1, Line 3"] = schedule1_values[0] if schedule1_values[0] != '' else '0'
        new_row["Rent/ Royalty (Schedule E) - Schedule 1, Line 5"] = schedule1_values[1] if schedule1_values[1] != '' else '0'
        new_row["Other Income - Schedule 1, Line 8"] = schedule1_values[2] if schedule1_values[2] != '' else '0'

    # Append and save
    if not ((df == new_row.iloc[0]).all(axis=1)).any():   
        df = pd.concat([df, pd.DataFrame([new_row])], ignore_index=True)
    print(f"CSV rows after append: {len(df)}")
    df.to_csv(csv_path, index=False)
    print(f" Data saved to CSV: {csv_path}")

# Gradio UI
iface = gr.Interface(
    fn=extract_numeric_values,
    inputs=[
        gr.File(label="Upload Main Form 1040 PDF (Required)", file_types=[".pdf"]),
        gr.File(label="Upload Schedule 1 PDF (Optional)", file_types=[".pdf"]),
        gr.Textbox(label="Client Name", placeholder="Enter client name"),
        gr.Checkbox(label="Draw Bounding Boxes on Form 1040 and Schedule 1?")
    ],
    outputs=[
    gr.Textbox(label="Extracted Numeric Values", lines=20),
    gr.File(label="Download Excel Output"),
    gr.File(label="1040 Bounding Boxes PDF"),
    gr.File(label="Schedule 1 Bounding Boxes PDF")
    ],

    title="Form 1040 & Schedule 1 Extractor",
    description="Upload a 2-page 1040 Form and (optionally) Schedule 1 to extract numeric fields. Optionally draw bounding boxes."
)

# with gr.Blocks(title="Tax PDF Extractor") as demo:
#     gr.Markdown("##  Tax PDF Extractor")
#     gr.Markdown("Upload Form 1040 (at least 2 pages). Optionally upload Schedule 1 for extra fields.")

#     client_name = gr.Textbox(label="Client Name (Required)", placeholder="Enter your full name")

#     form_1040 = gr.File(label="Upload Main Form 1040 PDF (Required)", file_types=[".pdf"])

#     has_schedule1 = gr.Radio(
#         choices=["Yes", "No"],
#         label="Do you have Schedule 1?",
#         value="No"
#     )

#     schedule1 = gr.File(label="Upload Schedule 1 PDF (Optional)", file_types=[".pdf"], visible=False)

#     # Show/hide schedule1 upload box
#     def toggle_schedule1(choice):
#         return gr.update(visible=choice == "Yes")

#     has_schedule1.change(fn=toggle_schedule1, inputs=has_schedule1, outputs=schedule1)

#     output_text = gr.Textbox(label="Extracted Numeric Values", lines=20)
#     output_file = gr.File(label="Download Excel Output")

#     def wrapper_extract(main_pdf, schedule1_pdf, client_name):
#         if not client_name:
#             return "Error: Client name is required.", None
#         return extract_numeric_values(main_pdf, schedule1_pdf)

#     submit_btn = gr.Button("Extract Data")

#     submit_btn.click(
#         fn=wrapper_extract,
#         inputs=[form_1040, schedule1, client_name],
#         outputs=[output_text, output_file]
#     )


iface.launch(share=True)