Update src/streamlit_app.py

src/streamlit_app.py

@@ -27,63 +27,111 @@ import pandas as pd
 from streamlit_drawable_canvas import st_canvas
 import pytesseract
 import numpy as np
+from datetime import datetime
+import fitz  # PyMuPDF for PDF handling
 
 # Set Tesseract path - auto-detect based on OS
 if os.name == 'nt':  # Windows
     pytesseract.pytesseract.tesseract_cmd = r"C:\Program Files\Tesseract-OCR\tesseract.exe"
 else:  # Linux/Mac (HF Spaces uses Linux)
-    # On HF Spaces with packages.txt, tesseract is in system PATH
-    # No need to set path explicitly
     pass
 
 # Page configuration
 st.set_page_config(page_title="Remittance Data Viewer", layout="wide")
 
-# Custom CSS to reduce gaps between form fields and style buttons
+# Custom CSS
 st.markdown("""
 <style>
-    /* Reduce spacing between form fields */
-    .stTextInput > div > div > input,
-    .stTextArea > div > div > textarea,
-    .stSelectbox > div > div > div {
-        margin-bottom: 0px !important;
-    }
-    
-    div[data-testid="stVerticalBlock"] > div:has(div[data-testid="stTextInput"]),
-    div[data-testid="stVerticalBlock"] > div:has(div[data-testid="stTextArea"]),
-    div[data-testid="stVerticalBlock"] > div:has(div[data-testid="stSelectbox"]) {
-        margin-bottom: 4px !important;
-    }
-    
-    /* Reduce gap between selectbox and following elements */
-    .stSelectbox {
-        margin-bottom: 4px !important;
-    }
-    
-    /* Style for small buttons */
-    .stButton > button {
-        padding: 0.25rem 0.5rem !important;
-        font-size: 1.2rem !important;
-        line-height: 1 !important;
-        min-height: 2rem !important;
-        height: 2rem !important;
-    }
-    
-    /* Reduce padding in form containers */
-    [data-testid="stVerticalBlock"] > [data-testid="stVerticalBlock"] {
-        gap: 0.25rem !important;
-    }
-    
-    /* REDUCE GAP BETWEEN COLUMNS */
-    [data-testid="column"] {
-        padding-left: 0.5rem !important;
-        padding-right: 0.5rem !important;
-    }
-    
-    [data-testid="stHorizontalBlock"] {
-        gap: 0.5rem !important;
-    }
+  /* Reduce spacing between form fields (unchanged) */
+  .stTextInput > div > div > input,
+  .stTextArea > div > div > textarea,
+  .stSelectbox > div > div > div {
+      margin-bottom: 0px !important;
+  }
+  div[data-testid="stVerticalBlock"] > div:has(div[data-testid="stTextInput"]),
+  div[data-testid="stVerticalBlock"] > div:has(div[data-testid="stTextArea"]),
+  div[data-testid="stVerticalBlock"] > div:has(div[data-testid="stSelectbox"]) {
+      margin-bottom: 4px !important;
+  }
+  .stSelectbox { margin-bottom: 4px !important; }
+
+  /* Button styling (unchanged) */
+  .stButton > button {
+      padding: 0.25rem 0.5rem !important;
+      font-size: 0.85rem !important;
+      line-height: 1 !important;
+      min-height: 1.8rem !important;
+      height: 1.8rem !important;
+  }
+  .stButton > button[kind="primary"] {
+      background-color: #FF0000 !important;
+      border-color: #FF0000 !important;
+      color: white !important;
+  }
+  .stButton > button[kind="primary"]:hover {
+      background-color: #CC0000 !important;
+      border-color: #CC0000 !important;
+  }
+
+  /* Small vertical gaps */
+  [data-testid="stVerticalBlock"] > [data-testid="stVerticalBlock"] { gap: 0.25rem !important; }
+  [data-testid="column"] { padding-left: 0.5rem !important; padding-right: 0.5rem !important; }
+  [data-testid="stHorizontalBlock"] { gap: 0.5rem !important; }
+
+  /* Active OCR field highlighting */
+  .ocr-active {
+      border: 2px solid #ff4b4b !important;
+      box-shadow: 0 0 5px rgba(255, 75, 75, 0.5) !important;
+  }
+
+  /* left/right helper classes (still usable if you wrap content in html containers) */
+  .left-scroll, .right-scroll {
+      max-height: calc(100vh - 160px);
+      overflow-y: auto;
+      overflow-x: auto;
+      padding: 8px;
+      border: 1px solid #e6e6e6;
+      border-radius: 6px;
+      background: #ffffff;
+  }
+  .left-scroll canvas, .left-scroll img, .right-scroll canvas, .right-scroll img {
+      max-width: 100% !important;
+      height: auto !important;
+      display: block;
+  }
+
+  /* Make column elements create independent scroll areas:
+     - give them a fixed-ish max-height relative to viewport
+     - force overflow-y: auto so they show their own scrollbar
+  */
+  div[data-testid="column"] {
+      /* room for header + padding; tune the 140px to your UI if needed */
+      max-height: calc(100vh - 140px) !important;
+      overflow-y: auto !important;
+      overflow-x: hidden !important;
+      position: relative !important;
+  }
+
+  /* keep canvas/images from forcing expansion */
+  div[data-testid="column"] img,
+  div[data-testid="column"] canvas {
+      max-width: 100% !important;
+      height: auto !important;
+      display: block !important;
+  }
+
+  /* Thin subtle scrollbars for WebKit */
+  div[data-testid="column"]::-webkit-scrollbar { width: 10px; height: 10px; }
+  div[data-testid="column"]::-webkit-scrollbar-thumb { border-radius: 8px; background-color: rgba(0,0,0,0.12); }
+  div[data-testid="column"]::-webkit-scrollbar-track { background: transparent; }
+
+  /* IMPORTANT: remove any rule that forcibly hides page scrollbars.
+     If you previously set overflow: hidden on the main container it prevents children from scrolling properly.
+     DO NOT set `overflow: hidden` on the top-level app container.
+  */
+  /* removed: [data-testid="stAppViewContainer"] > .main { overflow: hidden !important; height: 100vh !important; } */
 </style>
+
 """, unsafe_allow_html=True)
 
 def load_jsonl(file):
@@ -100,66 +148,88 @@ def save_to_jsonl(data):
     jsonl_content = '\n'.join([json.dumps(record) for record in data])
     return jsonl_content
 
+def pdf_to_images(pdf_file):
+    """Convert PDF to list of PIL Images (one per page)"""
+    try:
+        pdf_bytes = pdf_file.read()
+        pdf_document = fitz.open(stream=pdf_bytes, filetype="pdf")
+        images = []
+        
+        for page_num in range(pdf_document.page_count):
+            page = pdf_document[page_num]
+            # Render page to an image (higher DPI for better quality)
+            # Using 3x zoom (300 DPI equivalent) for much better clarity
+            pix = page.get_pixmap(matrix=fitz.Matrix(3, 3), alpha=False)
+            img_data = pix.tobytes("png")
+            img = Image.open(io.BytesIO(img_data))
+            images.append(img)
+        
+        pdf_document.close()
+        return images
+    except Exception as e:
+        st.error(f"Error converting PDF: {str(e)}")
+        return []
+
 def perform_ocr(image, bbox):
     """Perform OCR on the selected region of the image"""
     try:
-        # bbox is [x1, y1, x2, y2]
         x1, y1, x2, y2 = int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3])
-        
-        # Ensure coordinates are within image bounds
         x1, y1 = max(0, x1), max(0, y1)
         x2, y2 = min(image.width, x2), min(image.height, y2)
-        
-        # Crop the image
         cropped = image.crop((x1, y1, x2, y2))
-        
-        # Perform OCR
         text = pytesseract.image_to_string(cropped, config='--psm 6').strip()
         return text
     except Exception as e:
         return f"OCR Error: {str(e)}"
 
-def scale_image_to_fixed_size(image, target_width=700, target_height=900):
-    """Scale and pad image to exact fixed size while maintaining aspect ratio and quality"""
-    # Convert image to RGB if it's not already (handles RGBA, L, etc.)
+def scale_image_to_fixed_size(image, max_width=900, max_height=1200):
+    """Scale image to fit within max dimensions while maintaining aspect ratio - NO PADDING"""
+    # Convert to RGB with proper handling
     if image.mode not in ('RGB', 'RGBA'):
         image = image.convert('RGB')
     elif image.mode == 'RGBA':
-        # Create white background for transparent images
         background = Image.new('RGB', image.size, (255, 255, 255))
-        background.paste(image, mask=image.split()[3])  # Use alpha channel as mask
+        background.paste(image, mask=image.split()[3])
         image = background
     
-    # Calculate scaling ratio to fit within target dimensions
-    width_ratio = target_width / image.width
-    height_ratio = target_height / image.height
-    
-    # Use the smaller ratio to ensure image fits within both constraints
+    # Calculate scaling ratio
+    width_ratio = max_width / image.width
+    height_ratio = max_height / image.height
     ratio = min(width_ratio, height_ratio)
     
-    # Calculate new dimensions
     new_width = int(image.width * ratio)
     new_height = int(image.height * ratio)
     
-    # Resize image with high-quality LANCZOS resampling
-    # Only resize if needed (don't upscale small images too much)
-    if ratio < 1.0 or (ratio > 1.0 and ratio < 1.5):
-        resized_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
-    else:
-        # For significant upscaling, use BICUBIC which can be sharper
-        resized_image = image.resize((new_width, new_height), Image.Resampling.BICUBIC)
+    # Always use LANCZOS for highest quality resampling
+    resized_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
+    
+    # Return without padding - image takes only the space it needs
+    return resized_image, ratio, 0, 0
+
+def swap_customer_supplier_details(index):
+    """Swap customer and supplier details"""
+    gt_parse = st.session_state.edited_data[index].get('gt_parse', {})
+    customer_supplier = gt_parse.get('Customer_supplier_details', {})
     
-    # Create a new image with target size and white background
-    final_image = Image.new('RGB', (target_width, target_height), (255, 255, 255))
+    # Store customer values
+    temp_customer_name = customer_supplier.get('Customer_name', '')
+    temp_customer_address = customer_supplier.get('Customer_address', '')
+    temp_customer_contact = customer_supplier.get('Customer_contact_info', '')
     
-    # Calculate position to paste resized image (center it)
-    paste_x = (target_width - new_width) // 2
-    paste_y = (target_height - new_height) // 2
+    # Swap: Customer ← Supplier
+    customer_supplier['Customer_name'] = customer_supplier.get('Supplier_name', '')
+    customer_supplier['Customer_address'] = customer_supplier.get('Supplier_address', '')
+    customer_supplier['Customer_contact_info'] = customer_supplier.get('Supplier_contact_info', '')
     
-    # Paste resized image onto white background
-    final_image.paste(resized_image, (paste_x, paste_y))
+    # Swap: Supplier ← Customer (from temp)
+    customer_supplier['Supplier_name'] = temp_customer_name
+    customer_supplier['Supplier_address'] = temp_customer_address
+    customer_supplier['Supplier_contact_info'] = temp_customer_contact
     
-    return final_image, ratio, paste_x, paste_y
+    # Update session state
+    gt_parse['Customer_supplier_details'] = customer_supplier
+    st.session_state.edited_data[index]['gt_parse'] = gt_parse
+    st.session_state.modified_indices.add(index)
 
 # Initialize session state
 if 'data' not in st.session_state:
@@ -172,6 +242,10 @@ if 'page' not in st.session_state:
     st.session_state.page = 'upload'
 if 'images' not in st.session_state:
     st.session_state.images = {}
+if 'pdf_metadata' not in st.session_state:
+    st.session_state.pdf_metadata = {}  # Store {filename: {'pages': [images], 'current_page': 0}}
+if 'current_page_num' not in st.session_state:
+    st.session_state.current_page_num = {}  # Track current page for each file
 if 'modified_indices' not in st.session_state:
     st.session_state.modified_indices = set()
 if 'ocr_active_section' not in st.session_state:
@@ -192,28 +266,56 @@ if 'save_message_time' not in st.session_state:
     st.session_state.save_message_time = None
 if 'just_saved' not in st.session_state:
     st.session_state.just_saved = False
+if 'just_swapped' not in st.session_state:
+    st.session_state.just_swapped = False
+if 'navigating_page' not in st.session_state:
+    st.session_state.navigating_page = False
 
-# Auto-save function
 def auto_save(index):
     """Automatically save changes to session state and mark as modified"""
     if st.session_state.edited_data:
         st.session_state.data = st.session_state.edited_data.copy()
         st.session_state.modified_indices.add(index)
 
-# Save button callback
-def save_changes_callback():
-    """Callback function for save button"""
-    auto_save(st.session_state.current_index)
-    st.session_state.save_message = "✅ Changes saved successfully!"
-    st.session_state.save_message_time = time.time()
+def activate_ocr_field(section, field, row_idx=None):
+    """Activate OCR for a specific field.
+       Toggle behavior: if the same field is already active, deactivate it to avoid repeated activations/looping.
+       Also ensures the line-item expander stays expanded when OCR is requested.
+    """
+    # If the requested field is already active, deactivate it (toggle off)
+    if (st.session_state.ocr_active_section == section and 
+        st.session_state.ocr_active_field == field and 
+        st.session_state.ocr_line_item_row == row_idx):
+        st.session_state.ocr_active_section = None
+        st.session_state.ocr_active_field = None
+        st.session_state.ocr_line_item_row = None
+    else:
+        # Activate new OCR target
+        st.session_state.ocr_active_section = section
+        st.session_state.ocr_active_field = field
+        st.session_state.ocr_line_item_row = row_idx
+
+        # If it's a line-item, mark that expander as expanded so it remains open after rerun
+        if section == 'Line_items' and row_idx is not None:
+            current_idx = st.session_state.get('current_index', 0)
+            expander_key = f"line_item_expander_{current_idx}_{row_idx}"
+            st.session_state[expander_key] = True
+
+    # Bump canvas_key to ensure canvas is refreshed/cleared when toggling OCR
+    st.session_state.canvas_key += 1
+    st.rerun()
+
+def is_ocr_active(section, field, row_idx=None):
+    """Check if this OCR button is currently active"""
+    return (st.session_state.ocr_active_section == section and 
+            st.session_state.ocr_active_field == field and 
+            st.session_state.ocr_line_item_row == row_idx)
 
 # PAGE 1: Upload Page
 if st.session_state.page == 'upload':
     st.title("📤 Remittance Data Viewer with OCR")
     st.markdown("### Upload your files to begin")
     
-    # Step 1: Upload JSONL
-    
     st.markdown("**Step 1: Upload JSONL File**")
     uploaded_file = st.file_uploader("Choose a JSONL file", type=['jsonl', 'json'])
     
@@ -226,65 +328,100 @@ if st.session_state.page == 'upload':
         except Exception as e:
             st.error(f"Error loading file: {str(e)}")
     
-    # Step 2: Upload Images
+    st.markdown("**Step 2: Upload Images/PDFs Folder**")
     
-    st.markdown("**Step 2: Upload Images Folder**")
-    
-    
-    uploaded_images = st.file_uploader(
-        "Choose image files", 
-        type=['png', 'jpg', 'jpeg', 'tiff', 'tif', 'bmp'],
+    uploaded_files = st.file_uploader(
+        "Choose image or PDF files", 
+        type=['png', 'jpg', 'jpeg', 'tiff', 'tif', 'bmp', 'pdf'],
         accept_multiple_files=True,
-        help="Select all images from your folder at once"
+        help="Select all images and PDFs from your folder at once"
     )
     
-    if uploaded_images:
-        # Load images into session state
+    if uploaded_files:
         images_dict = {}
-        for img_file in uploaded_images:
+        pdf_metadata = {}
+        
+        for file in uploaded_files:
             try:
-                image = Image.open(img_file)
-                images_dict[img_file.name] = image
+                file_ext = file.name.lower().split('.')[-1]
+                
+                if file_ext == 'pdf':
+                    # Convert PDF to images
+                    pdf_images = pdf_to_images(file)
+                    if pdf_images:
+                        # Store first page as the main image
+                        images_dict[file.name] = pdf_images[0]
+                        # Store all pages in metadata
+                        pdf_metadata[file.name] = {
+                            'pages': pdf_images,
+                            'total_pages': len(pdf_images),
+                            'current_page': 0
+                        }
+                        #st.info(f"📄 Converted PDF '{file.name}' ({len(pdf_images)} pages)")
+                else:
+                    # Handle regular images
+                    image = Image.open(file)
+                    images_dict[file.name] = image
+                    
             except Exception as e:
-                st.warning(f"Could not load image {img_file.name}: {str(e)}")
+                st.warning(f"Could not load file {file.name}: {str(e)}")
         
         st.session_state.images = images_dict
-        # Show summary of loaded images and matches with ground truth
+        st.session_state.pdf_metadata = pdf_metadata
+        
+        # Initialize current page tracking
+        for filename in pdf_metadata.keys():
+            if filename not in st.session_state.current_page_num:
+                st.session_state.current_page_num[filename] = 0
+        
         if st.session_state.data is not None:
-            # gather ground truth file names
             gt_file_names = [rec.get('file_name', '') for rec in st.session_state.data]
             matched_images = set()
             unmatched_gt_files = []
             
-            # Find matched images - CASE SENSITIVE EXACT MATCH ONLY
+            # Try to match with and without extensions
             for fname in gt_file_names:
                 if not fname:
                     continue
-                # Check for exact match in uploaded images
+                
+                # Try exact match first
                 if fname in images_dict:
                     matched_images.add(fname)
+                else:
+                    # Try adding common extensions
+                    found = False
+                    for ext in ['.pdf', '.png', '.jpg', '.jpeg', '.tiff', '.tif', '.bmp']:
+                        if fname + ext in images_dict:
+                            matched_images.add(fname)
+                            found = True
+                            break
+                    
+                    # Try matching filename without extension from uploaded files
+                    if not found:
+                        for uploaded_name in images_dict.keys():
+                            uploaded_base = uploaded_name.rsplit('.', 1)[0]
+                            if uploaded_base == fname:
+                                matched_images.add(fname)
+                                found = True
+                                break
             
-            # Find unmatched ground truth file names
             for fname in gt_file_names:
                 if fname and fname not in matched_images:
                     unmatched_gt_files.append(fname)
             
-            st.success(f"✅ Successfully loaded {len(images_dict)} images!")
+            st.success(f"✅ Successfully loaded {len(images_dict)} files ({len(pdf_metadata)} PDFs)!")
             st.info(f"🔎 Exact matches: {len(matched_images)}/{len([f for f in gt_file_names if f])}")
             
-            # Show unmatched files
             if unmatched_gt_files:
-                st.warning(f"⚠️ {len(unmatched_gt_files)} file(s) from JSONL not matched to images:")
+                st.warning(f"⚠️ {len(unmatched_gt_files)} file(s) from JSONL not matched:")
                 with st.expander(f"Show {len(unmatched_gt_files)} unmatched file names"):
                     for fname in unmatched_gt_files:
                         st.text(f"  • {fname}")
             else:
-                st.success("✅ All JSONL file names matched to images!")
+                st.success("✅ All JSONL file names matched to files!")
         else:
-            st.success(f"✅ Successfully loaded {len(images_dict)} images!")
-            st.info("ℹ️ Upload a JSONL file to see how many images match the ground truth 'file_name' field.")
-    
-    # Continue Button
+            st.success(f"✅ Successfully loaded {len(images_dict)} files ({len(pdf_metadata)} PDFs)!")
+            st.info("ℹ️ Upload a JSONL file to see how many files match the ground truth 'file_name' field.")
     
     if st.session_state.data is not None:
         col1, col2, col3 = st.columns([1, 1, 1])
@@ -292,17 +429,18 @@ if st.session_state.page == 'upload':
             if st.button("Continue to Viewer →", type="primary", use_container_width=True):
                 st.session_state.page = 'viewer'
                 st.session_state.modified_indices = set()
+                st.session_state.navigating_page = False
                 st.rerun()
 
 # PAGE 2: Viewer Page
 elif st.session_state.page == 'viewer':
-    # Clear old save messages (after 3 seconds)
     if st.session_state.save_message_time is not None:
         if time.time() - st.session_state.save_message_time > 3:
             st.session_state.save_message = None
             st.session_state.save_message_time = None
     
-    # Header with back button and download options
+    today_date = datetime.now().strftime("%Y-%m-%d")
+    
     col1, col2, col3, col4 = st.columns([1, 2, 2, 2])
     
     with col1:
@@ -312,9 +450,9 @@ elif st.session_state.page == 'viewer':
             st.session_state.ocr_active_field = None
             st.session_state.save_message = None
             st.session_state.save_message_time = None
+            st.session_state.navigating_page = False
             st.rerun()
     
-    # Download modified records and unmodified records separately
     with col2:
         if st.session_state.modified_indices:
             modified_data = [st.session_state.edited_data[i] for i in sorted(st.session_state.modified_indices)]
@@ -322,98 +460,175 @@ elif st.session_state.page == 'viewer':
             st.download_button(
                 label=f"⬇️ Download Modified ({len(modified_data)})",
                 data=jsonl_modified,
-                file_name="modified_remittance_data.jsonl",
+                file_name=f"modified_remittance_data_{today_date}.jsonl",
                 mime="application/jsonl",
                 type="primary",
                 use_container_width=True
             )
         else:
-            st.button(
-                "⬇️ No Modified Records",
-                disabled=True,
-                use_container_width=True
-            )
+            st.button("⬇️ No Modified Records", disabled=True, use_container_width=True)
     
-    # Download unmodified records (original data excluding modified)
     with col3:
         if st.session_state.modified_indices:
-            # Get original unmodified data
             unmodified_data = [st.session_state.data[i] for i in range(len(st.session_state.data)) 
                               if i not in st.session_state.modified_indices]
             jsonl_unmodified = save_to_jsonl(unmodified_data)
             st.download_button(
                 label=f"⬇️ Download Unmodified ({len(unmodified_data)})",
                 data=jsonl_unmodified,
-                file_name="unmodified_remittance_data.jsonl",
+                file_name=f"unmodified_remittance_data_{today_date}.jsonl",
                 mime="application/jsonl",
                 use_container_width=True
             )
         else:
-            st.button(
-                "⬇️ No Unmodified Records",
-                disabled=True,
-                use_container_width=True
-            )
+            st.button("⬇️ No Unmodified Records", disabled=True, use_container_width=True)
     
-    # Download all edited data
     with col4:
         jsonl_all = save_to_jsonl(st.session_state.edited_data)
         st.download_button(
             label=f"⬇️ Download All ({len(st.session_state.edited_data)})",
             data=jsonl_all,
-            file_name="all_remittance_data.jsonl",
+            file_name=f"all_remittance_data_{today_date}.jsonl",
             mime="application/jsonl",
             use_container_width=True
         )
-
-    
-    
-    
-    # File selector dropdown
-    file_names = [record.get('file_name', f'Record {i}') for i, record in enumerate(st.session_state.data)]
-    
-    selected_file = st.selectbox(
-        "Select a file to view:",
-        options=range(len(file_names)),
-        format_func=lambda x: f"{'✏️ ' if x in st.session_state.modified_indices else ''}{file_names[x]}",
-        index=st.session_state.current_index
-    )
-    
-    st.session_state.current_index = selected_file
-    current_record = st.session_state.edited_data[selected_file]
-    
     
+    file_names = [record.get('file_name', f'Record {i}') for i, record in enumerate(st.session_state.data or [])]
+
+    # Guard: no records at all
+    if not file_names:
+        st.error("No records loaded. Please upload a JSONL file on the Upload page.")
+        if st.button("← Back to Upload"):
+            st.session_state.page = 'upload'
+            st.rerun()
+    else:
+        # Build options (list is safer than range for length checks)
+        options = list(range(len(file_names)))
+
+        # Ensure edited_data exists and has consistent length
+        if not st.session_state.edited_data or len(st.session_state.edited_data) != len(file_names):
+            # try to sync edited_data to data
+            st.session_state.edited_data = (st.session_state.data or []).copy()
+
+        # Clamp current_index into valid range
+        cur_idx = st.session_state.get('current_index', 0)
+        try:
+            cur_idx = int(cur_idx)
+        except Exception:
+            cur_idx = 0
+        if cur_idx < 0:
+            cur_idx = 0
+        if cur_idx >= len(options):
+            cur_idx = len(options) - 1
+
+        # Show selectbox with a safe index
+        selected_file = st.selectbox(
+            "Select a file to view:",
+            options=options,
+            format_func=lambda x: f"{'✏️ ' if x in st.session_state.modified_indices else ''}{file_names[x]}",
+            index=cur_idx
+        )
+
+        # Persist chosen index
+        st.session_state.current_index = selected_file
+
+        # Safe access to the current record
+        current_record = st.session_state.edited_data[selected_file]
     
-    # Main layout: LHS (Image) and RHS (Details) - REDUCED GAP
-    left_col, right_col = st.columns([1.3, 1], gap="small")
+    left_col, right_col = st.columns([1.8, 1.0], gap="small")
     
     # LEFT SIDE: Image Display with OCR Canvas
     with left_col:
-        st.markdown("### 🖼️ Document Image")
-        
         file_name = current_record.get('file_name', '')
         
         if file_name:
-            st.caption(f"**File:** {file_name}")
-            
-            # Try to find matching image - CASE SENSITIVE EXACT MATCH ONLY
-            current_image = None
+            # Find the actual file name (handle cases where extension is missing)
+            actual_file_name = None
             if file_name in st.session_state.images:
-                current_image = st.session_state.images[file_name]
+                actual_file_name = file_name
+            else:
+                # Try adding common extensions
+                for ext in ['.pdf', '.png', '.jpg', '.jpeg', '.tiff', '.tif', '.bmp']:
+                    if file_name + ext in st.session_state.images:
+                        actual_file_name = file_name + ext
+                        break
+                
+                # Try matching without extension
+                if not actual_file_name:
+                    for uploaded_name in st.session_state.images.keys():
+                        uploaded_base = uploaded_name.rsplit('.', 1)[0]
+                        if uploaded_base == file_name:
+                            actual_file_name = uploaded_name
+                            break
+            
+            if actual_file_name:
+                # Check if this is a PDF with multiple pages
+                is_pdf = actual_file_name in st.session_state.pdf_metadata
+                
+                if is_pdf:
+                    pdf_meta = st.session_state.pdf_metadata[actual_file_name]
+                    total_pages = pdf_meta['total_pages']
+                    current_page = st.session_state.current_page_num.get(actual_file_name, 0)
+                    
+                    # PDF Navigation Header
+                    col_prev, col_info, col_next = st.columns([1, 2, 1])
+                    
+                    with col_prev:
+                        prev_clicked = st.button("⬅️ Previous", key=f"prev_page_{selected_file}_{actual_file_name}", 
+                                   disabled=(current_page == 0), use_container_width=True)
+                    
+                    with col_info:
+                        st.markdown(f"<div style='text-align: center; padding: 5px;'><b>📄 Page {current_page + 1} of {total_pages}</b></div>", unsafe_allow_html=True)
+                    
+                    with col_next:
+                        next_clicked = st.button("Next ➡️", key=f"next_page_{selected_file}_{actual_file_name}",
+                                   disabled=(current_page >= total_pages - 1), use_container_width=True)
+                    
+                    # Handle navigation only if not already navigating
+                    if not st.session_state.navigating_page:
+                        if prev_clicked:
+                            st.session_state.navigating_page = True
+                            st.session_state.current_page_num[actual_file_name] = max(0, current_page - 1)
+                            st.session_state.canvas_key += 1
+                            st.session_state.ocr_active_section = None
+                            st.session_state.ocr_active_field = None
+                            st.rerun()
+                        elif next_clicked:
+                            st.session_state.navigating_page = True
+                            st.session_state.current_page_num[actual_file_name] = min(total_pages - 1, current_page + 1)
+                            st.session_state.canvas_key += 1
+                            st.session_state.ocr_active_section = None
+                            st.session_state.ocr_active_field = None
+                            st.rerun()
+                    else:
+                        # Reset the flag after rerun
+                        st.session_state.navigating_page = False
+            
+            if actual_file_name:
+                # Determine if PDF and get the appropriate image
+                is_pdf = actual_file_name in st.session_state.pdf_metadata
+                
+                if is_pdf:
+                    # Get the current page image
+                    current_page = st.session_state.current_page_num.get(actual_file_name, 0)
+                    pdf_meta = st.session_state.pdf_metadata[actual_file_name]
+                    current_image = pdf_meta['pages'][current_page]
+                else:
+                    current_image = st.session_state.images[actual_file_name]
             else:
-                st.error(f"❌ Image '{file_name}' not found in uploaded images")
-                st.info("💡 Available images:")
-                with st.expander("Show available images"):
+                st.error(f"❌ File '{file_name}' not found in uploaded files")
+                st.info("💡 Available files:")
+                with st.expander("Show available files"):
                     for img_name in list(st.session_state.images.keys())[:20]:
                         st.text(f"  • {img_name}")
                     if len(st.session_state.images) > 20:
                         st.text(f"  ... and {len(st.session_state.images) - 20} more")
             
             if current_image:
-                # Scale image to fixed size
+                # Scale to a reasonable size so canvas doesn't become excessively large
                 scaled_image, scale_ratio, paste_x, paste_y = scale_image_to_fixed_size(current_image)
                 
-                # Always show canvas for drawing rectangles
+                # Render the canvas. Its internal canvas will be constrained by the wrapper due to CSS above.
                 canvas_result = st_canvas(
                     fill_color="rgba(255, 165, 0, 0.3)",
                     stroke_width=2,
@@ -426,14 +641,12 @@ elif st.session_state.page == 'viewer':
                     key=f"canvas_{selected_file}_{st.session_state.canvas_key}",
                 )
                 
-                # Process OCR when rectangle is drawn and field is selected
+                # Only attempt OCR if there's an active OCR target AND the user has drawn something (objects exist)
                 if canvas_result.json_data is not None and st.session_state.ocr_active_field:
-                    objects = canvas_result.json_data["objects"]
+                    objects = canvas_result.json_data.get("objects", [])
                     if len(objects) > 0:
-                        # Get the last drawn rectangle
                         rect = objects[-1]
                         
-                        # Adjust coordinates for padding and scale back to original image coordinates
                         bbox = [
                             (rect["left"] - paste_x) / scale_ratio,
                             (rect["top"] - paste_y) / scale_ratio,
@@ -441,25 +654,25 @@ elif st.session_state.page == 'viewer':
                             (rect["top"] + rect["height"] - paste_y) / scale_ratio
                         ]
                         
-                        # Perform OCR on original image
                         with st.spinner("Performing OCR..."):
                             ocr_text = perform_ocr(current_image, bbox)
                         
                         if ocr_text and not ocr_text.startswith("OCR Error"):
                             st.success(f"✅ OCR Result: {ocr_text}")
                             
-                            # Update the field value
                             gt_parse = st.session_state.edited_data[selected_file].get('gt_parse', {})
                             
                             if st.session_state.ocr_active_section == 'Line_items':
-                                # Handle line items
                                 line_items = gt_parse.get('Line_items', [])
                                 row_idx = st.session_state.ocr_line_item_row
                                 if row_idx is not None and row_idx < len(line_items):
                                     line_items[row_idx][st.session_state.ocr_active_field] = ocr_text
                                     gt_parse['Line_items'] = line_items
+
+                                    # ensure expander stays open for this row after OCR
+                                    expander_key = f"line_item_expander_{selected_file}_{row_idx}"
+                                    st.session_state[expander_key] = True
                             else:
-                                # Handle other sections
                                 section = st.session_state.ocr_active_section
                                 field = st.session_state.ocr_active_field
                                 if section not in gt_parse:
@@ -467,10 +680,15 @@ elif st.session_state.page == 'viewer':
                                 gt_parse[section][field] = ocr_text
                             
                             st.session_state.edited_data[selected_file]['gt_parse'] = gt_parse
+                            st.session_state.modified_indices.add(selected_file)
+                            
+                            # After successful OCR, deactivate OCR target to prevent repeated triggers/loops
+                            st.session_state.ocr_active_section = None
+                            st.session_state.ocr_active_field = None
+                            st.session_state.ocr_line_item_row = None
                             
-                            # Clear canvas and reset
+                            # Clear canvas for next OCR by bumping canvas_key then rerun
                             st.session_state.canvas_key += 1
-                            time.sleep(0.3)
                             st.rerun()
                         else:
                             st.error(ocr_text)
@@ -481,9 +699,8 @@ elif st.session_state.page == 'viewer':
     with right_col:
         st.markdown("### 📝 Document Details")
         
-        gt_parse = current_record.get('gt_parse', {})
+        gt_parse = st.session_state.edited_data[selected_file].get('gt_parse', {})
         
-        # Create tabs for each section
         tab1, tab2, tab3, tab4 = st.tabs([
             "📄 Remittance Details", 
             "👥 Party Details", 
@@ -493,340 +710,504 @@ elif st.session_state.page == 'viewer':
         
         # TAB 1: Remittance Details
         with tab1:
+            remittance = gt_parse.get('Remittance_details', {})
             
+            # Each field with OCR button
+            col_input, col_btn = st.columns([5, 1])
+            with col_input:
+                remittance['Remittance_adv_no'] = st.text_input(
+                    "Remittance Advice No", 
+                    value=remittance.get('Remittance_adv_no', ''),
+                    key=f"rem_adv_no_{selected_file}"
+                )
+            with col_btn:
+                st.markdown("<br>", unsafe_allow_html=True)
+                if st.button("🔍", key=f"ocr_rem_adv_no_{selected_file}", 
+                           type="primary" if is_ocr_active('Remittance_details', 'Remittance_adv_no') else "secondary"):
+                    activate_ocr_field('Remittance_details', 'Remittance_adv_no')
             
-            # OCR Field Selector
-            remittance_fields = [
-                'Select fields',
-                'Remittance_adv_no',
-                'Remittance_adv_date',
-                'Payment_method',
-                'FCY',
-                'Total_payment_amt_FCY',
-                'Payment_date',
-                'Payment_ref_no'
-            ]
-            
-            selected_rem_field = st.selectbox(
-                "🔍 Select field to populate via OCR:",
-                options=remittance_fields,
-                key=f"rem_ocr_select_{selected_file}"
-            )
+            col_input, col_btn = st.columns([5, 1])
+            with col_input:
+                remittance['Remittance_adv_date'] = st.text_input(
+                    "Remittance Advice Date", 
+                    value=remittance.get('Remittance_adv_date', ''),
+                    key=f"rem_adv_date_{selected_file}"
+                )
+            with col_btn:
+                st.markdown("<br>", unsafe_allow_html=True)
+                if st.button("🔍", key=f"ocr_rem_adv_date_{selected_file}", 
+                           type="primary" if is_ocr_active('Remittance_details', 'Remittance_adv_date') else "secondary"):
+                    activate_ocr_field('Remittance_details', 'Remittance_adv_date')
             
-            if selected_rem_field != 'Select fields':
-                st.session_state.ocr_active_section = 'Remittance_details'
-                st.session_state.ocr_active_field = selected_rem_field
-                st.session_state.ocr_line_item_row = None
-            else:
-                if st.session_state.ocr_active_section == 'Remittance_details':
-                    st.session_state.ocr_active_section = None
-                    st.session_state.ocr_active_field = None
+            col_input, col_btn = st.columns([5, 1])
+            with col_input:
+                remittance['Payment_method'] = st.text_input(
+                    "Payment Method", 
+                    value=remittance.get('Payment_method', ''),
+                    key=f"payment_method_{selected_file}"
+                )
+            with col_btn:
+                st.markdown("<br>", unsafe_allow_html=True)
+                if st.button("🔍", key=f"ocr_payment_method_{selected_file}", 
+                           type="primary" if is_ocr_active('Remittance_details', 'Payment_method') else "secondary"):
+                    activate_ocr_field('Remittance_details', 'Payment_method')
             
-            remittance = gt_parse.get('Remittance_details', {})
+            col_input, col_btn = st.columns([5, 1])
+            with col_input:
+                remittance['FCY'] = st.text_input(
+                    "FCY (Foreign Currency)", 
+                    value=remittance.get('FCY', ''),
+                    key=f"fcy_{selected_file}"
+                )
+            with col_btn:
+                st.markdown("<br>", unsafe_allow_html=True)
+                if st.button("🔍", key=f"ocr_fcy_{selected_file}", 
+                           type="primary" if is_ocr_active('Remittance_details', 'FCY') else "secondary"):
+                    activate_ocr_field('Remittance_details', 'FCY')
             
-            remittance['Remittance_adv_no'] = st.text_input(
-                "Remittance Advice No", 
-                value=remittance.get('Remittance_adv_no', ''),
-                key=f"rem_adv_no_{selected_file}"
-            )
-            remittance['Remittance_adv_date'] = st.text_input(
-                "Remittance Advice Date", 
-                value=remittance.get('Remittance_adv_date', ''),
-                key=f"rem_adv_date_{selected_file}"
-            )
-            remittance['Payment_method'] = st.text_input(
-                "Payment Method", 
-                value=remittance.get('Payment_method', ''),
-                key=f"payment_method_{selected_file}"
-            )
-            remittance['FCY'] = st.text_input(
-                "FCY (Foreign Currency)", 
-                value=remittance.get('FCY', ''),
-                key=f"fcy_{selected_file}"
-            )
-            remittance['Total_payment_amt_FCY'] = st.text_input(
-                "Total Payment Amount (FCY)", 
-                value=remittance.get('Total_payment_amt_FCY', ''),
-                key=f"total_payment_{selected_file}"
-            )
-            remittance['Payment_date'] = st.text_input(
-                "Payment Date", 
-                value=remittance.get('Payment_date', ''),
-                key=f"payment_date_{selected_file}"
-            )
-            remittance['Payment_ref_no'] = st.text_input(
-                "Payment Reference No", 
-                value=remittance.get('Payment_ref_no', ''),
-                key=f"payment_ref_{selected_file}"
-            )
+            col_input, col_btn = st.columns([5, 1])
+            with col_input:
+                remittance['Total_payment_amt_FCY'] = st.text_input(
+                    "Total Payment Amount (FCY)", 
+                    value=remittance.get('Total_payment_amt_FCY', ''),
+                    key=f"total_payment_{selected_file}"
+                )
+            with col_btn:
+                st.markdown("<br>", unsafe_allow_html=True)
+                if st.button("🔍", key=f"ocr_total_payment_{selected_file}", 
+                           type="primary" if is_ocr_active('Remittance_details', 'Total_payment_amt_FCY') else "secondary"):
+                    activate_ocr_field('Remittance_details', 'Total_payment_amt_FCY')
+            
+            col_input, col_btn = st.columns([5, 1])
+            with col_input:
+                remittance['Payment_date'] = st.text_input(
+                    "Payment Date", 
+                    value=remittance.get('Payment_date', ''),
+                    key=f"payment_date_{selected_file}"
+                )
+            with col_btn:
+                st.markdown("<br>", unsafe_allow_html=True)
+                if st.button("🔍", key=f"ocr_payment_date_{selected_file}", 
+                           type="primary" if is_ocr_active('Remittance_details', 'Payment_date') else "secondary"):
+                    activate_ocr_field('Remittance_details', 'Payment_date')
+            
+            col_input, col_btn = st.columns([5, 1])
+            with col_input:
+                remittance['Payment_ref_no'] = st.text_input(
+                    "Payment Reference No", 
+                    value=remittance.get('Payment_ref_no', ''),
+                    key=f"payment_ref_{selected_file}"
+                )
+            with col_btn:
+                st.markdown("<br>", unsafe_allow_html=True)
+                if st.button("🔍", key=f"ocr_payment_ref_{selected_file}", 
+                           type="primary" if is_ocr_active('Remittance_details', 'Payment_ref_no') else "secondary"):
+                    activate_ocr_field('Remittance_details', 'Payment_ref_no')
             
             gt_parse['Remittance_details'] = remittance
         
-        # TAB 2: Customer/Supplier Details
+        # TAB 2: Customer/Supplier Details with SWAP button
         with tab2:
+            # SWAP BUTTON - Centered and prominent
+            col1, col2, col3 = st.columns([1, 2, 1])
+            with col2:
+                if st.button("🔄 Swap Customer ↔ Supplier", key=f"swap_btn_{selected_file}", 
+                           type="primary", use_container_width=True):
+                    if not st.session_state.just_swapped:
+                        st.session_state.just_swapped = True
+                        swap_customer_supplier_details(selected_file)
+                        st.rerun()
             
-            
-            # OCR Field Selector
-            customer_fields = [
-                'Select fields',
-                'Customer_name',
-                'Customer_address',
-                'Customer_contact_info',
-                'Supplier_name',
-                'Supplier_address',
-                'Supplier_contact_info'
-            ]
-            
-            selected_cust_field = st.selectbox(
-                "🔍 Select field to populate via OCR:",
-                options=customer_fields,
-                key=f"cust_ocr_select_{selected_file}"
-            )
-            
-            if selected_cust_field != 'Select fields':
-                st.session_state.ocr_active_section = 'Customer_supplier_details'
-                st.session_state.ocr_active_field = selected_cust_field
-                st.session_state.ocr_line_item_row = None
-            else:
-                if st.session_state.ocr_active_section == 'Customer_supplier_details':
-                    st.session_state.ocr_active_section = None
-                    st.session_state.ocr_active_field = None
+            # Reset the flag after rerun
+            if st.session_state.just_swapped:
+                st.session_state.just_swapped = False
             
             st.markdown("**Customer Details**")
             customer_supplier = gt_parse.get('Customer_supplier_details', {})
             
-            customer_supplier['Customer_name'] = st.text_input(
-                "Customer Name", 
-                value=customer_supplier.get('Customer_name', ''),
-                key=f"cust_name_{selected_file}"
-            )
-            customer_supplier['Customer_address'] = st.text_area(
-                "Customer Address", 
-                value=customer_supplier.get('Customer_address', ''),
-                key=f"cust_addr_{selected_file}",
-                height=60
-            )
-            customer_supplier['Customer_contact_info'] = st.text_input(
-                "Customer Contact Info", 
-                value=customer_supplier.get('Customer_contact_info', ''),
-                key=f"cust_contact_{selected_file}"
-            )
+            col_input, col_btn = st.columns([5, 1])
+            with col_input:
+                customer_supplier['Customer_name'] = st.text_input(
+                    "Customer Name", 
+                    value=customer_supplier.get('Customer_name', ''),
+                    key=f"cust_name_{selected_file}"
+                )
+            with col_btn:
+                st.markdown("<br>", unsafe_allow_html=True)
+                if st.button("🔍", key=f"ocr_cust_name_{selected_file}", 
+                           type="primary" if is_ocr_active('Customer_supplier_details', 'Customer_name') else "secondary"):
+                    activate_ocr_field('Customer_supplier_details', 'Customer_name')
+            
+            col_input, col_btn = st.columns([5, 1])
+            with col_input:
+                customer_supplier['Customer_address'] = st.text_area(
+                    "Customer Address", 
+                    value=customer_supplier.get('Customer_address', ''),
+                    key=f"cust_addr_{selected_file}",
+                    height=60
+                )
+            with col_btn:
+                st.markdown("<br>", unsafe_allow_html=True)
+                if st.button("🔍", key=f"ocr_cust_addr_{selected_file}", 
+                           type="primary" if is_ocr_active('Customer_supplier_details', 'Customer_address') else "secondary"):
+                    activate_ocr_field('Customer_supplier_details', 'Customer_address')
+            
+            col_input, col_btn = st.columns([5, 1])
+            with col_input:
+                customer_supplier['Customer_contact_info'] = st.text_input(
+                    "Customer Contact Info", 
+                    value=customer_supplier.get('Customer_contact_info', ''),
+                    key=f"cust_contact_{selected_file}"
+                )
+            with col_btn:
+                st.markdown("<br>", unsafe_allow_html=True)
+                if st.button("🔍", key=f"ocr_cust_contact_{selected_file}", 
+                           type="primary" if is_ocr_active('Customer_supplier_details', 'Customer_contact_info') else "secondary"):
+                    activate_ocr_field('Customer_supplier_details', 'Customer_contact_info')
             
             st.markdown("**Supplier Details**")
-            customer_supplier['Supplier_name'] = st.text_input(
-                "Supplier Name", 
-                value=customer_supplier.get('Supplier_name', ''),
-                key=f"supp_name_{selected_file}"
-            )
-            customer_supplier['Supplier_address'] = st.text_area(
-                "Supplier Address", 
-                value=customer_supplier.get('Supplier_address', ''),
-                key=f"supp_addr_{selected_file}",
-                height=60
-            )
-            customer_supplier['Supplier_contact_info'] = st.text_input(
-                "Supplier Contact Info", 
-                value=customer_supplier.get('Supplier_contact_info', ''),
-                key=f"supp_contact_{selected_file}"
-            )
+            
+            col_input, col_btn = st.columns([5, 1])
+            with col_input:
+                customer_supplier['Supplier_name'] = st.text_input(
+                    "Supplier Name", 
+                    value=customer_supplier.get('Supplier_name', ''),
+                    key=f"supp_name_{selected_file}"
+                )
+            with col_btn:
+                st.markdown("<br>", unsafe_allow_html=True)
+                if st.button("🔍", key=f"ocr_supp_name_{selected_file}", 
+                           type="primary" if is_ocr_active('Customer_supplier_details', 'Supplier_name') else "secondary"):
+                    activate_ocr_field('Customer_supplier_details', 'Supplier_name')
+            
+            col_input, col_btn = st.columns([5, 1])
+            with col_input:
+                customer_supplier['Supplier_address'] = st.text_area(
+                    "Supplier Address", 
+                    value=customer_supplier.get('Supplier_address', ''),
+                    key=f"supp_addr_{selected_file}",
+                    height=60
+                )
+            with col_btn:
+                st.markdown("<br>", unsafe_allow_html=True)
+                if st.button("🔍", key=f"ocr_supp_addr_{selected_file}", 
+                           type="primary" if is_ocr_active('Customer_supplier_details', 'Supplier_address') else "secondary"):
+                    activate_ocr_field('Customer_supplier_details', 'Supplier_address')
+            
+            col_input, col_btn = st.columns([5, 1])
+            with col_input:
+                customer_supplier['Supplier_contact_info'] = st.text_input(
+                    "Supplier Contact Info", 
+                    value=customer_supplier.get('Supplier_contact_info', ''),
+                    key=f"supp_contact_{selected_file}"
+                )
+            with col_btn:
+                st.markdown("<br>", unsafe_allow_html=True)
+                if st.button("🔍", key=f"ocr_supp_contact_{selected_file}", 
+                           type="primary" if is_ocr_active('Customer_supplier_details', 'Supplier_contact_info') else "secondary"):
+                    activate_ocr_field('Customer_supplier_details', 'Supplier_contact_info')
             
             gt_parse['Customer_supplier_details'] = customer_supplier
         
         # TAB 3: Bank Details
         with tab3:
+            bank = gt_parse.get('Bank_details', {})
             
+            col_input, col_btn = st.columns([5, 1])
+            with col_input:
+                bank['Bank_name'] = st.text_input(
+                    "Bank Name", 
+                    value=bank.get('Bank_name', ''),
+                    key=f"bank_name_{selected_file}"
+                )
+            with col_btn:
+                st.markdown("<br>", unsafe_allow_html=True)
+                if st.button("🔍", key=f"ocr_bank_name_{selected_file}", 
+                           type="primary" if is_ocr_active('Bank_details', 'Bank_name') else "secondary"):
+                    activate_ocr_field('Bank_details', 'Bank_name')
             
-            # OCR Field Selector
-            bank_fields = [
-                'Select fields',
-                'Bank_name',
-                'Bank_acc_no',
-                'Bank_routing_no',
-                'Swift_code'
-            ]
-            
-            selected_bank_field = st.selectbox(
-                "🔍 Select field to populate via OCR:",
-                options=bank_fields,
-                key=f"bank_ocr_select_{selected_file}"
-            )
-            
-            if selected_bank_field != 'Select fields':
-                st.session_state.ocr_active_section = 'Bank_details'
-                st.session_state.ocr_active_field = selected_bank_field
-                st.session_state.ocr_line_item_row = None
-            else:
-                if st.session_state.ocr_active_section == 'Bank_details':
-                    st.session_state.ocr_active_section = None
-                    st.session_state.ocr_active_field = None
+            col_input, col_btn = st.columns([5, 1])
+            with col_input:
+                bank['Bank_acc_no'] = st.text_input(
+                    "Bank Account No", 
+                    value=bank.get('Bank_acc_no', ''),
+                    key=f"bank_acc_{selected_file}"
+                )
+            with col_btn:
+                st.markdown("<br>", unsafe_allow_html=True)
+                if st.button("🔍", key=f"ocr_bank_acc_{selected_file}", 
+                           type="primary" if is_ocr_active('Bank_details', 'Bank_acc_no') else "secondary"):
+                    activate_ocr_field('Bank_details', 'Bank_acc_no')
             
-            bank = gt_parse.get('Bank_details', {})
+            col_input, col_btn = st.columns([5, 1])
+            with col_input:
+                bank['Bank_routing_no'] = st.text_input(
+                    "Bank Routing No", 
+                    value=bank.get('Bank_routing_no', ''),
+                    key=f"bank_routing_{selected_file}"
+                )
+            with col_btn:
+                st.markdown("<br>", unsafe_allow_html=True)
+                if st.button("🔍", key=f"ocr_bank_routing_{selected_file}", 
+                           type="primary" if is_ocr_active('Bank_details', 'Bank_routing_no') else "secondary"):
+                    activate_ocr_field('Bank_details', 'Bank_routing_no')
             
-            bank['Bank_name'] = st.text_input(
-                "Bank Name", 
-                value=bank.get('Bank_name', ''),
-                key=f"bank_name_{selected_file}"
-            )
-            bank['Bank_acc_no'] = st.text_input(
-                "Bank Account No", 
-                value=bank.get('Bank_acc_no', ''),
-                key=f"bank_acc_{selected_file}"
-            )
-            bank['Bank_routing_no'] = st.text_input(
-                "Bank Routing No", 
-                value=bank.get('Bank_routing_no', ''),
-                key=f"bank_routing_{selected_file}"
-            )
-            bank['Swift_code'] = st.text_input(
-                "SWIFT Code", 
-                value=bank.get('Swift_code', ''),
-                key=f"swift_{selected_file}"
-            )
+            col_input, col_btn = st.columns([5, 1])
+            with col_input:
+                bank['Swift_code'] = st.text_input(
+                    "SWIFT Code", 
+                    value=bank.get('Swift_code', ''),
+                    key=f"swift_{selected_file}"
+                )
+            with col_btn:
+                st.markdown("<br>", unsafe_allow_html=True)
+                if st.button("🔍", key=f"ocr_swift_{selected_file}", 
+                           type="primary" if is_ocr_active('Bank_details', 'Swift_code') else "secondary"):
+                    activate_ocr_field('Bank_details', 'Swift_code')
             
             gt_parse['Bank_details'] = bank
         
         # TAB 4: Line Items
         with tab4:
+            current_gt_parse = st.session_state.edited_data[selected_file].get('gt_parse', {})
+            line_items = current_gt_parse.get('Line_items', [])
             
-            
-            # OCR Controls for Line Items - Fixed layout
-            line_items = gt_parse.get('Line_items', [])
-            
-            # Adjusted column widths - all controls in single compact line
-            col_field, col_row, col_add, col_remove = st.columns([1.5, 0.7, 0.30, 0.30])
-            
-            line_item_fields = [
-                'Select fields',
-                'Po_number',
-                'Invoice_no',
-                'Other_doc_ref_no',
-                'Invoice_date',
-                'Invoice_amount_FCY',
-                'Amount_paid_for_each_invoice',
-                'Outstanding_balance_FCY',
-                'Discounts_taken_FCY',
-                'Adjustments(without_holding_tax)_FCY',
-                'Descriptions'
-            ]
-            
-            with col_field:
-                selected_line_field = st.selectbox(
-                    "🔍 Field:",
-                    options=line_item_fields,
-                    key=f"line_ocr_field_{selected_file}"
-                )
-            
-            with col_row:
-                if len(line_items) > 0:
-                    selected_row = st.selectbox(
-                        "Row:",
-                        options=list(range(len(line_items))),
-                        format_func=lambda x: f"Row {x + 1}",
-                        key=f"line_ocr_row_{selected_file}"
-                    )
-                else:
-                    st.selectbox("Row:", options=[], disabled=True, key=f"line_ocr_row_empty_{selected_file}")
-                    selected_row = None
-            
+            # Add/Remove row buttons
+            col_add, col_remove = st.columns([1, 1])
             with col_add:
-                # Use button with on_click callback to prevent loop
-                if st.button("➕", key=f"add_row_{selected_file}", help="Add new row"):
+                if st.button("➕ Add New Row", key=f"add_row_{selected_file}", use_container_width=True):
                     if not st.session_state.button_clicked:
                         st.session_state.button_clicked = True
                         new_row = {
-                            "Po_number": "",
-                            "Invoice_no": "",
-                            "Other_doc_ref_no": "",
-                            "Invoice_date": "",
-                            "Invoice_amount_FCY": "",
-                            "Amount_paid_for_each_invoice": "",
-                            "Outstanding_balance_FCY": "",
-                            "Discounts_taken_FCY": "",
-                            "Adjustments(without_holding_tax)_FCY": "",
+                            "Po_number": "", "Invoice_no": "", "Other_doc_ref_no": "",
+                            "Invoice_date": "", "Invoice_amount_FCY": "",
+                            "Amount_paid_for_each_invoice": "", "Outstanding_balance_FCY": "",
+                            "Discounts_taken_FCY": "", "Adjustments(without_holding_tax)_FCY": "",
                             "Descriptions": ""
                         }
-                        line_items.append(new_row)
-                        gt_parse['Line_items'] = line_items
-                        st.session_state.edited_data[selected_file]['gt_parse'] = gt_parse
+                        current_gt_parse = st.session_state.edited_data[selected_file].get('gt_parse', {})
+                        current_line_items = current_gt_parse.get('Line_items', [])
+                        current_line_items.append(new_row)
+                        current_gt_parse['Line_items'] = current_line_items
+                        st.session_state.edited_data[selected_file]['gt_parse'] = current_gt_parse
                         st.session_state.modified_indices.add(selected_file)
+
+                        # Ensure the newly added row's expander is open
+                        new_idx = len(current_line_items) - 1
+                        expander_key_new = f"line_item_expander_{selected_file}_{new_idx}"
+                        st.session_state[expander_key_new] = True
+
                         st.rerun()
             
             with col_remove:
-                if st.button("➖", key=f"remove_row_{selected_file}", help="Remove selected row", disabled=(len(line_items) == 0)):
-                    if not st.session_state.button_clicked and len(line_items) > 0 and selected_row is not None:
+                if st.button("➖ Remove Last Row", key=f"remove_row_{selected_file}", 
+                           disabled=(len(line_items) == 0), use_container_width=True):
+                    if not st.session_state.button_clicked and len(line_items) > 0:
                         st.session_state.button_clicked = True
-                        line_items.pop(selected_row)
-                        gt_parse['Line_items'] = line_items
-                        st.session_state.edited_data[selected_file]['gt_parse'] = gt_parse
+                        current_gt_parse = st.session_state.edited_data[selected_file].get('gt_parse', {})
+                        current_line_items = current_gt_parse.get('Line_items', [])
+                        N = len(current_line_items)
+                        current_line_items.pop()
+                        current_gt_parse['Line_items'] = current_line_items
+                        st.session_state.edited_data[selected_file]['gt_parse'] = current_gt_parse
                         st.session_state.modified_indices.add(selected_file)
+
+                        # Remove the expander flag for the popped row (if present)
+                        popped_idx = N - 1
+                        expander_key_popped = f"line_item_expander_{selected_file}_{popped_idx}"
+                        if expander_key_popped in st.session_state:
+                            del st.session_state[expander_key_popped]
+
                         st.rerun()
             
-            # Reset button clicked flag after processing
             if st.session_state.button_clicked:
                 st.session_state.button_clicked = False
             
-            # Set OCR state for line items
-            if selected_line_field != 'Select fields' and selected_row is not None:
-                st.session_state.ocr_active_section = 'Line_items'
-                st.session_state.ocr_active_field = selected_line_field
-                st.session_state.ocr_line_item_row = selected_row
-            else:
-                if st.session_state.ocr_active_section == 'Line_items':
-                    st.session_state.ocr_active_section = None
-                    st.session_state.ocr_active_field = None
-                    st.session_state.ocr_line_item_row = None
-            
+            # Display each row as an expander with OCR buttons
+            current_gt_parse = st.session_state.edited_data[selected_file].get('gt_parse', {})
+            line_items = current_gt_parse.get('Line_items', [])
             
-            
-            # Display line items table
             if line_items:
-                df = pd.DataFrame(line_items)
+                for idx, item in enumerate(line_items):
+                    # Use a persistent session_state flag so expansion state is preserved across reruns.
+                    expander_key = f"line_item_expander_{selected_file}_{idx}"
+                    expanded_default = st.session_state.get(expander_key, False)
+
+                    # Note: do NOT pass a 'key' arg to st.expander to maintain compatibility; control expanded via session_state flag.
+                    with st.expander(f"**Row {idx + 1}** - Invoice: {item.get('Invoice_no', 'N/A')}", expanded=expanded_default):
+                        # PO Number
+                        col_input, col_btn = st.columns([5, 1])
+                        with col_input:
+                            item['Po_number'] = st.text_input(
+                                "PO Number", 
+                                value=item.get('Po_number', ''),
+                                key=f"po_num_{selected_file}_{idx}"
+                            )
+                        with col_btn:
+                            st.markdown("<br>", unsafe_allow_html=True)
+                            if st.button("🔍", key=f"ocr_po_{selected_file}_{idx}", 
+                                       type="primary" if is_ocr_active('Line_items', 'Po_number', idx) else "secondary"):
+                                # ensure expander stays open when user explicitly requests OCR
+                                st.session_state[expander_key] = True
+                                activate_ocr_field('Line_items', 'Po_number', idx)
+                        
+                        # Invoice No
+                        col_input, col_btn = st.columns([5, 1])
+                        with col_input:
+                            item['Invoice_no'] = st.text_input(
+                                "Invoice No", 
+                                value=item.get('Invoice_no', ''),
+                                key=f"inv_no_{selected_file}_{idx}"
+                            )
+                        with col_btn:
+                            st.markdown("<br>", unsafe_allow_html=True)
+                            if st.button("🔍", key=f"ocr_inv_{selected_file}_{idx}", 
+                                       type="primary" if is_ocr_active('Line_items', 'Invoice_no', idx) else "secondary"):
+                                st.session_state[expander_key] = True
+                                activate_ocr_field('Line_items', 'Invoice_no', idx)
+                        
+                        # Other Doc Ref No
+                        col_input, col_btn = st.columns([5, 1])
+                        with col_input:
+                            item['Other_doc_ref_no'] = st.text_input(
+                                "Other Doc Ref No", 
+                                value=item.get('Other_doc_ref_no', ''),
+                                key=f"other_doc_{selected_file}_{idx}"
+                            )
+                        with col_btn:
+                            st.markdown("<br>", unsafe_allow_html=True)
+                            if st.button("🔍", key=f"ocr_other_{selected_file}_{idx}", 
+                                       type="primary" if is_ocr_active('Line_items', 'Other_doc_ref_no', idx) else "secondary"):
+                                st.session_state[expander_key] = True
+                                activate_ocr_field('Line_items', 'Other_doc_ref_no', idx)
+                        
+                        # Invoice Date
+                        col_input, col_btn = st.columns([5, 1])
+                        with col_input:
+                            item['Invoice_date'] = st.text_input(
+                                "Invoice Date", 
+                                value=item.get('Invoice_date', ''),
+                                key=f"inv_date_{selected_file}_{idx}"
+                            )
+                        with col_btn:
+                            st.markdown("<br>", unsafe_allow_html=True)
+                            if st.button("🔍", key=f"ocr_inv_date_{selected_file}_{idx}", 
+                                       type="primary" if is_ocr_active('Line_items', 'Invoice_date', idx) else "secondary"):
+                                st.session_state[expander_key] = True
+                                activate_ocr_field('Line_items', 'Invoice_date', idx)
+                        
+                        # Invoice Amount FCY
+                        col_input, col_btn = st.columns([5, 1])
+                        with col_input:
+                            item['Invoice_amount_FCY'] = st.text_input(
+                                "Invoice Amount FCY", 
+                                value=item.get('Invoice_amount_FCY', ''),
+                                key=f"inv_amt_{selected_file}_{idx}"
+                            )
+                        with col_btn:
+                            st.markdown("<br>", unsafe_allow_html=True)
+                            if st.button("🔍", key=f"ocr_inv_amt_{selected_file}_{idx}", 
+                                       type="primary" if is_ocr_active('Line_items', 'Invoice_amount_FCY', idx) else "secondary"):
+                                st.session_state[expander_key] = True
+                                activate_ocr_field('Line_items', 'Invoice_amount_FCY', idx)
+                        
+                        # Amount Paid
+                        col_input, col_btn = st.columns([5, 1])
+                        with col_input:
+                            item['Amount_paid_for_each_invoice'] = st.text_input(
+                                "Amount Paid", 
+                                value=item.get('Amount_paid_for_each_invoice', ''),
+                                key=f"amt_paid_{selected_file}_{idx}"
+                            )
+                        with col_btn:
+                            st.markdown("<br>", unsafe_allow_html=True)
+                            if st.button("🔍", key=f"ocr_amt_paid_{selected_file}_{idx}", 
+                                       type="primary" if is_ocr_active('Line_items', 'Amount_paid_for_each_invoice', idx) else "secondary"):
+                                st.session_state[expander_key] = True
+                                activate_ocr_field('Line_items', 'Amount_paid_for_each_invoice', idx)
+                        
+                        # Outstanding Balance
+                        col_input, col_btn = st.columns([5, 1])
+                        with col_input:
+                            item['Outstanding_balance_FCY'] = st.text_input(
+                                "Outstanding Balance FCY", 
+                                value=item.get('Outstanding_balance_FCY', ''),
+                                key=f"out_bal_{selected_file}_{idx}"
+                            )
+                        with col_btn:
+                            st.markdown("<br>", unsafe_allow_html=True)
+                            if st.button("🔍", key=f"ocr_out_bal_{selected_file}_{idx}", 
+                                       type="primary" if is_ocr_active('Line_items', 'Outstanding_balance_FCY', idx) else "secondary"):
+                                st.session_state[expander_key] = True
+                                activate_ocr_field('Line_items', 'Outstanding_balance_FCY', idx)
+                        
+                        # Discounts
+                        col_input, col_btn = st.columns([5, 1])
+                        with col_input:
+                            item['Discounts_taken_FCY'] = st.text_input(
+                                "Discounts Taken FCY", 
+                                value=item.get('Discounts_taken_FCY', ''),
+                                key=f"disc_{selected_file}_{idx}"
+                            )
+                        with col_btn:
+                            st.markdown("<br>", unsafe_allow_html=True)
+                            if st.button("🔍", key=f"ocr_disc_{selected_file}_{idx}", 
+                                       type="primary" if is_ocr_active('Line_items', 'Discounts_taken_FCY', idx) else "secondary"):
+                                st.session_state[expander_key] = True
+                                activate_ocr_field('Line_items', 'Discounts_taken_FCY', idx)
+                        
+                        # Adjustments
+                        col_input, col_btn = st.columns([5, 1])
+                        with col_input:
+                            item['Adjustments(without_holding_tax)_FCY'] = st.text_input(
+                                "Adjustments FCY", 
+                                value=item.get('Adjustments(without_holding_tax)_FCY', ''),
+                                key=f"adj_{selected_file}_{idx}"
+                            )
+                        with col_btn:
+                            st.markdown("<br>", unsafe_allow_html=True)
+                            if st.button("🔍", key=f"ocr_adj_{selected_file}_{idx}", 
+                                       type="primary" if is_ocr_active('Line_items', 'Adjustments(without_holding_tax)_FCY', idx) else "secondary"):
+                                st.session_state[expander_key] = True
+                                activate_ocr_field('Line_items', 'Adjustments(without_holding_tax)_FCY', idx)
+                        
+                        # Descriptions
+                        col_input, col_btn = st.columns([5, 1])
+                        with col_input:
+                            item['Descriptions'] = st.text_area(
+                                "Descriptions", 
+                                value=item.get('Descriptions', ''),
+                                key=f"desc_{selected_file}_{idx}",
+                                height=60
+                            )
+                        with col_btn:
+                            st.markdown("<br>", unsafe_allow_html=True)
+                            if st.button("🔍", key=f"ocr_desc_{selected_file}_{idx}", 
+                                       type="primary" if is_ocr_active('Line_items', 'Descriptions', idx) else "secondary"):
+                                st.session_state[expander_key] = True
+                                activate_ocr_field('Line_items', 'Descriptions', idx)
                 
-                # Convert amount fields to numeric
-                amount_fields = ['Invoice_amount_FCY', 'Amount_paid_for_each_invoice', 
-                               'Outstanding_balance_FCY', 'Discounts_taken_FCY', 
-                               'Adjustments(without_holding_tax)_FCY']
+                # Update line items back to gt_parse
+                current_gt_parse['Line_items'] = line_items
                 
-                for field in amount_fields:
-                    if field in df.columns:
-                        df[field] = pd.to_numeric(df[field].replace('', None), errors='coerce')
+                st.markdown("**📊 Line Items Summary Table**")
                 
-                column_config = {
-                    "Po_number": st.column_config.TextColumn("PO Number", width="small"),
-                    "Invoice_no": st.column_config.TextColumn("Invoice No", width="small"),
-                    "Other_doc_ref_no": st.column_config.TextColumn("Other Doc Ref No", width="small"),
-                    "Invoice_date": st.column_config.TextColumn("Invoice Date", width="small"),
-                    "Invoice_amount_FCY": st.column_config.NumberColumn("Invoice Amt FCY", width="small", format="%.2f"),
-                    "Amount_paid_for_each_invoice": st.column_config.NumberColumn("Amount Paid", width="small", format="%.2f"),
-                    "Outstanding_balance_FCY": st.column_config.NumberColumn("Outstanding FCY", width="small", format="%.2f"),
-                    "Discounts_taken_FCY": st.column_config.NumberColumn("Discounts FCY", width="small", format="%.2f"),
-                    "Adjustments(without_holding_tax)_FCY": st.column_config.NumberColumn("Adjustments FCY", width="small", format="%.2f"),
-                    "Descriptions": st.column_config.TextColumn("Descriptions", width="medium"),
-                }
+                # Display summary table with index starting from 1
+                df = pd.DataFrame(line_items)
+                df.index = df.index + 1  # Start index from 1
+                df.index.name = 'SL No'
                 
-                edited_df = st.data_editor(
+                st.dataframe(
                     df,
-                    column_config=column_config,
-                    num_rows="fixed",
                     use_container_width=True,
-                    key=f"line_items_table_{selected_file}",
-                    hide_index=False
+                    height=300
                 )
-                
-                # Convert back to string
-                for field in amount_fields:
-                    if field in edited_df.columns:
-                        edited_df[field] = edited_df[field].apply(lambda x: str(x) if pd.notna(x) else '')
-                
-                gt_parse['Line_items'] = edited_df.to_dict('records')
             else:
-                st.info("No line items. Click ➕ to add a new row.")
+                st.info("No line items. Click '➕ Add New Row' to add a new row.")
         
-        # Update the edited data
         st.session_state.edited_data[selected_file]['gt_parse'] = gt_parse
         
         # Save button
-        st.markdown("---")
         col1, col2 = st.columns([1, 1])
         with col1:
             if st.button("💾 Save Changes", type="primary", use_container_width=True, key=f"save_btn_{selected_file}"):
@@ -837,10 +1218,8 @@ elif st.session_state.page == 'viewer':
                     st.session_state.save_message_time = time.time()
                     st.rerun()
         
-        # Reset the just_saved flag after rerun
         if st.session_state.just_saved:
             st.session_state.just_saved = False
         
-        # Display save message under the button (appears after rerun)
         if st.session_state.save_message:
-            st.success(st.session_state.save_message)
+            st.success(st.session_state.save_message)