src/cross_reference/test_cross_reference.py

import pandas as pd
import os
import logging
import re
import json

from src.segmentation.parser import LegalDocumentParser
from src.segmentation.writer import SegmentWriter
from src.segmentation.confidence import ConfidenceScorer
from src.cross_reference.extractor import CrossReferenceExtractor

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s — %(message)s")
logger = logging.getLogger(__name__)

NEO4J_URI = os.getenv("NEO4J_URI", "bolt://localhost:7687")
NEO4J_USER = os.getenv("NEO4J_USER", "neo4j")
NEO4J_PASS = os.getenv("NEO4J_PASSWORD", "password")

def run_unified_pipeline():
    # 1.1 Infrastructure Skeleton
    scorer = ConfidenceScorer()
    parser = LegalDocumentParser()
    
    logger.info("Đọc metadata_deduped.parquet...")
    meta_df = pd.read_parquet("data/metadata_deduped.parquet")
    meta_df['id'] = meta_df['id'].astype(str)
    
    # 1.2 Filter loai_van_ban and ngay_ban_hanh
    core_types = ['Thông tư', 'Nghị định', 'Luật', 'Bộ luật']
    validity = ['Còn hiệu lực','Hết hiệu lực một phần']
    meta_df['ngay_ban_hanh'] = pd.to_datetime(meta_df['ngay_ban_hanh'], errors='coerce', dayfirst=True)
    filtered_meta = meta_df[
        (meta_df['loai_van_ban'].isin(core_types)) &
        (meta_df['tinh_trang_hieu_luc'].isin(validity)) &
        (meta_df['ngay_ban_hanh'] >= '2000-01-01')
    ]
    logger.info(f"Tổng số văn bản sau khi lọc: {len(filtered_meta)}")
    
    # Load lookup table from JSON (Centralized for all scripts)
    lookup_path = "data/so_ky_hieu_lookup.json"
    if not os.path.exists(lookup_path):
        logger.info("Không thấy bảng tra cứu JSON, đang tạo mới...")
        from src.data_pipeline.build_lookup_json import build
        build()
        
    with open(lookup_path, "r", encoding="utf-8") as f:
        lookup = json.load(f)
    logger.info(f"Đã nạp bảng tra cứu từ JSON ({len(lookup)} mục).")
    
    extractor = CrossReferenceExtractor(lookup_table=lookup)
    
    logger.info("Đọc content_clean.parquet...")
    content_df = pd.read_parquet("data/content_clean.parquet")
    content_df['id'] = content_df['id'].astype(str)
    content_dict = content_df.set_index('id')['clean_html'].to_dict()
    
    total_docs = 0
    batch_size = 50  # 3.1 Batching to prevent memory issues
    batch_segments = []
    part_idx = 0
    
    # Lists for relationships
    internal_refs_data = []
    external_refs_data = []
    modifies_refs_data = []
    
    doc_id = '153913'
    meta_row = filtered_meta[filtered_meta['id'] == doc_id]
    if meta_row.empty:
        logger.error(f"Văn bản ID {doc_id} không tìm thấy trong metadata.")
        return
    row = meta_row.iloc[0]
    
    html = content_dict.get(doc_id, "")
        
    try:
        # 2.1 Stage 2: Preamble Extraction
        logger.info(f"Start stage 2: Preamble Extraction...")
        preamble_text = ""
        parts = re.split(r'(<[^>]+>\s*Điều\s+1[\.:\s])', html, maxsplit=1, flags=re.IGNORECASE)
        if len(parts) > 1:
            preamble_text = parts[0]
        
        primary_target_ref = extractor._extract_preamble_anchor(preamble_text)
        if primary_target_ref:
            logger.info(f"Tìm thấy văn bản đích từ lời nói đầu: {primary_target_ref.raw_so_ky_hieu} (ID: {primary_target_ref.target_doc_id})")
        
        # 2.2 Stage 3: Segmentation
        logger.info(f"Start stage 3: Segmentation...")
        result = parser.parse(
            doc_id=doc_id,
            clean_html=html,
            loai_van_ban=row.get('loai_van_ban', '')
        )
        logger.info(f"Phân tích segmentation: {len(result.segments)} segments")
        
        # 2.3 & 2.4 Stage 4: Cross-Reference & Context-Aware Extraction
        logger.info(f"Start stage 4: Cross-Reference & Context-Aware Extraction...")
        is_modifying = (primary_target_ref is not None) or ("sửa đổi" in str(row.get('title', '')).lower())
        
        from src.segmentation.models import HierarchyType
        all_relationships = []

        def parse_uid_parts(uid):
            if not uid: return "", "", ""
            p = uid.split('_')
            d = k = di = ""
            if 'dieu' in p: d = p[p.index('dieu')+1]
            if 'khoan' in p: k = p[p.index('khoan')+1]
            if 'diem' in p: di = p[p.index('diem')+1]
            return d, k, di
        logger.info(f"Executing stage 4: Parse done...")
        uid_to_seg = {s.uid: s for s in result.segments if s.uid}
        last_target_doc_id = None
        last_target_article = None

        for seg in result.segments:
            if seg.hierarchy_type not in [HierarchyType.DIEU, HierarchyType.KHOAN, HierarchyType.DIEM]:
                continue
            
            # Xác định article_uid, clause_uid, point_uid cho source
            art_uid = cl_uid = pt_uid = None
            curr = seg
            if curr.hierarchy_type == HierarchyType.DIEM:
                pt_uid = curr.uid
                curr = uid_to_seg.get(curr.parent_uid)
            if curr and curr.hierarchy_type == HierarchyType.KHOAN:
                cl_uid = curr.uid
                curr = uid_to_seg.get(curr.parent_uid)
            if curr and curr.hierarchy_type == HierarchyType.DIEU:
                art_uid = curr.uid
            
            if not art_uid: continue

            ext_result = extractor.extract_from_article(
                doc_id=doc_id,
                article_uid=art_uid,
                clause_uid=cl_uid,
                point_uid=pt_uid,
                article_text=seg.clean_text,
                is_modifying_doc=is_modifying
            )

            src_art, src_cl, src_pt = parse_uid_parts(seg.uid)

            # Thu thập Internal Refs
            for r in ext_result.internal_refs:
                all_relationships.append({
                    "src_doc": doc_id, "src_art": src_art, "src_cl": src_cl, "src_pt": src_pt,
                    "tgt_doc": doc_id, "tgt_art": r.target_article_index or "", 
                    "tgt_cl": r.target_clause_index or "", "tgt_pt": r.target_point_label or "",
                    "type": "Internal", "context": r.context_text.replace('\n', ' ')
                })

            # Thu thập External Refs
            for r in ext_result.external_refs:
                all_relationships.append({
                    "src_doc": doc_id, "src_art": src_art, "src_cl": src_cl, "src_pt": src_pt,
                    "tgt_doc": r.target_doc_id or r.raw_so_ky_hieu, "tgt_art": r.target_article_index or "", 
                    "tgt_cl": r.target_clause_index or "", "tgt_pt": r.target_point_label or "",
                    "type": "External", "context": r.context_text.replace('\n', ' ')
                })

            # Thu thập Modification Refs
            for r in ext_result.modification_refs:
                # Bổ sung thông tin target_doc_id nếu thiếu từ primary target (lời nói đầu)
                if not r.target_doc_id and primary_target_ref:
                    r.target_doc_id = primary_target_ref.target_doc_id

                # Logic ROLL BACK: Nếu ref thiếu Điều đích, lấy từ ref trước đó
                if r.is_partial_ref and last_target_article:
                    r.target_article_index = last_target_article
                    if not r.target_doc_id:
                        r.target_doc_id = last_target_doc_id
                
                # Cập nhật state cho các ref tiếp theo
                if r.target_article_index:
                    last_target_article = r.target_article_index
                    last_target_doc_id = r.target_doc_id

                all_relationships.append({
                    "src_doc": doc_id, "src_art": src_art, "src_cl": src_cl, "src_pt": src_pt,
                    "tgt_doc": r.target_doc_id or r.raw_target_so_ky_hieu, "tgt_art": r.target_article_index or "", 
                    "tgt_cl": r.target_clause_index or "", "tgt_pt": r.target_point_label or "",
                    "type": f"Modification ({r.action.value if hasattr(r.action, 'value') else r.action})", 
                    "context": r.context_text.replace('\n', ' ')
                })
        logger.info(f"Executing stage 4: Extract done...")

        # Xuất file Markdown
        output_file = f"test_{doc_id}.md"
        with open(output_file, "w", encoding="utf-8") as f:
            f.write(f"# Kết quả trích dẫn quan hệ - Văn bản {doc_id}\n\n")
            f.write("| Source Docs | Article | Clause | Point | Target Docs | Article | Clause | Point | Type | Context |\n")
            f.write("|-------------|---------|--------|-------|-------------|---------|--------|-------|------|---------|\n")
            for rel in all_relationships:
                f.write(f"| {rel['src_doc']} | {rel['src_art']} | {rel['src_cl']} | {rel['src_pt']} | "
                        f"{rel['tgt_doc']} | {rel['tgt_art']} | {rel['tgt_cl']} | {rel['tgt_pt']} | "
                        f"{rel['type']} | {rel['context']} |\n")
        
        logger.info(f"Đã xuất kết quả ra file: {output_file}")

    except Exception as e:
        logger.error(f"Lỗi xử lý nội dung văn bản {doc_id}: {e}")
        import traceback
        logger.error(traceback.format_exc())

    logger.info("=== HOÀN TẤT UNIFIED PIPELINE ===")

if __name__ == "__main__":
    run_unified_pipeline()