DoLuEdu-WordFileParse/scripts/omml_mathml2latex/strategies/chemistry_strategy.py

"""
Chemistry strategy for chemical equation notation.

This strategy handles formatting for chemical equations, including element
symbols, reaction arrows, ion charges, and state symbols.
"""

from lxml import etree as ET
from typing import Set
from ..core.interfaces import Strategy

# Import from existing data modules
from ..data import CHEMICAL_ELEMENTS, UNICODE_TO_LATEX, MATH_FUNCTIONS, UNICODE_PASSTHROUGH


class ChemistryStrategy(Strategy):
    """
    Chemical equation notation strategy.

    This strategy provides formatting for chemical content:
    - Chemical elements → \\mathrm{} (upright, not italic)
    - Reaction arrows → \\xlongequal{condition} or \\rightarrow
    - Ion charges → proper superscript formatting
    - State symbols → \\uparrow, \\downarrow
    - Chinese condition text → \\text{}
    """

    # Chemical elements (118 elements)
    ELEMENTS: Set[str] = CHEMICAL_ELEMENTS

    # Mathematical functions still get special treatment
    FUNCTIONS: Set[str] = MATH_FUNCTIONS

    def __init__(self):
        """Initialize the chemistry strategy."""
        # Subject mode identifier for probability scoring
        self.subject_mode = 'chemistry'

    def format_identifier(self, text: str, node: ET.Element) -> str:
        """
        Format an identifier (<mi> element) for chemistry.

        UPDATED: Now ONLY respects mathvariant attribute, no automatic element wrapping.
        User has full control through explicit mathvariant settings in Word.

        Args:
            text: The identifier text content
            node: The <mi> element node

        Returns:
            Formatted LaTeX string

        Logic (SIMPLIFIED):
            1. Mathematical functions → \\function (e.g., "sin" → "\\sin")
            2. mathvariant="normal" → \\mathrm{text} (user explicitly set upright)
            3. Otherwise → return as-is (remain italic, default for math)

        Examples:
            "Fe" with mathvariant="normal" → "\\mathrm{Fe}"
            "Fe" without mathvariant → "Fe" (italic)
            "H" with mathvariant="normal" → "\\mathrm{H}"
            "F" without mathvariant → "F" (italic, NOT wrapped)
            "sin" → "\\sin" (still a math function)
        """
        # Check if it's a mathematical function (rare in chemistry, but possible)
        if text in self.FUNCTIONS:
            return f"\\{text}"

        # DISABLED: Automatic element wrapping (user wants full control)
        # Old logic that auto-wrapped elements is now commented out
        # # Check if it's a chemical element
        # if text in self.ELEMENTS:
        #     return f"\\mathrm{{{text}}}"

        # DISABLED: Automatic capital letter wrapping
        # # Single CAPITAL letters → wrap in \mathrm{} (likely element symbols)
        # if len(text) == 1 and text.isupper():
        #     return f"\\mathrm{{{text}}}"

        # ONLY respect mathvariant attribute (user's explicit setting)
        mathvariant = node.get('mathvariant', '')
        if mathvariant == 'normal':
            # User explicitly set this to upright in Word (via m:sty='p' in OMML)
            return f"\\mathrm{{{text}}}"

        # UNICODE_PASSTHROUGH: Keep certain symbols as Unicode for MathJax compatibility
        if text in UNICODE_PASSTHROUGH:
            return text

        # Check for Unicode symbols
        if text in UNICODE_TO_LATEX:
            return UNICODE_TO_LATEX[text]

        # Handle multi-character text with Unicode
        if len(text) > 1:
            # Check if it contains Unicode (passthrough or convertible)
            has_unicode = any(char in UNICODE_TO_LATEX for char in text)
            has_passthrough = any(char in UNICODE_PASSTHROUGH for char in text)
            if has_unicode or has_passthrough:
                converted_parts = []
                for char in text:
                    if char in UNICODE_PASSTHROUGH:
                        converted_parts.append(char)
                    elif char in UNICODE_TO_LATEX:
                        converted_parts.append(UNICODE_TO_LATEX[char])
                    elif char.isspace():
                        converted_parts.append(' ')
                    else:
                        converted_parts.append(char)
                return ' '.join(converted_parts)

        # Default: return as-is (italic, standard for math variables)
        return text

    def format_operator(self, text: str, node: ET.Element) -> str:
        """
        Format an operator (<mo> element) for chemistry.

        Args:
            text: The operator text content
            node: The <mo> element node

        Returns:
            Formatted LaTeX string

        Logic:
            1. Reaction arrows → \\xlongequal{condition} or \\rightarrow
            2. State symbols → \\uparrow, \\downarrow
            3. Unicode symbols → LaTeX equivalents (via UNICODE_TO_LATEX)
            4. Standard operators → unchanged (+, -, =)

        Examples:
            "→" → "\\rightarrow"
            "↑" → "\\uparrow"
            "↓" → "\\downarrow"
            "=" → "="
            "+" → "+"
        """
        # Handle < and > specifically (LaTeX convention)
        if text == "<":
            return "\\lt"
        if text == ">":
            return "\\gt"

        # Reaction arrows
        if text == '→':
            return '\\rightarrow'
        elif text == '←':
            return '\\leftarrow'
        elif text == '⇌' or text == '⇄':
            return '\\rightleftharpoons'
        elif text == '↑':
            return '\\uparrow'
        elif text == '↓':
            return '\\downarrow'

        # UNICODE_PASSTHROUGH: Keep certain symbols as Unicode for MathJax compatibility
        if text in UNICODE_PASSTHROUGH:
            return text

        # Check for Unicode symbol mapping
        if text in UNICODE_TO_LATEX:
            return UNICODE_TO_LATEX[text]

        # Handle multi-character operators with Unicode
        if len(text) > 1:
            has_unicode = any(char in UNICODE_TO_LATEX for char in text)
            has_passthrough = any(char in UNICODE_PASSTHROUGH for char in text)
            if has_unicode or has_passthrough:
                converted_parts = []
                for char in text:
                    if char in UNICODE_PASSTHROUGH:
                        converted_parts.append(char)
                    elif char in UNICODE_TO_LATEX:
                        converted_parts.append(UNICODE_TO_LATEX[char])
                    else:
                        converted_parts.append(char)
                return ' '.join(converted_parts)

        # Default: return as-is
        return text

    def format_number(self, text: str, node: ET.Element) -> str:
        """
        Format a number (<mn> element) for chemistry.

        Args:
            text: The number text content
            node: The <mn> element node

        Returns:
            Formatted LaTeX string

        Logic:
            Numbers are returned as-is in chemistry mode.

        Examples:
            "2" → "2" (stoichiometric coefficient)
            "3.5" → "3.5" (pH value)
        """
        # Numbers are returned as-is
        return text

    def detect_context(self, equation_text: str) -> bool:
        """
        Detect if an equation is chemical content.

        Args:
            equation_text: The full equation MathML text

        Returns:
            True if likely chemical, False otherwise

        Detection indicators:
            - Chemical elements (Fe, Cu, H2O, etc.)
            - Reaction arrows (→, ⇌, ↑, ↓)
            - Ion charges (2+, 3-, etc.)
            - Multiple element symbols in sequence

        Examples:
            "Fe + H2SO4 → FeSO4 + H2" → True
            "2H2O ⇌ H3O+ + OH-" → True
            "sin(x) + cos(x)" → False
            "P = C" → False (just math variables, not Phosphorus = Carbon)
        """
        # CRITICAL FIX: Use proper XML parsing instead of naive substring search
        # The old method falsely detected "ln" in "xmlns" and "sup" in "<msup>" tags
        # as mathematical functions, causing chemistry equations to be rejected.

        try:
            # Parse MathML XML to ElementTree
            root = ET.fromstring(equation_text.encode('utf-8'))

            # Use the robust chemistry detection from utils module
            from ..utils.chemistry_detection import detect_chemistry_equation
            return detect_chemistry_equation(root)

        except (ET.XMLSyntaxError, ValueError):
            # Fallback to string-based detection if XML parsing fails
            # This should rarely happen with valid MathML

            # Check for UNAMBIGUOUS multi-character elements
            unambiguous_elements = {e for e in self.ELEMENTS if len(e) > 1}
            element_count = sum(1 for element in unambiguous_elements if element in equation_text)
            if element_count >= 1:
                return True

            # Check for reaction arrows (strong indicator)
            if any(arrow in equation_text for arrow in ['→', '←', '⇌', '⇄', '↑', '↓']):
                return True

            return False

    @property
    def name(self) -> str:
        """Get the strategy name."""
        return "chemistry"

    def __repr__(self) -> str:
        """String representation for debugging."""
        return f"ChemistryStrategy(elements={len(self.ELEMENTS)})"