evaluation.py

import os
import sys
import re
import string
import json
import tempfile
import ast
import xml.sax.saxutils as saxutils
from collections import defaultdict, Counter
from PIL import Image
import jiwer
import pandas as pd
from lxml import etree  # STRICTLY REQUIRED for reference paper logic

# --- Fix imports when running from project root ---
# Add the project root and src directory to Python path
script_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.abspath(os.path.join(script_dir, "../.."))
src_dir = os.path.join(project_root, "src")
dots_ocr_dir = os.path.join(project_root, "dots.ocr")

if src_dir not in sys.path:
    sys.path.insert(0, src_dir)
if dots_ocr_dir not in sys.path:
    sys.path.insert(0, dots_ocr_dir)

# --- Imports from your project structure ---
from chamberact.chamberact.config import VLM_SAMPLING_PARAMS, CONVOCATIONS_FILE
from chamberact.vlm_handler import VLMHandler
from chamberact.postprocessor import TranscriptionPostprocessor
from chamberact.kb_retriever import KnowledgeBaseRetriever
from chamberact.utils import save_json
from chamberact.speaker_matcher import match_speakers_to_uris

# --- Imports for DotsOCR ---
from dots_ocr.parser import DotsOCRParser
from dots_ocr.utils.consts import MIN_PIXELS, MAX_PIXELS

# --- Configuration ---
DOTS_HOST = "127.0.0.1"
DOTS_PORT = 8000
PROMPT_MODE = "prompt_layout_all_en"

PREDICTIONS_PATH = "./evaluation/predictions/"
JPG_PATH = "./evaluation/gold_standard_jpg/"
TXT_PATH = "./evaluation/gold_standard/"
XML_PATH = "./evaluation/gold_standard_xml/"


def normalize_text(text):
    """Lowercase, remove punctuation, and collapse whitespace."""
    if not text:
        return ""
    text = text.lower()
    text = text.translate(str.maketrans("", "", string.punctuation + "«»—"))
    text = re.sub(r"\s+", " ", text).strip()
    return text


def postprocess_dots_ocr_output(text):
    """Post-process DotsOCR output to clean up formatting artifacts.

    Removes:
    - Markdown bold formatting (**text**)
    - Markdown italic formatting (*text* and _text_)
    - Markdown headers (# ## ### etc.)
    - End-of-line hyphenation (re-joins hyphenated words)
    """
    if not text:
        return ""

    # Remove markdown bold (**text**)
    text = re.sub(r"\*\*(.+?)\*\*", r"\1", text)

    # Remove markdown italic (*text* but not inside words)
    text = re.sub(r"(?<!\w)\*(.+?)\*(?!\w)", r"\1", text)

    # Remove markdown italic with underscores (_text_)
    text = re.sub(r"(?<!\w)_(.+?)_(?!\w)", r"\1", text)

    # Remove markdown headers (# ## ### etc. at start of lines)
    text = re.sub(r"^#+\s*", "", text, flags=re.MULTILINE)

    # Remove end-of-line hyphenation (rejoin hyphenated words)
    # Matches lowercase letter followed by hyphen, whitespace/newline, then lowercase letter
    text = re.sub(r"(?<=[a-z])-\s+(?=[a-z])", "", text)

    # Also handle hyphenation at line breaks specifically
    text = re.sub(r"(?<=[a-z])-\n(?=[a-z])", "", text)

    return text


def extract_text_from_dots_output(output_dir, filename_base):
    """Parses JSON output from DotsOCR."""
    json_filename = f"{filename_base}_page_0.json"
    md_filename = f"{filename_base}_page_0.md"
    json_filepath = os.path.join(output_dir, json_filename)
    md_filepath = os.path.join(output_dir, md_filename)
    extracted_text = ""

    try:
        if not os.path.exists(json_filepath):
            raise FileNotFoundError
        with open(json_filepath, "r", encoding="utf-8") as f:
            data = json.load(f)
        if not isinstance(data, list):
            data = ast.literal_eval(data)
        texts = [item["text"] for item in data if "text" in item and item["text"]]
        extracted_text = "\n\n".join(texts)
    except (json.JSONDecodeError, FileNotFoundError, TypeError, SyntaxError):
        if os.path.exists(md_filepath):
            with open(md_filepath, "r", encoding="utf-8") as f:
                extracted_text = f.read()
    return extracted_text


def run_dots_ocr_on_image(parser: DotsOCRParser, image_path: str) -> tuple[str, str]:
    """Runs DotsOCR on a single image via temp PDF."""
    with tempfile.TemporaryDirectory() as temp_dir:
        try:
            img = Image.open(image_path)
            temp_pdf_path = os.path.join(temp_dir, "temp_input.pdf")
            img.save(temp_pdf_path, "PDF", resolution=200.0)

            filename_base = "temp_input"
            parser.parse_file(
                input_path=temp_pdf_path,
                output_dir=temp_dir,
                prompt_mode=PROMPT_MODE,
            )

            parser_output_dir = os.path.join(temp_dir, filename_base)
            ocr_text = extract_text_from_dots_output(parser_output_dir, filename_base)
            # Apply post-processing to clean up DotsOCR output
            postprocessed_ocr_text = postprocess_dots_ocr_output(ocr_text)
            return ocr_text, postprocessed_ocr_text
        except Exception as e:
            print(
                f"Error processing image {image_path} with DotsOCR: {e}",
                file=sys.stderr,
            )
            return "", ""


def get_speakers_from_xml(xml_content):
    """
    Parses XML string and extracts speakers using the EXACT logic
    from the reference paper:
    1. Uses lxml.
    2. Finds <speech is_president='true'>, counts as 'president', removes element.
    3. Finds remaining <speech>, extracts 'speaker' attr, splits by '/'.

    For presidents (is_president="true"): Always count as "president" regardless of speaker value.
    For non-presidents: Only count if speaker is identified (not "none" or "unknown").
    """
    try:
        root = etree.fromstring(bytes(xml_content, encoding="utf-8"))
    except etree.XMLSyntaxError:
        return []

    speaker_list = []

    # Logic Step 1: Handle Presidents
    # Always count presidents, even if speaker="none", because is_president is what matters
    for element in root.findall(".//speech[@is_president='true']"):
        speaker_list.append("president")
        parent = element.getparent()
        if parent is not None:
            parent.remove(element)

    # Logic Step 2: Handle remaining speakers (non-presidents)
    # Only count if speaker is identified
    for element in root.findall(".//speech"):
        spk_attr = element.get("speaker")
        if spk_attr:
            # Reference logic: split URI and take last part
            clean_speaker = str(spk_attr).split("/")[-1]
            # Exclude "none" and "unknown" speakers for non-presidents
            if clean_speaker not in ["none", "unknown"]:
                speaker_list.append(clean_speaker)

    return speaker_list


def evaluate_tags_exact_match(gold_xml_folder, pred_xml_folder, gold_txt_folder):
    """
    Calculates Precision, Recall, and F1 using Bag-of-Entities intersection.
    Maps file names by sorting the Gold Standard TXT folder and using 1-based indexing.
    """
    results = {
        "all": {"tp": 0, "fp": 0, "fn": 0},
        "pre": {"tp": 0, "fp": 0, "fn": 0},
        "post": {"tp": 0, "fp": 0, "fn": 0},
    }

    if not os.path.exists(gold_txt_folder):
        print(f"Error: TXT folder not found at {gold_txt_folder}")
        return results

    # 1. Establish the full sequence (Ground Truth Order)
    all_files_sorted = sorted(
        [f for f in os.listdir(gold_txt_folder) if f.endswith(".txt")]
    )

    # 2. Get available Gold Standard Annotations
    gold_xml_files = [f for f in os.listdir(gold_xml_folder) if f.endswith(".xml")]

    print(
        f"Evaluating {len(gold_xml_files)} annotated files out of {len(all_files_sorted)} total sequence."
    )

    for gold_file in gold_xml_files:
        # Reconstruct the base filename (e.g., 'doc.xml' -> 'doc.txt')
        base_name_txt = gold_file.replace(".xml", ".txt")

        # Find the index in the sorted list to determine the numeric ID
        if base_name_txt not in all_files_sorted:
            continue

        file_index = all_files_sorted.index(base_name_txt)

        # STRICT 1-BASED INDEXING (as requested)
        pred_filename = f"{file_index + 1}.xml"
        pred_path = os.path.join(pred_xml_folder, pred_filename)

        if not os.path.exists(pred_path):
            print(
                f"Warning: Prediction file {pred_filename} missing for Gold Standard {gold_file}"
            )
            continue

        # Determine Era
        if "repubblica" in gold_file.lower():
            era = "post"
        else:
            era = "pre"

        # Read Files
        with open(os.path.join(gold_xml_folder, gold_file), "r", encoding="utf-8") as f:
            gold_content = f.read()

        with open(pred_path, "r", encoding="utf-8") as f:
            pred_content = f.read()

        # Extract Speakers
        gold_list = get_speakers_from_xml(gold_content)
        pred_list = get_speakers_from_xml(pred_content)

        # Create Counters (Bag of Entities)
        gold_speakers = Counter(gold_list)
        pred_speakers = Counter(pred_list)

        # Intersection Logic (Reference Method)
        tp = sum((gold_speakers & pred_speakers).values())
        fp = sum((pred_speakers - gold_speakers).values())
        fn = sum((gold_speakers - pred_speakers).values())

        # Accumulate
        results[era]["tp"] += tp
        results[era]["fp"] += fp
        results[era]["fn"] += fn

        results["all"]["tp"] += tp
        results["all"]["fp"] += fp
        results["all"]["fn"] += fn

    # Calculate final metrics
    final_metrics = {}
    for cat, stats in results.items():
        tp, fp, fn = stats["tp"], stats["fp"], stats["fn"]
        precision = tp / (tp + fp) if (tp + fp) > 0 else 0.0
        recall = tp / (tp + fn) if (tp + fn) > 0 else 0.0
        f1 = (
            (2 * (precision * recall) / (precision + recall))
            if (precision + recall) > 0
            else 0.0
        )
        final_metrics[cat] = {"p": precision, "r": recall, "f1": f1}

    return final_metrics


def main():
    # --- Initialization ---
    print("Initializing DotsOCR Parser...")
    ocr_parser = DotsOCRParser(
        ip=DOTS_HOST,
        port=DOTS_PORT,
        dpi=200,
        num_thread=16,
        min_pixels=MIN_PIXELS,
        max_pixels=MAX_PIXELS,
        max_completion_tokens=16384,
    )

    print("Initializing Main Pipeline components...")
    postproc = TranscriptionPostprocessor()
    kb_retriever = KnowledgeBaseRetriever()
    convocations = pd.read_csv(CONVOCATIONS_FILE, sep=";")
    vlm = VLMHandler()

    os.makedirs(PREDICTIONS_PATH, exist_ok=True)
    os.makedirs(PREDICTIONS_PATH + "raw", exist_ok=True)

    # Load gold standard filenames (sorted) to establish the sequence
    textual_filenames = []
    gold_standard_transcriptions = []

    sorted_txts = sorted([f for f in os.listdir(TXT_PATH) if f.endswith(".txt")])
    for txt_file in sorted_txts:
        with open(os.path.join(TXT_PATH, txt_file), "r", encoding="utf-8") as f:
            gold_standard_transcriptions.append(f.read().strip())
        textual_filenames.append(txt_file)

    # --- STEP 1: PREPROCESSING (Small VLM) ---
    print("\n--- Starting Preprocessing (DotsOCR) ---")
    pages_data = []
    dots_ocr_outputs = []  # Store DotsOCR outputs for intermediate evaluation
    sorted_images = sorted([f for f in os.listdir(JPG_PATH) if f.endswith(".jpg")])

    for i, img_file in enumerate(sorted_images):
        img_path = os.path.join(JPG_PATH, img_file)
        print(f"Preprocessing image {i+1}/{len(sorted_images)}: {img_file}")

        image = Image.open(img_path)
        ocr_text, postprocessed_ocr_text = run_dots_ocr_on_image(ocr_parser, img_path)

        # Append data. Note: i is 0-based index here.
        pages_data.append((image, ocr_text, i))
        dots_ocr_outputs.append(
            postprocessed_ocr_text
        )  # Save for intermediate evaluation

        # Save DotsOCR intermediate output to file (1-based indexing)
        dots_ocr_filename = f"{PREDICTIONS_PATH}/{i + 1}_dots_ocr.txt"
        with open(dots_ocr_filename, "w", encoding="utf-8") as f:
            f.write(ocr_text)

    # --- STEP 2: TRANSCRIPTION (Large VLM) ---
    print("\n--- Starting Transcription (Large VLM) ---")
    transcribed_items = vlm.transcribe_document_pages(
        pages_data, VLM_SAMPLING_PARAMS, PREDICTIONS_PATH + "raw"
    )

    documents = defaultdict(list)
    for item in transcribed_items:
        documents[item["page_number_in_pdf"]].append(item)

    # Sort documents by page number to align with textual_filenames
    documents = [documents[i] for i in sorted(documents.keys())]

    # --- STEP 3: POST-PROCESSING & ENTITY LINKING ---
    print("\n--- Starting Post-processing & Entity Linking ---")
    postprocessed_documents = [
        postproc.process_transcriptions(doc) for doc in documents
    ]

    wer = []
    cer = []
    norm_wer = []
    norm_cer = []

    print("\n--- Processing Files, Saving XMLs and Computing Text Metrics ---")
    for idx, document in enumerate(postprocessed_documents):
        # Retrieve original filename based on index
        filename = textual_filenames[idx]

        # Define output numeric filename (1-based index)
        output_numeric_name = f"{idx + 1}"

        print(f"Processing {output_numeric_name}.xml (Original: {filename})")

        # --- KB Retrieval Logic ---
        known_people_in_session = {}

        if "camera" in filename.lower():
            # Expected format: camera-leg-date-...
            parts = filename.split("-")
            # Logic to extract date. E.g. camera-repubblica_06-19720705-...
            # Attempt to find the date part (8 digits)
            date_part = None
            for p in parts:
                if p.isdigit() and len(p) == 8:
                    date_part = p
                    break

            if date_part:
                convocation_matches = convocations[
                    (convocations["seduta_date"] == int(date_part))
                    & (convocations["is_assemblea"] == 1)
                ]
                if not convocation_matches.empty:
                    convocation = convocation_matches.iloc[0]
                    known_people_in_session = kb_retriever.get_people_in_session(
                        convocation["seduta_uri"], convocation["seduta_date"]
                    )
        else:
            # Senate or other
            parts = filename.split("-")
            if len(parts) > 1:
                leg = parts[1]
                # Assuming 'leg' extraction logic matches KB requirements
                known_people_in_session = kb_retriever.get_people_in_senate(leg)

        # --- Speaker Linking ---
        speakers_to_link = set()
        for item in document:
            extracted_speaker = item.get("speaker", "none")
            if extracted_speaker and extracted_speaker.lower() not in [
                "none",
                "unknown",
            ]:
                speakers_to_link.add(extracted_speaker)

        matched_uris = match_speakers_to_uris(
            speakers_to_link, known_people_in_session, document
        )

        for item in document:
            item["speaker_uri"] = list(
                matched_uris.get(item.get("speaker", "none"), ["none"])
            )

        # Save JSON prediction (for debug)
        save_json(document, f"{PREDICTIONS_PATH}/{output_numeric_name}.json")

        # --- XML Generation (CRITICAL) ---
        xml_parts = ['<?xml version="1.0" encoding="UTF-8"?>\n<document>\n']
        active_speech_uri = None

        for item in document:
            if item["type"] == "page-header":
                continue

            raw_speaker = item.get("speaker", "none")
            content = saxutils.escape(item.get("content", ""))
            uris = item.get("speaker_uri", ["none"])
            current_uri = uris[0] if uris else "none"

            is_speech_block = raw_speaker and raw_speaker.lower() not in [
                "none",
                "unknown",
            ]

            if is_speech_block:
                if current_uri != active_speech_uri:
                    if active_speech_uri is not None:
                        xml_parts.append("</speech>\n")

                    is_pres = "true" if "PRESIDENTE" in raw_speaker.upper() else "false"
                    xml_parts.append(
                        f'<speech speaker="{current_uri}" is_president="{is_pres}">'
                    )
                    xml_parts.append(f"{saxutils.escape(raw_speaker)}. {content}")
                    active_speech_uri = current_uri
                else:
                    xml_parts.append(f" {content}")
            else:
                if active_speech_uri is not None:
                    xml_parts.append("</speech>\n")
                    active_speech_uri = None
                xml_parts.append(f"{content}\n")

        if active_speech_uri is not None:
            xml_parts.append("</speech>\n")

        xml_parts.append("</document>")
        xml_output = "".join(xml_parts)

        # SAVE XML USING 1-BASED INDEX (e.g. 1.xml)
        with open(
            f"{PREDICTIONS_PATH}/{output_numeric_name}.xml", "w", encoding="utf-8"
        ) as f:
            f.write(xml_output)

        # --- Text Reconstruction for WER/CER ---
        output_parts = []
        last_speaker = None
        for item in document:
            if item["type"] == "page-header":
                continue
            current_speaker = item.get("speaker", "none")
            content = item.get("content", "")

            if current_speaker == last_speaker:
                if output_parts:
                    output_parts.append(" " + content)
            else:
                if output_parts:
                    output_parts.append("\n")
                if current_speaker and current_speaker not in ["none", "unknown"]:
                    output_parts.append(f"{current_speaker}. {content}")
                else:
                    output_parts.append(content)
            last_speaker = current_speaker

        ocr_text = "".join(output_parts)

        # Save OCR text
        with open(
            f"{PREDICTIONS_PATH}/{output_numeric_name}_ocr.txt", "w", encoding="utf-8"
        ) as f:
            f.write(ocr_text)

        # --- WER/CER Calculation ---
        gt_text = gold_standard_transcriptions[idx]
        wer.append(jiwer.wer(gt_text, ocr_text))
        cer.append(jiwer.cer(gt_text, ocr_text))

        normalized_ground_truth = normalize_text(gt_text)
        normalized_extracted = normalize_text(ocr_text)
        norm_wer.append(jiwer.wer(normalized_ground_truth, normalized_extracted))
        norm_cer.append(jiwer.cer(normalized_ground_truth, normalized_extracted))

    # --- DOTS OCR INTERMEDIATE EVALUATION ---
    print("\n--- Computing DotsOCR Intermediate Metrics ---")
    dots_wer = []
    dots_cer = []
    dots_norm_wer = []
    dots_norm_cer = []

    for idx, dots_ocr_text in enumerate(dots_ocr_outputs):
        gt_text = gold_standard_transcriptions[idx]
        dots_wer.append(jiwer.wer(gt_text, dots_ocr_text))
        dots_cer.append(jiwer.cer(gt_text, dots_ocr_text))

        normalized_ground_truth = normalize_text(gt_text)
        normalized_dots = normalize_text(dots_ocr_text)
        dots_norm_wer.append(jiwer.wer(normalized_ground_truth, normalized_dots))
        dots_norm_cer.append(jiwer.cer(normalized_ground_truth, normalized_dots))

    # --- FINAL EVALUATION ---
    print("\n--- Computing Tagging Metrics (Reference Paper Method) ---")

    # Calls the evaluation with exact logic, passing the TXT path for index alignment
    tag_metrics = evaluate_tags_exact_match(XML_PATH, PREDICTIONS_PATH, TXT_PATH)

    # --- Metrics Reporting ---
    print("\n--- Final Results ---")

    with open(f"{PREDICTIONS_PATH}/metrics.txt", "w") as f:
        f.write("=" * 50 + "\n")
        f.write("FULL PIPELINE (DotsOCR + VLM) METRICS\n")
        f.write("=" * 50 + "\n\n")

        f.write("Word Error Rate (WER):\n")
        f.write(f"Average: {sum(wer) / len(wer) if wer else 0}\n")
        f.write("Character Error Rate (CER):\n")
        f.write(f"Average: {sum(cer) / len(cer) if cer else 0}\n")
        f.write("Normalized WER:\n")
        f.write(f"Average: {sum(norm_wer) / len(norm_wer) if norm_wer else 0}\n")
        f.write("Normalized CER:\n")
        f.write(f"Average: {sum(norm_cer) / len(norm_cer) if norm_cer else 0}\n\n")

        f.write("--- Speaker Tagging Metrics ---\n")

        for cat in ["all", "pre", "post"]:
            m = tag_metrics[cat]
            header = cat.upper() if cat != "all" else "OVERALL"

            output_str = (
                f"{header} F1:        {m['f1']:.4f}\n"
                f"{header} Precision: {m['p']:.4f}\n"
                f"{header} Recall:    {m['r']:.4f}\n"
            )
            print(output_str.strip())
            f.write(output_str + "\n")

        f.write("\n" + "=" * 50 + "\n")
        f.write("DOTS OCR ONLY (INTERMEDIATE STEP) METRICS\n")
        f.write("=" * 50 + "\n\n")

        f.write("Word Error Rate (WER):\n")
        f.write(f"Average: {sum(dots_wer) / len(dots_wer) if dots_wer else 0}\n")
        f.write("Character Error Rate (CER):\n")
        f.write(f"Average: {sum(dots_cer) / len(dots_cer) if dots_cer else 0}\n")
        f.write("Normalized WER:\n")
        f.write(
            f"Average: {sum(dots_norm_wer) / len(dots_norm_wer) if dots_norm_wer else 0}\n"
        )
        f.write("Normalized CER:\n")
        f.write(
            f"Average: {sum(dots_norm_cer) / len(dots_norm_cer) if dots_norm_cer else 0}\n\n"
        )

    print("\n--- Evaluation Finished ---")


if __name__ == "__main__":
    os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
    main()