"""
Formula extraction chain for identifying and organizing mathematical formulas.
"""

import re
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple

from ..core.retrieval import RetrievalService, RetrievalResult
from ..core.models import LLMClient, GenerationConfig
from ..core.prompts import get_prompt
from ..infra.logging import get_logger
from ..infra.tracing import trace_function

logger = get_logger(__name__)


@dataclass
class Formula:
    """Represents an extracted formula."""
    name: str
    latex: str
    variables: Dict[str, str] = field(default_factory=dict)
    context: str = ""
    depends_on: List[str] = field(default_factory=list)
    topic: str = ""
    importance_score: float = 0.0
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "name": self.name,
            "latex": self.latex,
            "variables": self.variables,
            "context": self.context,
            "depends_on": self.depends_on,
            "topic": self.topic,
            "importance_score": self.importance_score
        }
    
    @property
    def clean_latex(self) -> str:
        """Get LaTeX without surrounding delimiters."""
        latex = self.latex.strip()
        
        # Remove common delimiters
        for start, end in [("$$", "$$"), ("$", "$"), (r"\[", r"\]"), (r"\(", r"\)")]:
            if latex.startswith(start) and latex.endswith(end):
                latex = latex[len(start):-len(end)].strip()
                break
        
        return latex


class FormulaExtractor:
    """
    Chain for extracting mathematical formulas from documents.
    
    Uses RAG to identify formulas and normalize them to
    consistent LaTeX format.
    """
    
    # Patterns for extracting formulas from text
    FORMULA_PATTERNS = [
        r'\$\$([^$]+)\$\$',                             # Display math
        r'\$([^$]+)\$',                                  # Inline math
        r'\\begin\{equation\}(.*?)\\end\{equation\}',   # equation env
        r'\\begin\{align\*?\}(.*?)\\end\{align\*?\}',   # align env
        r'\\\[(.*?)\\\]',                                # \[ \] display
        r'\\\((.*?)\\\)',                                # \( \) inline
    ]
    
    def __init__(
        self,
        retrieval_service: RetrievalService,
        llm_client: LLMClient,
        max_formulas: int = 100,
        normalize_latex: bool = True,
        include_derivations: bool = False,
        group_by_topic: bool = True
    ):
        """
        Initialize formula extractor.
        
        Args:
            retrieval_service: Service for retrieving context
            llm_client: LLM for generation
            max_formulas: Maximum formulas to extract
            normalize_latex: Normalize LaTeX syntax
            include_derivations: Include derivation steps
            group_by_topic: Group formulas by topic
        """
        self.retrieval_service = retrieval_service
        self.llm_client = llm_client
        self.max_formulas = max_formulas
        self.normalize_latex = normalize_latex
        self.include_derivations = include_derivations
        self.group_by_topic = group_by_topic
    
    @trace_function(name="extract_formulas")
    def extract(
        self,
        queries: Optional[List[str]] = None,
        config: Optional[GenerationConfig] = None
    ) -> List[Formula]:
        """
        Extract formulas from the indexed document.
        
        Args:
            queries: Optional specific queries for formula retrieval
            config: Generation configuration
            
        Returns:
            List of extracted formulas
        """
        config = config or GenerationConfig()
        
        # Default queries for formula extraction
        if queries is None:
            queries = [
                "equations formulas mathematical expressions",
                "key equations derivations",
                "loss function objective",
                "theorem lemma proposition corollary",
                "algorithm update rule",
            ]
        
        all_formulas = []
        seen_latex = set()
        
        for query in queries:
            # Retrieve math-heavy content
            results = self.retrieval_service.retrieve_math(query, top_k=8)
            
            if not results:
                # Fall back to general retrieval
                results = self.retrieval_service.retrieve(query, top_k=5)
            
            if not results:
                continue
            
            # Extract formulas from context
            formulas = self._extract_from_context(results, config)
            
            # Deduplicate by LaTeX content
            for formula in formulas:
                normalized = self._normalize_latex_key(formula.latex)
                if normalized not in seen_latex:
                    all_formulas.append(formula)
                    seen_latex.add(normalized)
        
        # Sort by importance
        all_formulas.sort(key=lambda f: f.importance_score, reverse=True)
        
        logger.info(f"Extracted {len(all_formulas)} formulas")
        return all_formulas[:self.max_formulas]
    
    def _extract_from_context(
        self,
        results: List[RetrievalResult],
        config: GenerationConfig
    ) -> List[Formula]:
        """Extract formulas from retrieval results."""
        context = "\n\n".join(r.content for r in results)
        
        # First, try to extract formulas directly from text
        direct_formulas = self._extract_direct(context)
        
        # Then use LLM for better organization and context
        prompt = get_prompt("formula_extraction")(context=context)
        result = self.llm_client.generate(prompt, config)
        
        # Parse LLM response
        llm_formulas = self._parse_response(result.text)
        
        # Merge direct and LLM-extracted formulas
        all_formulas = self._merge_formulas(direct_formulas, llm_formulas)
        
        return all_formulas
    
    def _extract_direct(self, text: str) -> List[Formula]:
        """Extract formulas directly using regex patterns."""
        formulas = []
        
        for pattern in self.FORMULA_PATTERNS:
            for match in re.finditer(pattern, text, re.DOTALL):
                latex = match.group(0)
                
                # Clean up the latex
                content = match.group(1) if match.lastindex else latex
                
                formulas.append(Formula(
                    name="",
                    latex=latex.strip(),
                    importance_score=0.3
                ))
        
        return formulas
    
    def _parse_response(self, text: str) -> List[Formula]:
        """Parse formula extraction response from LLM."""
        formulas = []
        current: Dict[str, Any] = {}
        in_variables = False
        
        for line in text.split("\n"):
            line = line.strip()
            
            if line.startswith("FORMULA:"):
                if current and "latex" in current:
                    formulas.append(self._create_formula(current))
                current = {"name": line[8:].strip(), "variables": {}}
                in_variables = False
            elif line.startswith("LATEX:"):
                current["latex"] = line[6:].strip()
            elif line.startswith("VARIABLES:"):
                in_variables = True
            elif in_variables and line.startswith("- "):
                # Parse variable definition
                var_line = line[2:].strip()
                if ":" in var_line:
                    var_name, var_desc = var_line.split(":", 1)
                    current["variables"][var_name.strip()] = var_desc.strip()
            elif line.startswith("CONTEXT:"):
                current["context"] = line[8:].strip()
                in_variables = False
            elif line.startswith("DEPENDS_ON:"):
                deps = line[11:].strip()
                current["depends_on"] = [d.strip() for d in deps.split(",") if d.strip()]
                in_variables = False
        
        # Don't forget the last formula
        if current and "latex" in current:
            formulas.append(self._create_formula(current))
        
        return formulas
    
    def _create_formula(self, data: Dict[str, Any]) -> Formula:
        """Create a Formula from parsed data."""
        # Calculate importance score
        score = 0.5  # Base score for LLM-extracted formulas
        
        if data.get("name"):
            score += 0.2
        if data.get("context"):
            score += 0.2
        if data.get("variables"):
            score += 0.1
        
        latex = data.get("latex", "")
        if self.normalize_latex:
            latex = self._normalize_latex(latex)
        
        return Formula(
            name=data.get("name", ""),
            latex=latex,
            variables=data.get("variables", {}),
            context=data.get("context", ""),
            depends_on=data.get("depends_on", []),
            importance_score=score
        )
    
    def _normalize_latex(self, latex: str) -> str:
        """Normalize LaTeX syntax for consistency."""
        # Ensure proper spacing
        latex = re.sub(r'\s+', ' ', latex)
        
        # Normalize common patterns
        replacements = [
            (r'\\frac(\S)', r'\\frac{\1}'),  # \frac1 -> \frac{1}
            (r'\\sqrt(\S)', r'\\sqrt{\1}'),   # \sqrt2 -> \sqrt{2}
        ]
        
        for pattern, replacement in replacements:
            latex = re.sub(pattern, replacement, latex)
        
        return latex.strip()
    
    def _normalize_latex_key(self, latex: str) -> str:
        """Create a normalized key for deduplication."""
        # Remove whitespace and common delimiters
        key = re.sub(r'\s+', '', latex)
        key = re.sub(r'[\$\\\[\]\(\)]', '', key)
        return key.lower()
    
    def _merge_formulas(
        self,
        direct: List[Formula],
        llm: List[Formula]
    ) -> List[Formula]:
        """Merge directly extracted and LLM-extracted formulas."""
        # Prefer LLM formulas as they have better context
        seen_keys = set()
        merged = []
        
        for formula in llm:
            key = self._normalize_latex_key(formula.latex)
            seen_keys.add(key)
            merged.append(formula)
        
        # Add direct formulas that weren't found by LLM
        for formula in direct:
            key = self._normalize_latex_key(formula.latex)
            if key not in seen_keys and len(key) > 5:  # Skip very short formulas
                merged.append(formula)
                seen_keys.add(key)
        
        return merged
    
    def organize_by_topic(
        self,
        formulas: List[Formula],
        topics: Optional[List[str]] = None,
        config: Optional[GenerationConfig] = None
    ) -> Dict[str, List[Formula]]:
        """
        Organize formulas by topic.
        
        Args:
            formulas: Formulas to organize
            topics: Optional topic categories
            config: Generation configuration
            
        Returns:
            Dictionary mapping topic to formulas
        """
        if not self.group_by_topic:
            return {"All Formulas": formulas}
        
        # Default topics
        if topics is None:
            topics = [
                "Definitions",
                "Loss Functions",
                "Update Rules",
                "Theorems",
                "Bounds",
                "Other"
            ]
        
        # Simple keyword-based grouping
        groups: Dict[str, List[Formula]] = {topic: [] for topic in topics}
        
        for formula in formulas:
            assigned = False
            name_lower = formula.name.lower()
            context_lower = formula.context.lower()
            
            # Match by keywords
            if any(kw in name_lower or kw in context_lower for kw in ["loss", "objective", "cost"]):
                groups["Loss Functions"].append(formula)
                assigned = True
            elif any(kw in name_lower or kw in context_lower for kw in ["update", "gradient", "step"]):
                groups["Update Rules"].append(formula)
                assigned = True
            elif any(kw in name_lower or kw in context_lower for kw in ["theorem", "lemma", "proposition"]):
                groups["Theorems"].append(formula)
                assigned = True
            elif any(kw in name_lower or kw in context_lower for kw in ["bound", "inequality"]):
                groups["Bounds"].append(formula)
                assigned = True
            elif any(kw in name_lower or kw in context_lower for kw in ["definition", "define"]):
                groups["Definitions"].append(formula)
                assigned = True
            
            if not assigned:
                groups["Other"].append(formula)
        
        # Remove empty groups
        return {k: v for k, v in groups.items() if v}
    
    def format_for_latex(
        self,
        formulas: List[Formula],
        environment: str = "align"
    ) -> str:
        """
        Format formulas for LaTeX document.
        
        Args:
            formulas: Formulas to format
            environment: LaTeX environment to use
            
        Returns:
            LaTeX formatted string
        """
        lines = [f"\\begin{{{environment}*}}"]
        
        for formula in formulas:
            latex = formula.clean_latex
            
            # Add label comment
            if formula.name:
                lines.append(f"  % {formula.name}")
            
            lines.append(f"  {latex} \\\\")
        
        lines.append(f"\\end{{{environment}*}}")
        
        return "\n".join(lines)
