"""Scale up evaluation report mapping against evaluation frameworks using agentic workflows"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/06_mappr.ipynb.

# %% auto 0
__all__ = ['GEMINI_API_KEY', 'cfg', 'lm', 'find_section_path', 'get_content_tool', 'format_enabler_theme', 'Overview',
           'Exploration', 'Assessment', 'Synthesis', 'ThemeAnalyzer']

# %% ../nbs/06_mappr.ipynb 5
from pathlib import Path
from functools import reduce
from toolslm.md_hier import *
from rich import print
import json
from fastcore.all import *

from typing import List
import dspy

from .frameworks import EvalData, IOMEvalData

# %% ../nbs/06_mappr.ipynb 6
from dotenv import load_dotenv
import os

load_dotenv()
GEMINI_API_KEY = os.getenv('GEMINI_API_KEY')

# %% ../nbs/06_mappr.ipynb 7
cfg = AttrDict({
    'lm': 'gemini/gemini-2.0-flash-exp',
    'api_key': GEMINI_API_KEY,
    'max_tokens': 8192,
    'track_usage': False,
})

# %% ../nbs/06_mappr.ipynb 8
lm = dspy.LM(cfg.lm, api_key=cfg.api_key)
dspy.configure(lm=lm)

# %% ../nbs/06_mappr.ipynb 13
def find_section_path(
    hdgs: dict, # The nested dictionary structure
    target_section: str # The section name to find
) -> list: # The nested key path for the given section name
    "Find the nested key path for a given section name."
    def search_recursive(current_dict, path=[]):
        for key, value in current_dict.items():
            current_path = path + [key]
            if key == target_section:
                return current_path
            if isinstance(value, dict):
                result = search_recursive(value, current_path)
                if result:
                    return result
        return None
    
    return search_recursive(hdgs)

# %% ../nbs/06_mappr.ipynb 17
def get_content_tool(
    hdgs: dict, # The nested dictionary structure
    keys_list: list, # The list of keys to navigate through
    ) -> str: # The content of the section
    "Navigate through nested levels using the exact key strings."
    return reduce(lambda current, key: current[key], keys_list, hdgs).text

# %% ../nbs/06_mappr.ipynb 21
def format_enabler_theme(
    theme: EvalData # The theme object
    ) -> str: # The formatted theme string
    "Format SRF enabler into structured text for LM processing."
    parts = [
        f'## Enabler {theme.id}: {theme.title}',
        '### Description', 
        theme.description
    ]
    return '\n'.join(parts)

# %% ../nbs/06_mappr.ipynb 26
class Overview(dspy.Signature):
    "Based on framework theme to map and report's TOC determine the sections to explore first."
    theme: str = dspy.InputField(desc="Theme being analyzed")
    all_headings: str = dspy.InputField(desc="Complete document structure")
    priority_sections: List[str] = dspy.OutputField(desc="Ordered list of section keys to explore first")
    strategy: str = dspy.OutputField(desc="Reasoning for this exploration strategy")

# %% ../nbs/06_mappr.ipynb 29
class Exploration(dspy.Signature):
    "Decide next exploration step for theme to be mapped based on current findings and available sections."
    theme: str = dspy.InputField(desc="Theme being analyzed")
    current_findings: str = dspy.InputField(desc="Evidence found so far")
    available_sections: str = dspy.InputField(desc="Remaining sections to explore")
    next_section: str = dspy.OutputField(desc="Next section key to explore, or 'DONE' if sufficient")
    reasoning: str = dspy.OutputField(desc="Why this section or why stopping")

# %% ../nbs/06_mappr.ipynb 31
class Assessment(dspy.Signature):
    "Assess if current evidence is sufficient for theme analysis."
    theme: str = dspy.InputField(desc="Theme being analyzed")
    evidence_so_far: str = dspy.InputField(desc="All evidence collected")
    sections_explored: str = dspy.InputField(desc="Sections already checked")
    sufficient: bool = dspy.OutputField(desc="Is evidence sufficient to make conclusion?")
    confidence_score: float = dspy.OutputField(desc="Confidence in current findings (0-1)")
    next_priority: str = dspy.OutputField(desc="If continuing, what type of section to prioritize")

# %% ../nbs/06_mappr.ipynb 32
class Synthesis(dspy.Signature):
    "Provide detailed rationale and synthesis of theme analysis."
    theme: str = dspy.InputField(desc="Theme being analyzed")
    all_evidence: str = dspy.InputField(desc="All collected evidence")
    sections_explored: str = dspy.InputField(desc="List of sections explored")
    theme_covered: bool = dspy.OutputField(desc="Final decision on theme coverage")
    confidence_explanation: str = dspy.OutputField(desc="Detailed explanation of confidence score")
    evidence_summary: str = dspy.OutputField(desc="Key evidence supporting the conclusion")
    gaps_identified: str = dspy.OutputField(desc="Any gaps or missing aspects")

# %% ../nbs/06_mappr.ipynb 35
class ThemeAnalyzer(dspy.Module):
    """
    Analyzes a theme across a document by iteratively exploring sections, collecting evidence, and synthesizing findings. 
    Uses a structured pipeline of overview -> exploration -> assessment -> synthesis.
    """
    def __init__(self, overview_sig, exploration_sig, assessment_sig, synthesis_sig, max_iter=10):
        self.overview = dspy.ChainOfThought(overview_sig)
        self.explore = dspy.ChainOfThought(exploration_sig)
        self.assess = dspy.ChainOfThought(assessment_sig)
        self.synthesize = dspy.ChainOfThought(synthesis_sig)
        self.max_iter = max_iter

# %% ../nbs/06_mappr.ipynb 36
@patch
def forward(
    self:ThemeAnalyzer, 
    theme: str, # The formatted theme to analyze
    headings: dict, # The headings TOC of the document
    get_content_fn=get_content_tool # The function to get the content of a section using `hdgs[keys_list].text` for instance
    ) -> Synthesis: # Synthesized analysis results including theme coverage, confidence, evidence and gaps
    "Executes a structured analysis process."
    priority_sections = self.get_overview(theme, headings)
    evidence = self.explore_iteratively(theme, priority_sections, headings, get_content_fn)
    return self.synthesize_findings(theme, evidence)

# %% ../nbs/06_mappr.ipynb 37
@patch
def get_overview(self:ThemeAnalyzer, theme, headings) -> Overview:
    overview = self.overview(theme=theme, all_headings=str(headings))
    print("Overview priority sections:", overview.priority_sections)
    print("Overview strategy:", overview.strategy)
    return overview.priority_sections


# %% ../nbs/06_mappr.ipynb 38
@patch
def explore_iteratively(self:ThemeAnalyzer, theme, priority_sections, headings, get_content_fn):
    evidence_collected = []
    sections_explored = []
    available_sections = priority_sections.copy()
    
    for i in range(self.max_iter):
        print(f"\n--- Iteration {i+1} ---")
        if not available_sections:
            print("No more sections to explore, stopping")
            break
        if self.should_stop_exploring(theme, evidence_collected, sections_explored):
            break
        decision = self.make_exploration_decision(theme, evidence_collected, available_sections)
        if decision.next_section == 'DONE':
            print("Decision says DONE, breaking")
            break
        evidence_collected, sections_explored = self.process_section(decision, 
                                                                     headings, 
                                                                     get_content_fn, 
                                                                     evidence_collected, 
                                                                     sections_explored, 
                                                                     available_sections)
    
    return {"evidence": evidence_collected, "sections": sections_explored}


# %% ../nbs/06_mappr.ipynb 39
@patch
def make_exploration_decision(self:ThemeAnalyzer, theme, evidence_collected, available_sections):
    decision = self.explore(
        theme=theme,
        current_findings="\n\n".join(evidence_collected) if evidence_collected else "No evidence collected yet",
        available_sections=str(available_sections)
    )
    print("Decision:", decision.next_section)
    print("Reasoning:", decision.reasoning)
    return decision


# %% ../nbs/06_mappr.ipynb 40
@patch
def should_stop_exploring(self:ThemeAnalyzer, theme, evidence_collected, sections_explored):
    if not evidence_collected:
        return False
        
    assessment = self.assess(
        theme=theme,
        evidence_so_far="\n\n".join(evidence_collected),
        sections_explored=str(sections_explored)
    )
    print("Assessment - Sufficient:", assessment.sufficient, "Confidence:", assessment.confidence_score)
    
    return assessment.sufficient and assessment.confidence_score > 0.8


# %% ../nbs/06_mappr.ipynb 41
@patch
def process_section(self:ThemeAnalyzer, decision, headings, get_content_fn, evidence_collected, sections_explored, available_sections):
    path = find_section_path(headings, decision.next_section)
    print("Path found:", path)
    
    if path:
        content = get_content_fn(headings, path)
        print("Content length:", len(content))
        evidence_collected.append(f"# Section: {decision.next_section}\n## Content\n{content}")
        sections_explored.append(decision.next_section)
        if decision.next_section in available_sections:
            available_sections.remove(decision.next_section)
    else:
        print("No path found for section!")
    
    return evidence_collected, sections_explored


# %% ../nbs/06_mappr.ipynb 42
@patch
def synthesize_findings(self:ThemeAnalyzer, theme, evidence):
    synthesis = self.synthesize(
        theme=theme,
        all_evidence="\n\n".join(evidence["evidence"]),
        sections_explored=str(evidence["sections"])
    )
    print("Synthesis result:", synthesis.theme_covered)
    print("Synthesis reasoning:", synthesis.confidence_explanation)
    print("Synthesis evidence:", synthesis.evidence_summary)
    print("Synthesis gaps:", synthesis.gaps_identified)
    return synthesis

