app.py

import streamlit as st
import pandas as pd
import numpy as np
from azure.core.credentials import AzureKeyCredential
from azure.ai.formrecognizer import DocumentAnalysisClient
import time
from dotenv import load_dotenv
import os

load_dotenv()

# Azure Form Recognizer credentials
endpoint = os.environ["endpoint"]
key = os.environ["key"]

# Initialize DocumentAnalysisClient
document_analysis_client = DocumentAnalysisClient(
    endpoint=endpoint, credential=AzureKeyCredential(key)
)

# Sidebar for document type selection
document_type = st.sidebar.selectbox(
    "Select Document Type", ["Invoice", "Receipt", "Identity Document"]
)

# Center area for file upload
uploaded_file = st.file_uploader(f"Choose or drag an {document_type} file", type=["pdf", "png", "jpg", "jpeg"])

# Initialize session state for processing results
if "processed" not in st.session_state:
    st.session_state.processed = False
    st.session_state.tables = []
    st.session_state.kv_df = pd.DataFrame()

# Process the uploaded file if the button is clicked
if uploaded_file is not None:
    if st.button("Upload and Process"):
        with st.spinner("Processing..."):
            time.sleep(1)  # Simulate some delay to show the spinner

            # Analyze the document
            poller = document_analysis_client.begin_analyze_document(
                "prebuilt-document", document=uploaded_file
            )
            result = poller.result()

            # Extract key-value pairs and store them in a dictionary
            kv_dict = {}
            for kv_pair in result.key_value_pairs:
                if kv_pair.key and kv_pair.value:  # Ensure both key and value exist
                    clean_key = kv_pair.key.content.rstrip(':')
                    kv_dict[clean_key] = kv_pair.value.content
            kv_df = pd.DataFrame(list(kv_dict.items()), columns=["Key", "Value"]).T

            # Set the first row as the header
            header = kv_df.iloc[0]

            # Create a new DataFrame with the header and the remaining rows
            kv_df = kv_df[1:].reset_index(drop=True)

            # Insert the new header row
            header_df = pd.DataFrame([header], columns=range(len(header)))

            # Concatenate the header row DataFrame with the original DataFrame
            kv_df = pd.concat([header_df, kv_df], ignore_index=True)

            st.session_state.kv_df = kv_df
            st.session_state.tables = []

            if result.tables:
                for table in result.tables:
                    data = []
                    for cell in table.cells:
                        data.append([cell.row_index, cell.column_index, cell.content])

                    table_df = pd.DataFrame(data, columns=["row_index", "column_index", "content"])
                    table_df = table_df.pivot(index="row_index", columns="column_index", values="content")
                    # Reset the column index
                    table_df.reset_index(drop=True, inplace=True)

                    # Set the first row as header
                    table_df.columns = table_df.iloc[0]  # Take the first row as column names
                    table_df = table_df.drop(table_df.index[0])  # Remove the first row after setting it as headers

                    # Optionally reset the index if required
                    table_df.reset_index(drop=True, inplace=True)

                    st.session_state.tables.append(table_df)

            st.session_state.processed = True

# Predefine default columns in the selected table
default_columns = ['Invoice No', 'Invoice Date', 'Customer Name']

# Sidebar for table and key-value pair selection
if st.session_state.processed:
    if not st.session_state.kv_df.empty:
        st.sidebar.write("Key-Value Pairs found.")
    else:
        st.sidebar.write("No Key-Value Pairs found in document.")

    if st.session_state.tables:
        # Display available tables for selection in the sidebar
        table_options = ["Select"] + [f"Table {i + 1}" for i in range(len(st.session_state.tables))]
        selected_table_option = st.sidebar.selectbox(
            "Select table to join key-value pairs:", table_options, index=0
        )

        # Proceed only if a table is selected
        if selected_table_option != "Select":
            selected_table_index = table_options.index(selected_table_option) - 1  # Adjust for "None"
            selected_table_df = st.session_state.tables[selected_table_index]

           
             # Add default columns to the left side of the table with NaN values initially
            for col in default_columns:
                if col not in selected_table_df.columns:
                    selected_table_df[col] = np.nan

            # Move default columns to the left
            # We include only those columns that exist in the table
            existing_columns = [col for col in default_columns if col in selected_table_df.columns]
            remaining_columns = [col for col in selected_table_df.columns if col not in existing_columns]
            selected_table_df = selected_table_df[[*existing_columns, *remaining_columns]]
            
            if not st.session_state.kv_df.empty:
                # Allow the user to select multiple columns using checkboxes
                first_row_values = st.session_state.kv_df.iloc[0].values
                second_row_values = st.session_state.kv_df.iloc[1].values
                key_value_options = [
                    f"{key}: {value}" for key, value in zip(first_row_values, second_row_values)
                ]
                
                selected_columns = []
                
                st.sidebar.write("Select key-value pairs to the default columns:")

                # Create draggable interface for each key-value pair
                for i, option in enumerate(key_value_options):
                    selected_column_option = st.sidebar.selectbox(
                        f"Insert {option} to the column:",
                        ["Select"] + default_columns,
                        key=f"col_select_{i}"
                    )
                    
                    if selected_column_option != "Select":
                        selected_columns.append((i, selected_column_option))

                # Insert selected key-value pairs into the appropriate columns
                for idx, selected_column in selected_columns:
                    st.session_state.kv_df[selected_column] = second_row_values[idx]
                    
                    # Fill the column in the table with the selected key-value pair
                    selected_table_df[selected_column] = st.session_state.kv_df.iloc[1, idx]

                st.write(f"Selected Table {selected_table_index + 1} with Key-Value Pairs Added:")

            # st.dataframe(selected_table_df)

            # Allow row and column selection for deletion
            st.sidebar.subheader("Manage Table")
            # Delete columns
            if st.sidebar.checkbox("Delete Columns"):
                col = list(selected_table_df.columns)
                col = col[3:]
                cols_to_delete = st.sidebar.multiselect(
                    "Select Columns to Delete", col
                )
                if st.sidebar.button("Remove Columns"):
                    st.session_state.tables[selected_table_index].drop(columns=cols_to_delete, inplace=True)

            # Delete rows
            if st.sidebar.checkbox("Delete Rows"):
                row_indices = st.sidebar.multiselect(
                    "Select Rows to Delete", selected_table_df.index.tolist()
                )
                if st.sidebar.button("Remove Rows"):
                    st.session_state.tables[selected_table_index].drop(index=row_indices, inplace=True)

            # Show the updated table
            # st.write(f"Updated Table {selected_table_index + 1}:")
            st.dataframe(selected_table_df)

            # Download updated table
            st.download_button(
                label=f"Download Updated Table {selected_table_index + 1} as CSV",
                data=selected_table_df.to_csv(index=False).encode('utf-8'),
                file_name=f"updated_table_{selected_table_index + 1}.csv",
                mime='text/csv',
            )
    else:
        st.sidebar.write("No tables found in the document.")

# Option to remove the file and clear the session
if st.session_state.processed and st.sidebar.button("Remove File"):
    st.session_state.processed = False
    st.session_state.tables = []
    st.session_state.kv_df = pd.DataFrame()
    st.rerun()