Update extract_red_text.py

extract_red_text.py

@@ -7,23 +7,25 @@ from docx.oxml.ns import qn
 from master_key import TABLE_SCHEMAS, HEADING_PATTERNS, PARAGRAPH_PATTERNS
 
 def is_red_font(run):
+    """Enhanced red font detection with better color checking"""
     col = run.font.color
     if col and col.rgb:
         r, g, b = col.rgb
-        if r>150 and g<100 and b<100 and (r-g)>30 and (r-b)>30:
+        if r > 150 and g < 100 and b < 100 and (r-g) > 30 and (r-b) > 30:
             return True
     rPr = getattr(run._element, "rPr", None)
     if rPr is not None:
         clr = rPr.find(qn('w:color'))
         if clr is not None:
             val = clr.get(qn('w:val'))
-            if re.fullmatch(r"[0-9A-Fa-f]{6}", val):
-                rr, gg, bb = int(val[:2],16), int(val[2:4],16), int(val[4:],16)
-                if rr>150 and gg<100 and bb<100 and (rr-gg)>30 and (rr-bb)>30:
+            if val and re.fullmatch(r"[0-9A-Fa-f]{6}", val):
+                rr, gg, bb = int(val[:2], 16), int(val[2:4], 16), int(val[4:], 16)
+                if rr > 150 and gg < 100 and bb < 100 and (rr-gg) > 30 and (rr-bb) > 30:
                     return True
     return False
 
 def _prev_para_text(tbl):
+    """Get text from previous paragraph before table"""
     prev = tbl._tbl.getprevious()
     while prev is not None and not prev.tag.endswith("}p"):
         prev = prev.getprevious()
@@ -31,127 +33,427 @@ def _prev_para_text(tbl):
         return ""
     return "".join(node.text for node in prev.iter() if node.tag.endswith("}t") and node.text).strip()
 
-def match_table_schema(tbl):
-    # look for explicit heading constraint
-    heading = _prev_para_text(tbl)
-    headers = [c.text.strip() for c in tbl.rows[0].cells]
-    col0    = [r.cells[0].text.strip() for r in tbl.rows]
-
-    # 1) exact first-cell name
-    first = tbl.rows[0].cells[0].text.strip()
-    if first in TABLE_SCHEMAS:
-        spec = TABLE_SCHEMAS[first]
-        if not spec.get("headings") or any(h["text"]==heading for h in spec.get("headings",[])):
-            return first
+def normalize_text(text):
+    """Normalize text for better matching"""
+    return re.sub(r'\s+', ' ', text.strip())
 
-    # 2) any other schema with explicit headings
-    for name, spec in TABLE_SCHEMAS.items():
-        if any(h["text"]==heading for h in spec.get("headings",[])):
-            return name
+def fuzzy_match_heading(heading, patterns):
+    """Check if heading matches any pattern with fuzzy matching"""
+    heading_norm = normalize_text(heading.upper())
+    for pattern in patterns:
+        if re.search(pattern, heading_norm, re.IGNORECASE):
+            return True
+    return False
 
-    # 3) by two-column 'columns'
-    for name, spec in TABLE_SCHEMAS.items():
-        cols = spec.get("columns")
-        if cols and all(col in headers for col in cols):
-            return name
+def get_table_context(tbl):
+    """Get comprehensive context information for table"""
+    heading = normalize_text(_prev_para_text(tbl))
+    headers = [normalize_text(c.text) for c in tbl.rows[0].cells if c.text.strip()]
+    col0 = [normalize_text(r.cells[0].text) for r in tbl.rows if r.cells[0].text.strip()]
+    first_cell = normalize_text(tbl.rows[0].cells[0].text) if tbl.rows else ""
+    all_cells = []
+    for row in tbl.rows:
+        for cell in row.cells:
+            text = normalize_text(cell.text)
+            if text:
+                all_cells.append(text)
+    return {
+        'heading': heading,
+        'headers': headers,
+        'col0': col0,
+        'first_cell': first_cell,
+        'all_cells': all_cells,
+        'num_rows': len(tbl.rows),
+        'num_cols': len(tbl.rows[0].cells) if tbl.rows else 0
+    }
 
-    # 4) row1 tables
-    for name, spec in TABLE_SCHEMAS.items():
-        if spec["orientation"]=="row1" and all(lbl in headers for lbl in spec["labels"]):
-            return name
+def calculate_schema_match_score(schema_name, spec, context):
+    """Enhanced calculate match score - IMPROVED for Vehicle Registration tables"""
+    score = 0
+    reasons = []
+    
+    # 🎯 VEHICLE REGISTRATION BOOST
+    if "Vehicle Registration" in schema_name:
+        vehicle_keywords = ["registration", "vehicle", "sub-contractor", "weight verification", "rfs suspension"]
+        table_text = " ".join(context['headers']).lower() + " " + context['heading'].lower()
+        
+        keyword_matches = sum(1 for keyword in vehicle_keywords if keyword in table_text)
+        if keyword_matches >= 2:
+            score += 150  # Very high boost for vehicle tables
+            reasons.append(f"Vehicle Registration keywords: {keyword_matches}/5")
+        elif keyword_matches >= 1:
+            score += 75   # Medium boost
+            reasons.append(f"Some Vehicle Registration keywords: {keyword_matches}/5")
+    
+    # 🎯 SUMMARY TABLE BOOST (existing logic)
+    if "Summary" in schema_name and "details" in " ".join(context['headers']).lower():
+        score += 100
+        reasons.append(f"Summary schema with DETAILS column - perfect match")
+    
+    if "Summary" not in schema_name and "details" in " ".join(context['headers']).lower():
+        score -= 75
+        reasons.append(f"Non-summary schema penalized for DETAILS column presence")
+    
+    # Context exclusions
+    if spec.get("context_exclusions"):
+        table_text = " ".join(context['headers']).lower() + " " + context['heading'].lower()
+        for exclusion in spec["context_exclusions"]:
+            if exclusion.lower() in table_text:
+                score -= 50
+                reasons.append(f"Context exclusion penalty: '{exclusion}' found")
+    
+    # Context keywords
+    if spec.get("context_keywords"):
+        table_text = " ".join(context['headers']).lower() + " " + context['heading'].lower()
+        keyword_matches = 0
+        for keyword in spec["context_keywords"]:
+            if keyword.lower() in table_text:
+                keyword_matches += 1
+        
+        if keyword_matches > 0:
+            score += keyword_matches * 15
+            reasons.append(f"Context keyword matches: {keyword_matches}/{len(spec['context_keywords'])}")
+    
+    # Direct first cell match
+    if context['first_cell'] and context['first_cell'].upper() == schema_name.upper():
+        score += 100
+        reasons.append(f"Direct first cell match: '{context['first_cell']}'")
+    
+    # Heading pattern matching
+    if spec.get("headings"):
+        for h in spec["headings"]:
+            if fuzzy_match_heading(context['heading'], [h["text"]]):
+                score += 50
+                reasons.append(f"Heading match: '{context['heading']}'")
+                break
+    
+    # Column header matching
+    if spec.get("columns"):
+        cols = [normalize_text(col) for col in spec["columns"]]
+        matches = 0
+        for col in cols:
+            if any(col.upper() in h.upper() for h in context['headers']):
+                matches += 1
+        if matches == len(cols):
+            score += 60
+            reasons.append(f"All column headers match: {cols}")
+        elif matches > 0:
+            score += matches * 20
+            reasons.append(f"Partial column matches: {matches}/{len(cols)}")
+    
+    # Label matching for left-oriented tables
+    if spec.get("orientation") == "left":
+        labels = [normalize_text(lbl) for lbl in spec["labels"]]
+        matches = 0
+        for lbl in labels:
+            if any(lbl.upper() in c.upper() or c.upper() in lbl.upper() for c in context['col0']):
+                matches += 1
+        if matches > 0:
+            score += (matches / len(labels)) * 30
+            reasons.append(f"Left orientation label matches: {matches}/{len(labels)}")
+    
+    # 🎯 ENHANCED Label matching for row1-oriented tables (Vehicle Registration)
+    elif spec.get("orientation") == "row1":
+        labels = [normalize_text(lbl) for lbl in spec["labels"]]
+        matches = 0
+        for lbl in labels:
+            # More flexible matching for vehicle tables
+            if any(lbl.upper() in h.upper() or h.upper() in lbl.upper() for h in context['headers']):
+                matches += 1
+            # Also check for partial keyword matches
+            elif any(word.upper() in " ".join(context['headers']).upper() for word in lbl.split() if len(word) > 3):
+                matches += 0.5  # Partial credit
+        
+        if matches > 0:
+            score += (matches / len(labels)) * 40  # Higher weight for row1 tables
+            reasons.append(f"Row1 orientation header matches: {matches}/{len(labels)}")
+    
+    # Special handling for Declaration tables (existing logic)
+    if schema_name == "Operator Declaration" and context['first_cell'].upper() == "PRINT NAME":
+        if "OPERATOR DECLARATION" in context['heading'].upper():
+            score += 80
+            reasons.append("Operator Declaration context match")
+        elif any("MANAGER" in cell.upper() for cell in context['all_cells']):
+            score += 60
+            reasons.append("Manager found in cells (likely Operator Declaration)")
+    
+    if schema_name == "NHVAS Approved Auditor Declaration" and context['first_cell'].upper() == "PRINT NAME":
+        if any("MANAGER" in cell.upper() for cell in context['all_cells']):
+            score -= 50
+            reasons.append("Penalty: Manager found (not auditor)")
+    
+    return score, reasons
 
-    # 5) left tables
+def match_table_schema(tbl):
+    """Improved table schema matching with scoring system"""
+    context = get_table_context(tbl)
+    best_match = None
+    best_score = 0
     for name, spec in TABLE_SCHEMAS.items():
-        if spec["orientation"]=="left" and all(lbl in col0 for lbl in spec["labels"]):
-            return name
-
+        score, reasons = calculate_schema_match_score(name, spec, context)
+        if score > best_score:
+            best_score = score
+            best_match = name
+    if best_score >= 20:
+        return best_match
     return None
 
-def extract_red_text(path):
-    doc = Document(path)
-    out = {}
-
-    # --- TABLES ---
-    for tbl in doc.tables:
-        schema = match_table_schema(tbl)
-        if not schema:
-            continue
-        spec = TABLE_SCHEMAS[schema]
-
-        # handle the special split_labels (row1 only)
-        if spec.get("split_labels") and spec["orientation"]=="row1":
-            cell_txt = tbl.rows[1].cells[0].text.strip()
-            first_lbl = spec["split_labels"][0]
-            narrative, _, tail = cell_txt.partition(first_lbl)
-            narrative = narrative.strip()
-            if narrative:
-                out.setdefault(schema, {}).setdefault(spec["labels"][0], []).append(narrative)
+def check_multi_schema_table(tbl):
+    """Check if table contains multiple schemas and split appropriately"""
+    context = get_table_context(tbl)
+    operator_labels = ["Operator name (Legal entity)", "NHVAS Accreditation No.", "Registered trading name/s", 
+                      "Australian Company Number", "NHVAS Manual"]
+    contact_labels = ["Operator business address", "Operator Postal address", "Email address", "Operator Telephone Number"]
+    has_operator = any(any(op_lbl.upper() in cell.upper() for op_lbl in operator_labels) for cell in context['col0'])
+    has_contact = any(any(cont_lbl.upper() in cell.upper() for cont_lbl in contact_labels) for cell in context['col0'])
+    if has_operator and has_contact:
+        return ["Operator Information", "Operator contact details"]
+    return None
 
-            for i, lbl in enumerate(spec["split_labels"]):
-                nxt = spec["split_labels"][i+1] if i+1<len(spec["split_labels"]) else None
-                pattern = rf"{re.escape(lbl)}\s*(.+?)(?={re.escape(nxt)})" if nxt else rf"{re.escape(lbl)}\s*(.+)$"
-                m = re.search(pattern, cell_txt, flags=re.DOTALL)
-                if m:
-                    val = m.group(1).strip()
-                    out.setdefault(schema, {}).setdefault(lbl, []).append(val)
+def extract_multi_schema_table(tbl, schemas):
+    """Extract data from table with multiple schemas"""
+    result = {}
+    for schema_name in schemas:
+        if schema_name not in TABLE_SCHEMAS:
             continue
-
-        # normal tables
-        labels = spec["labels"] + [schema]
-        collected = {lbl: [] for lbl in labels}
-        seen      = {lbl: set() for lbl in labels}
-        by_col = (spec["orientation"]=="row1")
-
-        rows = tbl.rows[1:]
-        for ri, row in enumerate(rows):
+        spec = TABLE_SCHEMAS[schema_name]
+        schema_data = {}
+        for ri, row in enumerate(tbl.rows):
+            if ri == 0:
+                continue
+            row_label = normalize_text(row.cells[0].text)
+            belongs_to_schema = False
+            matched_label = None
+            for spec_label in spec["labels"]:
+                spec_norm = normalize_text(spec_label).upper()
+                row_norm = row_label.upper()
+                if spec_norm == row_norm or spec_norm in row_norm or row_norm in spec_norm:
+                    belongs_to_schema = True
+                    matched_label = spec_label
+                    break
+            if not belongs_to_schema:
+                continue
             for ci, cell in enumerate(row.cells):
                 red_txt = "".join(run.text for p in cell.paragraphs for run in p.runs if is_red_font(run)).strip()
-                if not red_txt:
-                    continue
+                if red_txt:
+                    if matched_label not in schema_data:
+                        schema_data[matched_label] = []
+                    if red_txt not in schema_data[matched_label]:
+                        schema_data[matched_label].append(red_txt)
+        if schema_data:
+            result[schema_name] = schema_data
+    return result
 
-                if by_col:
-                    # column header → your defined label
-                    lbl = spec["labels"][ci] if ci < len(spec["labels"]) else schema
+def extract_table_data(tbl, schema_name, spec):
+    """Extract red text data from table based on schema - ENHANCED for Vehicle Registration"""
+    
+    # 🎯 SPECIAL HANDLING for Vehicle Registration tables
+    if "Vehicle Registration" in schema_name:
+        print(f"    🚗 EXTRACTION FIX: Processing Vehicle Registration table")
+        
+        labels = spec["labels"]
+        collected = {lbl: [] for lbl in labels}
+        seen = {lbl: set() for lbl in labels}
+        
+        # For Vehicle Registration, orientation is "row1" - headers in first row
+        if len(tbl.rows) < 2:
+            print(f"    ❌ Vehicle table has less than 2 rows")
+            return {}
+        
+        # Map header cells to labels
+        header_row = tbl.rows[0]
+        column_mapping = {}
+        
+        print(f"    📋 Mapping {len(header_row.cells)} header cells to labels")
+        
+        for col_idx, cell in enumerate(header_row.cells):
+            header_text = normalize_text(cell.text).strip()
+            if not header_text:
+                continue
+                
+            print(f"      Column {col_idx}: '{header_text}'")
+            
+            # Find best matching label
+            best_match = None
+            best_score = 0
+            
+            for label in labels:
+                # Direct match
+                if header_text.upper() == label.upper():
+                    best_match = label
+                    best_score = 1.0
+                    break
+                
+                # Partial keyword matching
+                header_words = set(word.upper() for word in header_text.split() if len(word) > 2)
+                label_words = set(word.upper() for word in label.split() if len(word) > 2)
+                
+                if header_words and label_words:
+                    common_words = header_words.intersection(label_words)
+                    if common_words:
+                        score = len(common_words) / max(len(header_words), len(label_words))
+                        if score > best_score and score >= 0.4:  # Lower threshold for vehicle tables
+                            best_score = score
+                            best_match = label
+            
+            if best_match:
+                column_mapping[col_idx] = best_match
+                print(f"        ✅ Mapped to: '{best_match}' (score: {best_score:.2f})")
+            else:
+                print(f"        ⚠️ No mapping found for '{header_text}'")
+        
+        print(f"    📊 Total column mappings: {len(column_mapping)}")
+        
+        # Extract red text from data rows (skip header)
+        for row_idx in range(1, len(tbl.rows)):
+            row = tbl.rows[row_idx]
+            print(f"      📌 Processing data row {row_idx}")
+            
+            for col_idx, cell in enumerate(row.cells):
+                if col_idx in column_mapping:
+                    label = column_mapping[col_idx]
+                    
+                    # Extract red text
+                    red_txt = "".join(run.text for p in cell.paragraphs for run in p.runs if is_red_font(run)).strip()
+                    
+                    if red_txt:
+                        print(f"        🔴 Found red text in '{label}': '{red_txt}'")
+                        
+                        if red_txt not in seen[label]:
+                            seen[label].add(red_txt)
+                            collected[label].append(red_txt)
+        
+        # Return only non-empty collections
+        result = {k: v for k, v in collected.items() if v}
+        print(f"    ✅ Vehicle Registration extracted: {len(result)} columns with data")
+        return result
+    
+    # 🎯 ORIGINAL CODE for all other tables (unchanged)
+    labels = spec["labels"] + [schema_name]
+    collected = {lbl: [] for lbl in labels}
+    seen = {lbl: set() for lbl in labels}
+    by_col = (spec["orientation"] == "row1")
+    start_row = 1 if by_col else 0
+    rows = tbl.rows[start_row:]
+    
+    for ri, row in enumerate(rows):
+        for ci, cell in enumerate(row.cells):
+            red_txt = "".join(run.text for p in cell.paragraphs for run in p.runs if is_red_font(run)).strip()
+            if not red_txt:
+                continue
+            if by_col:
+                if ci < len(spec["labels"]):
+                    lbl = spec["labels"][ci]
                 else:
-                    # first cell in this row → must be one of your labels
-                    raw = row.cells[0].text.strip()
-                    lbl = raw if raw in spec["labels"] else schema
+                    lbl = schema_name
+            else:
+                raw_label = normalize_text(row.cells[0].text)
+                lbl = None
+                for spec_label in spec["labels"]:
+                    if normalize_text(spec_label).upper() == raw_label.upper():
+                        lbl = spec_label
+                        break
+                if not lbl:
+                    for spec_label in spec["labels"]:
+                        spec_norm = normalize_text(spec_label).upper()
+                        raw_norm = raw_label.upper()
+                        if spec_norm in raw_norm or raw_norm in spec_norm:
+                            lbl = spec_label
+                            break
+                if not lbl:
+                    lbl = schema_name
+            if red_txt not in seen[lbl]:
+                seen[lbl].add(red_txt)
+                collected[lbl].append(red_txt)
+    return {k: v for k, v in collected.items() if v}
 
-                if red_txt not in seen[lbl]:
-                    seen[lbl].add(red_txt)
-                    collected[lbl].append(red_txt)
-
-        # keep only non-empty
-        data = {k:v for k,v in collected.items() if v}
+def extract_red_text(input_doc):
+    # input_doc: docx.Document object or file path
+    if isinstance(input_doc, str):
+        doc = Document(input_doc)
+    else:
+        doc = input_doc
+    out = {}
+    table_count = 0
+    for tbl in doc.tables:
+        table_count += 1
+        multi_schemas = check_multi_schema_table(tbl)
+        if multi_schemas:
+            multi_data = extract_multi_schema_table(tbl, multi_schemas)
+            for schema_name, schema_data in multi_data.items():
+                if schema_data:
+                    if schema_name in out:
+                        for k, v in schema_data.items():
+                            if k in out[schema_name]:
+                                out[schema_name][k].extend(v)
+                            else:
+                                out[schema_name][k] = v
+                    else:
+                        out[schema_name] = schema_data
+            continue
+        schema = match_table_schema(tbl)
+        if not schema:
+            continue
+        spec = TABLE_SCHEMAS[schema]
+        data = extract_table_data(tbl, schema, spec)
         if data:
-            out[schema] = data
-
-    # --- PARAGRAPHS ---
+            if schema in out:
+                for k, v in data.items():
+                    if k in out[schema]:
+                        out[schema][k].extend(v)
+                    else:
+                        out[schema][k] = v
+            else:
+                out[schema] = data
     paras = {}
     for idx, para in enumerate(doc.paragraphs):
         red_txt = "".join(r.text for r in para.runs if is_red_font(r)).strip()
         if not red_txt:
             continue
-
-        # find nearest heading above
         context = None
         for j in range(idx-1, -1, -1):
-            txt = doc.paragraphs[j].text.strip()
-            if txt and any(re.search(p, txt) for p in HEADING_PATTERNS["main"]+HEADING_PATTERNS["sub"]):
-                context = txt
-                break
-
-        # fallback for date line
+            txt = normalize_text(doc.paragraphs[j].text)
+            if txt:
+                all_patterns = HEADING_PATTERNS["main"] + HEADING_PATTERNS["sub"]
+                if any(re.search(p, txt, re.IGNORECASE) for p in all_patterns):
+                    context = txt
+                    break
         if not context and re.fullmatch(PARAGRAPH_PATTERNS["date_line"], red_txt):
             context = "Date"
-
-        paras.setdefault(context or "(para)", []).append(red_txt)
-
+        if not context:
+            context = "(para)"
+        paras.setdefault(context, []).append(red_txt)
     if paras:
         out["paragraphs"] = paras
-
     return out
 
+def extract_red_text_filelike(input_file, output_file):
+    """
+    Accepts:
+      input_file: file-like object (BytesIO/File) or path
+      output_file: file-like object (opened for writing text) or path
+    """
+    if hasattr(input_file, "seek"):
+        input_file.seek(0)
+    doc = Document(input_file)
+    result = extract_red_text(doc)
+    if hasattr(output_file, "write"):
+        json.dump(result, output_file, indent=2, ensure_ascii=False)
+        output_file.flush()
+    else:
+        with open(output_file, "w", encoding="utf-8") as f:
+            json.dump(result, f, indent=2, ensure_ascii=False)
+    return result
+
 if __name__ == "__main__":
-    fn = sys.argv[1] if len(sys.argv)>1 else "test.docx"
-    print(json.dumps(extract_red_text(fn), indent=2, ensure_ascii=False))
+    # Support both script and app/file-like usage
+    if len(sys.argv) == 3:
+        input_docx = sys.argv[1]
+        output_json = sys.argv[2]
+        doc = Document(input_docx)
+        word_data = extract_red_text(doc)
+        with open(output_json, 'w', encoding='utf-8') as f:
+            json.dump(word_data, f, indent=2, ensure_ascii=False)
+        print(json.dumps(word_data, indent=2, ensure_ascii=False))
+    else:
+        print("To use as a module: extract_red_text_filelike(input_file, output_file)")