"""
Model summarization chain for extracting model descriptions and theories.
"""

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

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 ModelSummary:
    """Represents a summarized model or theory."""
    name: str
    assumptions: List[str] = field(default_factory=list)
    objective: str = ""
    formulation: str = ""
    training: str = ""
    properties: List[str] = field(default_factory=list)
    practical_notes: List[str] = field(default_factory=list)
    full_summary: str = ""
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "name": self.name,
            "assumptions": self.assumptions,
            "objective": self.objective,
            "formulation": self.formulation,
            "training": self.training,
            "properties": self.properties,
            "practical_notes": self.practical_notes,
            "full_summary": self.full_summary
        }


class ModelSummarizer:
    """
    Chain for summarizing models and theoretical frameworks.
    
    Extracts key components:
    - Assumptions
    - Objectives/Loss functions
    - Model formulation
    - Training/Optimization procedures
    - Key properties
    """
    
    def __init__(
        self,
        retrieval_service: RetrievalService,
        llm_client: LLMClient,
        include_assumptions: bool = True,
        include_objectives: bool = True,
        include_training: bool = True
    ):
        """
        Initialize model summarizer.
        
        Args:
            retrieval_service: Service for retrieving context
            llm_client: LLM for generation
            include_assumptions: Include model assumptions
            include_objectives: Include objectives/loss functions
            include_training: Include training procedures
        """
        self.retrieval_service = retrieval_service
        self.llm_client = llm_client
        self.include_assumptions = include_assumptions
        self.include_objectives = include_objectives
        self.include_training = include_training
    
    @trace_function(name="summarize_model")
    def summarize(
        self,
        model_name: str,
        queries: Optional[List[str]] = None,
        config: Optional[GenerationConfig] = None
    ) -> ModelSummary:
        """
        Summarize a model or theoretical framework.
        
        Args:
            model_name: Name of the model to summarize
            queries: Optional specific queries
            config: Generation configuration
            
        Returns:
            ModelSummary
        """
        config = config or GenerationConfig()
        
        # Build queries based on model name
        if queries is None:
            queries = [
                f"{model_name} model description",
                f"{model_name} assumptions",
                f"{model_name} objective loss function",
                f"{model_name} training optimization algorithm",
                f"{model_name} properties theoretical results",
            ]
        
        # Retrieve relevant content
        all_results = []
        for query in queries:
            results = self.retrieval_service.retrieve(query, top_k=3)
            all_results.extend(results)
        
        # Deduplicate results
        seen = set()
        unique_results = []
        for r in all_results:
            key = r.content[:200]
            if key not in seen:
                unique_results.append(r)
                seen.add(key)
        
        if not unique_results:
            logger.warning(f"No content found for model: {model_name}")
            return ModelSummary(name=model_name)
        
        # Generate summary
        summary = self._generate_summary(model_name, unique_results, config)
        
        logger.info(f"Summarized model: {model_name}")
        return summary
    
    def _generate_summary(
        self,
        model_name: str,
        results: List[RetrievalResult],
        config: GenerationConfig
    ) -> ModelSummary:
        """Generate model summary from retrieval results."""
        context = "\n\n".join(r.content for r in results)
        
        prompt = get_prompt("model_extraction")(
            context=context,
            model_name=model_name
        )
        
        result = self.llm_client.generate(prompt, config)
        
        # Parse the structured response
        return self._parse_response(model_name, result.text)
    
    def _parse_response(self, model_name: str, text: str) -> ModelSummary:
        """Parse model summarization response."""
        summary = ModelSummary(name=model_name, full_summary=text)
        
        current_section = None
        current_content = []
        
        for line in text.split("\n"):
            line_stripped = line.strip()
            
            # Detect section headers
            if line_stripped.startswith("### Assumptions"):
                if current_section and current_content:
                    self._assign_section(summary, current_section, current_content)
                current_section = "assumptions"
                current_content = []
            elif line_stripped.startswith("### Objective"):
                if current_section and current_content:
                    self._assign_section(summary, current_section, current_content)
                current_section = "objective"
                current_content = []
            elif line_stripped.startswith("### Formulation"):
                if current_section and current_content:
                    self._assign_section(summary, current_section, current_content)
                current_section = "formulation"
                current_content = []
            elif line_stripped.startswith("### Training") or line_stripped.startswith("### Optimization"):
                if current_section and current_content:
                    self._assign_section(summary, current_section, current_content)
                current_section = "training"
                current_content = []
            elif line_stripped.startswith("### Key Properties") or line_stripped.startswith("### Properties"):
                if current_section and current_content:
                    self._assign_section(summary, current_section, current_content)
                current_section = "properties"
                current_content = []
            elif line_stripped.startswith("### Practical"):
                if current_section and current_content:
                    self._assign_section(summary, current_section, current_content)
                current_section = "practical"
                current_content = []
            elif line_stripped and not line_stripped.startswith("##"):
                current_content.append(line_stripped)
        
        # Assign final section
        if current_section and current_content:
            self._assign_section(summary, current_section, current_content)
        
        return summary
    
    def _assign_section(
        self,
        summary: ModelSummary,
        section: str,
        content: List[str]
    ) -> None:
        """Assign content to the appropriate section."""
        if section == "assumptions" and self.include_assumptions:
            summary.assumptions = self._parse_list(content)
        elif section == "objective" and self.include_objectives:
            summary.objective = "\n".join(content)
        elif section == "formulation":
            summary.formulation = "\n".join(content)
        elif section == "training" and self.include_training:
            summary.training = "\n".join(content)
        elif section == "properties":
            summary.properties = self._parse_list(content)
        elif section == "practical":
            summary.practical_notes = self._parse_list(content)
    
    def _parse_list(self, content: List[str]) -> List[str]:
        """Parse content as a list of items."""
        items = []
        current_item = []
        
        for line in content:
            if line.startswith("- ") or line.startswith("* "):
                if current_item:
                    items.append(" ".join(current_item))
                current_item = [line[2:]]
            elif line.startswith(tuple(f"{i}." for i in range(1, 10))):
                if current_item:
                    items.append(" ".join(current_item))
                current_item = [line.split(".", 1)[1].strip()]
            elif current_item:
                current_item.append(line)
        
        if current_item:
            items.append(" ".join(current_item))
        
        return items if items else content
    
    @trace_function(name="extract_all_models")
    def extract_all(
        self,
        config: Optional[GenerationConfig] = None
    ) -> List[ModelSummary]:
        """
        Extract all models/theories from the document.
        
        Args:
            config: Generation configuration
            
        Returns:
            List of ModelSummary
        """
        config = config or GenerationConfig()
        
        # First, identify model names
        results = self.retrieval_service.retrieve(
            "model architecture method approach algorithm",
            top_k=10
        )
        
        # Extract potential model names from content
        model_names = self._identify_models(results)
        
        if not model_names:
            # Use generic approach
            model_names = ["Main Model"]
        
        # Summarize each model
        summaries = []
        for name in model_names[:5]:  # Limit to 5 models
            summary = self.summarize(name, config=config)
            if summary.objective or summary.formulation:
                summaries.append(summary)
        
        return summaries
    
    def _identify_models(self, results: List[RetrievalResult]) -> List[str]:
        """Identify model names from retrieval results."""
        import re
        
        model_names = set()
        
        # Patterns for model names
        patterns = [
            r'(\w+(?:\s+\w+)?)\s+(?:model|algorithm|method|approach)',
            r'(?:the|our)\s+(\w+(?:\s+\w+)?)\s+(?:model|algorithm)',
            r'(\w+(?:-\w+)?)\s+(?:network|architecture)',
        ]
        
        for result in results:
            text = result.content
            for pattern in patterns:
                for match in re.finditer(pattern, text, re.IGNORECASE):
                    name = match.group(1).strip()
                    if len(name) > 2 and name.lower() not in ["the", "our", "this", "that"]:
                        model_names.add(name)
        
        return list(model_names)[:10]
    
    def format_for_latex(self, summary: ModelSummary) -> str:
        """
        Format model summary for LaTeX.
        
        Args:
            summary: ModelSummary to format
            
        Returns:
            LaTeX formatted string
        """
        lines = [f"\\subsection{{{summary.name}}}"]
        
        if summary.assumptions:
            lines.append("\\paragraph{Assumptions}")
            lines.append("\\begin{itemize}")
            for assumption in summary.assumptions:
                lines.append(f"  \\item {assumption}")
            lines.append("\\end{itemize}")
        
        if summary.objective:
            lines.append("\\paragraph{Objective}")
            lines.append(summary.objective)
        
        if summary.formulation:
            lines.append("\\paragraph{Formulation}")
            lines.append(summary.formulation)
        
        if summary.training:
            lines.append("\\paragraph{Training}")
            lines.append(summary.training)
        
        if summary.properties:
            lines.append("\\paragraph{Key Properties}")
            lines.append("\\begin{itemize}")
            for prop in summary.properties:
                lines.append(f"  \\item {prop}")
            lines.append("\\end{itemize}")
        
        return "\n".join(lines)
