"""
Statistics strategy for statistical notation.

This strategy handles formatting for statistical equations, including
probability operators, distribution notation, and statistical symbols.
"""

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

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


class StatisticsStrategy(Strategy):
    """
    Statistical notation strategy.

    This strategy provides formatting for statistical content:
    - Probability operators → \\Pr, \\operatorname{Var}, \\operatorname{Cov}
    - Expected value E remains italic (use mathvariant="double-struck" for \\mathbb{E})
    - Distribution notation → proper formatting
    - Statistical symbols → Greek letters, operators
    - Variables remain italic (LaTeX default)
    """

    # Mathematical functions that get special LaTeX commands
    FUNCTIONS: Set[str] = MATH_FUNCTIONS

    # Statistical functions and operators (upright)
    # CRITICAL: 'E' is NOT in this list - it's just a regular identifier
    # Only use \mathbb{E} when there's explicit mathvariant="double-struck" attribute
    STATISTICAL_OPERATORS: Set[str] = {
        'Pr',       # Probability
        # 'E',      # Expected value - REMOVED: should remain italic unless mathvariant specified
        'Var',      # Variance
        'Cov',      # Covariance
        'Cor',      # Correlation
        'SD',       # Standard deviation
        'SE',       # Standard error
        'MSE',      # Mean squared error
        'RSS',      # Residual sum of squares
        'SSE',      # Sum of squared errors
        'Bias',     # Bias
        'df',       # degrees of freedom
        'logit',    # Logit function
        'probit',   # Probit function
    }

    # Distribution names (upright)
    DISTRIBUTIONS: Set[str] = {
        'Normal',
        'Binomial',
        'Poisson',
        'Exponential',
        'Uniform',
        'Gamma',
        'Beta',
        'Bernoulli',
        'Geometric',
        'Hypergeometric',
        'Multinomial',
        'NegativeBinomial',
        'Chi',
        'ChiSquared',
        'Student',
        'F',
        'Weibull',
        'Cauchy',
        'Laplace',
        'Logistic',
        'Pareto',
    }

    def __init__(self):
        """Initialize the statistics strategy."""
        pass

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

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

        Returns:
            Formatted LaTeX string

        Logic:
            1. Mathematical functions → \\function (e.g., "sin" → "\\sin")
            2. Statistical operators → \\operatorname{op} (e.g., "Var" → "\\operatorname{Var}")
            3. Distribution names → \\mathrm{dist} (e.g., "Normal" → "\\mathrm{Normal}")
            4. Single letters → unchanged (variables remain italic)
            5. Greek letters → LaTeX symbols

        Examples:
            "Pr" → "\\Pr"
            "E" → "E" (expected value - remains italic unless mathvariant="double-struck")
            "Var" → "\\operatorname{Var}"
            "Normal" → "\\mathrm{Normal}"
            "μ" → "\\mu"
            "σ" → "\\sigma"
            "x" → "x" (random variable)
        """
        # Check if it's a mathematical function
        if text in self.FUNCTIONS:
            return f"\\{text}"

        # Special statistical operators with LaTeX commands
        if text == 'Pr':
            return '\\Pr'
        elif text in self.STATISTICAL_OPERATORS:
            # CRITICAL: Do NOT automatically convert 'E' to \mathbb{E}!
            # Only apply \mathbb{} when there's explicit mathvariant="double-struck" attribute
            # The mi_visitor will handle mathvariant attributes properly.
            # In statistics mode, treat 'E' like other operators (upright but not blackboard bold)
            return f"\\operatorname{{{text}}}"

        # Distribution names (upright)
        if text in self.DISTRIBUTIONS:
            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 (Greek letters)
        if text in UNICODE_TO_LATEX:
            return UNICODE_TO_LATEX[text]

        # Handle multi-character strings
        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 (variables remain italic in LaTeX)
        return text

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

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

        Returns:
            Formatted LaTeX string

        Logic:
            1. Delimiters → __DELIM__ markers for later processing
            2. Unicode symbols → LaTeX equivalents (via UNICODE_TO_LATEX)
            3. Statistical symbols → proper notation (~, ∼, ≈, etc.)
            4. Standard operators → unchanged (+, -, *, /, =)

        Examples:
            "(" → "__DELIM__(__DELIM__"
            "~" → "\\sim" (distributed as)
            "∼" → "\\sim"
            "≈" → "\\approx"
            "∝" → "\\propto" (proportional to)
            "∈" → "\\in"
        """
        # CRITICAL: Mark delimiters for potential scaling
        # This matches MathStrategy behavior
        if text in ("(", ")", "[", "]", "{", "}"):
            return f"__DELIM__{text}__DELIM__"

        # Special case: pipe can be absolute value delimiter
        if text == "|":
            return f"__DELIM__|__DELIM__"

        # Handle < and > specifically (LaTeX convention)
        if text == "<":
            return "\\lt"
        if text == ">":
            return "\\gt"

        # Tilde (distributed as)
        if text == '~' or text == '∼':
            return '\\sim'

        # 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 statistics.

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

        Returns:
            Formatted LaTeX string

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

        Examples:
            "0.05" → "0.05" (p-value)
            "95" → "95" (confidence level)
        """
        # Numbers are returned as-is
        return text

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

        Args:
            equation_text: The full equation text

        Returns:
            True if likely statistics, False otherwise

        Detection indicators:
            - Statistical operators (Pr, E, Var, Cov)
            - Distribution notation (Normal, Binomial, etc.)
            - Greek letters commonly used in stats (μ, σ, β, α)
            - Tilde notation (X ~ Normal)
            - Statistical symbols (≈, ∝, ∼)

        Examples:
            "X ~ Normal(μ, σ²)" → True
            "E[X] = μ" → True
            "Var(X) = σ²" → True
            "sin(x) + cos(x)" → False
        """
        # Check for statistical operators
        op_count = sum(1 for op in self.STATISTICAL_OPERATORS if op in equation_text)
        if op_count >= 1:
            return True

        # Check for distribution names
        dist_count = sum(1 for dist in self.DISTRIBUTIONS if dist in equation_text)
        if dist_count >= 1:
            return True

        # Check for tilde notation (distributed as)
        if '~' in equation_text or '∼' in equation_text:
            return True

        # Check for common Greek letters in statistics context
        stats_greek = ['μ', 'σ', 'β', 'α', 'θ', 'λ', 'ρ', 'τ', 'χ']
        greek_count = sum(1 for letter in stats_greek if letter in equation_text)
        if greek_count >= 2:  # At least 2 Greek letters
            return True

        # Check for statistical patterns (p-value, confidence interval)
        import re
        if re.search(r'p\s*[<>=]?\s*0\.\d+', equation_text):  # p-value
            return True
        if re.search(r'\d+%\s*(CI|confidence)', equation_text, re.IGNORECASE):
            return True

        return False

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

    def __repr__(self) -> str:
        """String representation for debugging."""
        return f"StatisticsStrategy(operators={len(self.STATISTICAL_OPERATORS)}, distributions={len(self.DISTRIBUTIONS)})"