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

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/08_mappr_lisette.ipynb.

# %% auto 0
__all__ = ['GEMINI_API_KEY', 'cfg', 'traces_dir', 'lm', 'select_section_sp', 'summarize_sp', 'evaluate_evidence_sp',
           'find_section_path', 'get_content_tool', 'flatten_sections', 'extract_content', 'format_enabler_theme',
           'format_crosscutting_theme', 'format_gcm_theme', 'format_srf_output', 'SelectSectionOutput',
           'SummarizeContentOutput', 'EvaluateEvidenceOutput', 'State', 'parse_response',
           'format_sections_for_selection', 'select_section', 'summarize_content', 'evaluate_evidence', 'limit',
           'Stage', 'TraceContext', 'setup_logger', 'setup_trace_logging', 'log_analysis_event', 'AnalysisResult',
           'PipelineResults', 'PipelineOrchestrator', 'get_stage1_covered_context', 'get_filtered_srf_output_ids',
           'get_combined_context']

# %% ../nbs/08_mappr_lisette.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 enum import Enum
import logging
import uuid
from datetime import datetime
from typing import List, Callable
import dspy
from asyncio import Semaphore, gather, sleep
import time
from collections import defaultdict
import copy

from pydantic import BaseModel, Field
from typing import List

from .frameworks import (EvalData, 
                                 IOMEvalData, 
                                 FrameworkInfo, 
                                 Framework,
                                 FrameworkCat,
                                 find_srf_output_by_id)

#from evaluatr.db_traces import TraceDB, Trace
from fastlite import database

from lisette import Chat, AsyncChat
import json

# %% ../nbs/08_mappr_lisette.ipynb 6
from dotenv import load_dotenv
import os

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

# %% ../nbs/08_mappr_lisette.ipynb 7
cfg = AttrDict({
    'lm': 'gemini/gemini-2.0-flash',
    'api_key': GEMINI_API_KEY,
    'max_tokens': 8192,
    'track_usage': False,
    'call_delay': 0.1, # in seconds
    'semaphore': 30,
    'dirs': AttrDict({
        'data': '.evaluatr',
        'trace': 'traces'
    }),
    'verbosity': 1,
    'cache': AttrDict({
        'is_active': False,
        'delay': 0.05 # threshold in seconds below which we consider the response is cached
    }),
    'max_iter': 10
})

# %% ../nbs/08_mappr_lisette.ipynb 8
traces_dir = Path.home() / cfg.dirs.data / cfg.dirs.trace
traces_dir.mkdir(parents=True, exist_ok=True)

# %% ../nbs/08_mappr_lisette.ipynb 12
lm = dspy.LM(cfg.lm, api_key=cfg.api_key, cache=cfg.cache.is_active)
dspy.configure(lm=lm)

# %% ../nbs/08_mappr_lisette.ipynb 17
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/08_mappr_lisette.ipynb 21
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/08_mappr_lisette.ipynb 23
def flatten_sections(hdgs, path=[]):
    """Extract flat list of (key, full_path) tuples from nested hdgs"""
    sections = []
    for key, value in hdgs.items():
        current_path = path + [key]
        sections.append((key, current_path))
        if isinstance(value, dict):
            sections.extend(flatten_sections(value, current_path))
    return sections

# %% ../nbs/08_mappr_lisette.ipynb 24
def extract_content(section_key: str, sections_lookup: dict, hdgs: dict) -> str:
    path = sections_lookup[section_key]
    return get_content_tool(hdgs, path)

# %% ../nbs/08_mappr_lisette.ipynb 27
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/08_mappr_lisette.ipynb 30
def format_crosscutting_theme(
    theme: EvalData # The theme object
    ) -> str: # The formatted theme string
    "Format SRF cross-cutting into structured text for LM processing."
    parts = [
        f'## Cross-cutting {theme.id}: {theme.title}',
        '### Description', 
        theme.description
    ]
    return '\n'.join(parts)

# %% ../nbs/08_mappr_lisette.ipynb 33
def format_gcm_theme(
    theme: dict # The GCM theme object from gcm_small
    ) -> str: # The formatted theme string
    "Format GCM objective into structured text for LM processing."
    parts = [
        f'## GCM Objective {theme["id"]}: {theme["title"]}',
        '### Core Theme', 
        theme["core_theme"]
    ]
    
    if theme.get("key_principles"):
        parts.extend(['### Key Principles', ', '.join(theme["key_principles"])])
    
    if theme.get("target_groups"):
        parts.extend(['### Target Groups', ', '.join(theme["target_groups"])])
        
    if theme.get("main_activities"):
        parts.extend(['### Main Activities', ', '.join(theme["main_activities"])])
    
    return '\n'.join(parts)

# %% ../nbs/08_mappr_lisette.ipynb 36
def format_srf_output(output_context: dict) -> str:
    "Format SRF output with full hierarchical context for LM processing."
    parts = [
        f'## SRF Output {output_context["output"]["id"]}: {output_context["output"]["title"]}',
        '### Strategic Context',
        f'**Objective {output_context["objective"]["id"]}**: {output_context["objective"]["title"]}',
        f'**Long    -term Outcome {output_context["long_outcome"]["id"]}**: {output_context["long_outcome"]["title"]}',
        f'**Short-term Outcome {output_context["short_outcome"]["id"]}**: {output_context["short_outcome"]["title"]}'
    ]
    
    return '\n'.join(parts)

# %% ../nbs/08_mappr_lisette.ipynb 40
class SelectSectionOutput(BaseModel):
    "Select the next most relevant section based on current evidence summary and gaps"
    section_index: int
    reasoning: str

# %% ../nbs/08_mappr_lisette.ipynb 41
class SelectSectionOutput(BaseModel):
    "Select the next most relevant section based on current evidence summary and gaps"
    section_index: int
    reasoning: str

# %% ../nbs/08_mappr_lisette.ipynb 42
class SummarizeContentOutput(BaseModel):
    "Summarize the content of a section and identify the key findings"
    summary: str
    key_findings: List[str]

# %% ../nbs/08_mappr_lisette.ipynb 43
class EvaluateEvidenceOutput(BaseModel):
    theme_covered: bool
    coverage_reasoning: str
    gaps_identified: str
    should_continue: bool

# %% ../nbs/08_mappr_lisette.ipynb 44
class State(BaseModel):
    theme: str
    prior_coverage_context: str = ""
    section_summaries: List[dict] = []  # Renamed from evidences
    explored_sections: List[str] = []
    available_sections: List[str]
    evaluation_history: List[dict] = []  # Track reasoning evolution
    iterations_completed: int = 0
    theme_covered: bool = False
    coverage_reasoning: str = ""
    gaps_identified: str = ""
    stop_reason: str = ""

# %% ../nbs/08_mappr_lisette.ipynb 46
select_section_sp = """### ROLE AND OBJECTIVE
You are a strategic document navigator. Your job is to select the next most relevant section to explore for evidence about a specific theme.

### CONTEXT
You will receive:
- The theme being analyzed
- Current evidence summary and identified gaps
- Available sections (as JSON array of [index, section_name] pairs)
- Sections already explored

### TASK INSTRUCTIONS
Select the next section most likely to contain NEW, relevant evidence about this theme.

### SELECTION RULES
1. Choose sections that directly address the theme (not tangentially related)
2. NEVER select a subsection if its parent section was already explored
   - Example: If "4. Findings" explored → skip "4.1 Relevance", "4.1.1 Needs", etc.
3. Prioritize sections that address gaps identified in previous evaluations
4. Select "DONE" only when no unexplored sections remain OR all remaining sections are clearly irrelevant

### EXAMPLES
**Example 1: Parent → Child is FORBIDDEN**
Explored: ["4. Findings"]
Available: ["4.1 Relevance", "4.2 Coherence", "5. Conclusions"]
✓ Valid: "5. Conclusions" (independent section)
✗ Invalid: "4.1 Relevance", "4.2 Coherence" (children of explored parent "4. Findings")

**Example 2: Child → Parent is ALLOWED**
Explored: ["4.1.1 Needs of migrants"]
Available: ["4.1 Relevance", "4. Findings", "5. Conclusions"]
✓ Valid: All three are valid (parent can add context after exploring child)

**Example 3: Siblings are always OK**
Explored: ["4.1 Relevance"]
Available: ["4.2 Coherence", "4.3 Effectiveness"]
✓ Valid: Both are OK (siblings are independent)

### OUTPUT FORMAT
JSON with:
- section_index: integer from the provided pairs, or "DONE"
- reasoning: explain why this section addresses current gaps and is not a subsection of explored content
"""

# %% ../nbs/08_mappr_lisette.ipynb 47
summarize_sp = """You are summarizing content from an evaluation report section.

Your task: Extract and condense the key points relevant to the theme being analyzed.

This summary will be used to:
- Maintain context across iterations without inflating the prompt
- Provide the evaluation step with essential information from this section

Keep it concise but capture:
- Main findings or claims related to the theme
- Supporting evidence (data, quotes, examples)
- Methodological details if relevant

Output JSON with:
- summary: concise summary of the content
- key_findings: list of specific findings relevant to the theme
"""

# %% ../nbs/08_mappr_lisette.ipynb 48
evaluate_evidence_sp = """### ROLE AND OBJECTIVE
You are a senior UN evaluation expert. Your task: Identify what's MISSING from this report that would prevent you from writing a substantive briefing on this theme.

### YOUR TASK
Assume you need to write a 2-page briefing for leadership on this theme covering:
- What the program achieved (or didn't achieve)
- Why it succeeded or failed  
- What lessons emerged
- Evidence-based recommendations for future programs

**Start by identifying gaps**: What critical information is missing?

### CRITICAL GAPS TO CHECK
For each, ask: "Is this gap present?"

1. **Causal understanding gap**: Can't explain WHY outcomes occurred?
2. **Outcome evidence gap**: Have activities but no results/impact data?
3. **Context gap**: Missing information on challenges/barriers faced?
4. **Specificity gap**: Only vague statements, no concrete examples/numbers?
5. **Balance gap**: Only successes OR only failures, not both?

### DECISION RULE
**If ANY critical gap exists → theme_covered=False**

Only mark theme_covered=True when:
- All 5 critical gaps have been ruled out
- You have enough evidence to write a substantive 2-page briefing
- A senior evaluator would find your briefing credible and actionable

When in doubt → mark as False

### OUTPUT FORMAT
JSON with:
- theme_covered: boolean
- coverage_reasoning: Start with gaps identified, then explain if evidence overcomes them
- gaps_identified: List specific critical gaps that remain
- should_continue: boolean
"""


# %% ../nbs/08_mappr_lisette.ipynb 50
def parse_response(result):
    "Extract JSON from Lisette response"
    return json.loads(result.choices[0].message.content)

# %% ../nbs/08_mappr_lisette.ipynb 51
def format_sections_for_selection(available_sections: List[str]) -> str:
    "Format available sections as indexed JSON array"
    return json.dumps([
        [i+1, s] for i, s in enumerate(available_sections)
    ], indent=2)

# %% ../nbs/08_mappr_lisette.ipynb 52
async def select_section(
    state: State,
    model: str = 'gemini/gemini-2.0-flash'
) -> dict:
    "Select next section to explore based on current state"
    chat = AsyncChat(model=model, sp=select_section_sp, temp=0)
    
    sections_json = format_sections_for_selection(state.available_sections)
    
    # Format evaluation history for context
    eval_summary = "\n".join([
        f"Iteration {ev['iteration']}: Theme covered={ev['theme_covered']}, Gaps: {ev['gaps_identified']}"
        for ev in state.evaluation_history
    ]) if state.evaluation_history else "No evaluations yet - initial exploration"
    
    parts = [
        state.prior_coverage_context,
        f"Theme being analyzed:\n{state.theme}",
        f"Evaluation history:\n{eval_summary}",
        f"Available sections:\n{sections_json}",
        f"Explored sections: {state.explored_sections}"
    ]
    prompt = "\n\n".join(p for p in parts if p)
    
    result = await chat(prompt, response_format=SelectSectionOutput)
    parsed = parse_response(result)
    
    section_key = state.available_sections[parsed['section_index'] - 1]
    return {'section_key': section_key, 'reasoning': parsed['reasoning']}

# %% ../nbs/08_mappr_lisette.ipynb 58
async def summarize_content(
    state: State, # The current state of the analysis
    section_key: str, # The key of the section to summarize
    content: str, # The content of the section to summarize
    # model: str = 'gemini/gemini-2.0-flash' # The model to use
    # model: str = 'claude-sonnet-4-20250514'
    model: str = 'gemini/gemini-2.5-flash'
) -> dict:
    "Summarize section content relevant to the theme"
    chat = AsyncChat(model=model, sp=summarize_sp, temp=0)
    
    parts = [
        state.prior_coverage_context,
        f"Theme: {state.theme}",
        f"Section: {section_key}",
        f"Content to summarize:\n{content}"
    ]
    prompt = "\n\n".join(p for p in parts if p)
    
    result = await chat(prompt, response_format=SummarizeContentOutput)
    return parse_response(result)

# %% ../nbs/08_mappr_lisette.ipynb 63
async def evaluate_evidence(
    state: State,
    new_content: str,
    # model: str = 'gemini/gemini-2.0-flash'
    model: str = 'gemini/gemini-2.5-flash'
    # model: str = 'claude-sonnet-4-20250514'
) -> dict:
    "Evaluate evidence collected and determine if more exploration needed"
    chat = AsyncChat(model=model, sp=evaluate_evidence_sp, temp=0)
    
    # Format previous summaries
    prev_summaries = "\n\n".join([
        f"Section: {state.explored_sections[i]}\n"
        f"Summary: {state.section_summaries[i]['summary']}\n"
        f"Key findings: {', '.join(state.section_summaries[i]['key_findings'])}"
        for i in range(len(state.section_summaries))
    ]) if state.section_summaries else "None yet"
    
    # Format evaluation history
    prev_evaluations = "\n\n".join([
        f"Iteration {ev['iteration']}:\n"
        f"Theme covered: {ev['theme_covered']}\n"
        f"Reasoning: {ev['coverage_reasoning']}\n"
        f"Gaps: {ev['gaps_identified']}"
        for ev in state.evaluation_history
    ]) if state.evaluation_history else "First evaluation - no previous assessments"
    
    parts = [
        state.prior_coverage_context,
        f"Theme being analyzed:\n{state.theme}",
        f"Previous evidence summaries:\n{prev_summaries}",
        f"Previous evaluation reasoning:\n{prev_evaluations}",
        f"Exploration progress: {len(state.explored_sections)} sections explored out of {len(state.explored_sections) + len(state.available_sections)} total available",
        f"New content to evaluate:\n{new_content}"
    ]
    prompt = "\n\n".join(p for p in parts if p)
    
    result = await chat(prompt, response_format=EvaluateEvidenceOutput)
    return parse_response(result)

# %% ../nbs/08_mappr_lisette.ipynb 69
async def limit(semaphore, coro, delay=None):
    "Execute coroutine with semaphore concurrency control"
    async with semaphore:
        result = await coro
        if delay: await sleep(delay)
        return result

# %% ../nbs/08_mappr_lisette.ipynb 76
class Stage(Enum):
    "Pipeline stage number."
    STAGE1 = "stage1"
    STAGE2 = "stage2"
    STAGE3 = "stage3"
    def __str__(self): return self.value

# %% ../nbs/08_mappr_lisette.ipynb 78
class TraceContext(AttrDict):
    "Context for tracing the mapping process"
    def __init__(self, 
                 report_id:str,  # Report identifier
                 stage:Stage,  # Pipeline stage number
                 framework:FrameworkInfo,  # Framework info (name, category, theme_id)
                 ): 
        # self.run_id = str(uuid.uuid4())[:8]  # Short unique ID
        store_attr()

    def __repr__(self):
        return f"TraceContext(report_id={self.report_id}, stage={self.stage}, framework={self.framework})"


# %% ../nbs/08_mappr_lisette.ipynb 80
def setup_logger(name, handler, level=logging.INFO, **kwargs):
    "Helper function to setup a logger with common configuration"
    logger = logging.getLogger(name)
    logger.handlers.clear()
    logger.addHandler(handler)
    logger.setLevel(level)
    for k,v in kwargs.items(): setattr(logger, k, v)
    return logger

# %% ../nbs/08_mappr_lisette.ipynb 81
def setup_trace_logging(report_id, verbosity=cfg.verbosity):
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f'{report_id}_{timestamp}.jsonl'
    file_handler = logging.FileHandler(traces_dir / filename, mode='w')
    setup_logger('trace.file', file_handler)    
    console_handler = logging.StreamHandler()
    setup_logger('trace.console', console_handler, verbosity=verbosity)

# %% ../nbs/08_mappr_lisette.ipynb 82
def log_analysis_event(event: str, trace_ctx: TraceContext, **extra_data):
    """Log an analysis event to file and console with different verbosity levels"""
    file_logger = logging.getLogger('trace.file')
    console_logger = logging.getLogger('trace.console')
    
    base_data = {
        "timestamp": datetime.now().isoformat(),
        "event": event,
        "report_id": trace_ctx.report_id,
        "stage": str(trace_ctx.stage),
        "framework": str(trace_ctx.framework.name),
        "framework_category": str(trace_ctx.framework.category),
        "framework_theme_id": str(trace_ctx.framework.theme_id),
    }
    base_data.update(extra_data)
    
    # File logger - always full JSON
    file_logger.info(json.dumps(base_data, indent=2))
    
    # Console logger - verbosity-based formatting
    if hasattr(console_logger, 'verbosity'):
        if console_logger.verbosity == 1:
            console_msg = f"{base_data['report_id']} - {base_data['stage']}"
        elif console_logger.verbosity == 2:
            console_msg = f"{base_data['report_id']} - {base_data['stage']} - {base_data['event']}"
        else:  # verbosity == 3
            console_msg = json.dumps(base_data, indent=2)
        
        console_logger.info(console_msg)

# %% ../nbs/08_mappr_lisette.ipynb 91
class AnalysisResult(AttrDict):
    "Extracts key results from State with framework metadata"
    def __init__(self, state: State, framework_info: FrameworkInfo):
        self.theme_covered = state.theme_covered
        self.coverage_reasoning = state.coverage_reasoning
        self.gaps_identified = state.gaps_identified
        self.explored_sections = state.explored_sections
        self.framework_name = framework_info.name
        self.framework_category = framework_info.category
        self.framework_theme_id = framework_info.theme_id

# %% ../nbs/08_mappr_lisette.ipynb 92
class PipelineResults(dict):
    def __init__(self):
        super().__init__()
        self[Stage.STAGE1] = defaultdict(lambda: defaultdict(dict))
        self[Stage.STAGE2] = defaultdict(lambda: defaultdict(dict))
        self[Stage.STAGE3] = defaultdict(lambda: defaultdict(dict))

# %% ../nbs/08_mappr_lisette.ipynb 93
@patch
def __call__(self:PipelineResults, stage=Stage.STAGE1, filter_type="all"):
    themes = []
    for frameworks in self[stage].values():
        for categories in frameworks.values():
            for theme in categories.values():
                if filter_type == "all" or \
                   (filter_type == "covered" and theme.theme_covered) or \
                   (filter_type == "uncovered" and not theme.theme_covered):
                    themes.append(theme)
    return themes

# %% ../nbs/08_mappr_lisette.ipynb 94
class PipelineOrchestrator:
    "Orchestrator for the IOM evaluation report mapping pipeline"
    def __init__(self, 
                 report_id:str, # Report identifier
                 headings:dict, # Report headings
                 get_content_fn:Callable, # Function to get the content of a section
                 eval_data:EvalData, # Evaluation data
                 verbosity:int=2, # Verbosity level
                 ):
        store_attr()
        setup_trace_logging(report_id, verbosity)
        self.results = PipelineResults()
        self.sections_lookup = {key: path for key, path in flatten_sections(headings)}

# %% ../nbs/08_mappr_lisette.ipynb 95
@patch
async def analyze_with_context(
    self: PipelineOrchestrator,
    theme: str,
    trace_ctx: TraceContext,
    semaphore: Semaphore,
    prior_coverage_context: str = ""
) -> AnalysisResult:
    "Analyze theme with orchestrator context and return wrapped result"
    
    # Create log function with trace context
    log_fn = lambda event, **kw: log_analysis_event(event, trace_ctx, **kw)
    
    # Call analyze_theme
    state = await analyze_theme(
        theme=theme,
        sections_lookup=self.sections_lookup,
        hdgs=self.headings,
        semaphore=semaphore,
        log_fn=log_fn,
        prior_coverage_context=prior_coverage_context,
        max_iterations=cfg.max_iter
    )
    
    # Wrap with framework metadata
    return AnalysisResult(state, trace_ctx.framework)

# %% ../nbs/08_mappr_lisette.ipynb 96
@patch
async def run_stage1(self:PipelineOrchestrator, semaphore):
    "Run stage 1 of the pipeline"
    tasks = []
    
    collections = [
        (self.eval_data.srf_enablers, FrameworkCat.ENABLERS, format_enabler_theme),
        (self.eval_data.srf_crosscutting_priorities, FrameworkCat.CROSSCUT, format_crosscutting_theme)
    ]

    for items, framework_cat, format_fn in collections:
        for item in items:
            trace_ctx = TraceContext(self.report_id, Stage.STAGE1, FrameworkInfo(Framework.SRF, framework_cat, item.id))
            theme = format_fn(item)
            tasks.append(self.analyze_with_context(theme, trace_ctx, semaphore))

    results = await gather(*tasks)
    for result in results: 
        self.results[Stage.STAGE1][result.framework_name][result.framework_category][result.framework_theme_id] = result

# %% ../nbs/08_mappr_lisette.ipynb 100
def get_stage1_covered_context(results: PipelineResults, eval_data: EvalData) -> str:
    "Get and format covered themes in Stage 1."
    covered_themes = results(Stage.STAGE1, filter_type="covered")
    if not covered_themes: return ""
    
    context_parts = []
    for theme in covered_themes:
        if theme.framework_category == str(FrameworkCat.ENABLERS):
            theme_data = next(t for t in eval_data.srf_enablers if t.id == theme.framework_theme_id)
        elif theme.framework_category == str(FrameworkCat.CROSSCUT):
            theme_data = next(t for t in eval_data.srf_crosscutting_priorities if t.id == theme.framework_theme_id)
        
        context_parts.append(f"- **{theme.framework_category} {theme_data.id}**: {theme_data.title}")
    
    return f"### Report Preliminary Context\nThis evaluation report covers the following Strategic Results Framework themes:\n" + "\n".join(context_parts)


# %% ../nbs/08_mappr_lisette.ipynb 103
@patch
async def run_stage2(self:PipelineOrchestrator, semaphore):
    "Run stage 2 of the pipeline - GCM objectives analysis"
    stage1_context = get_stage1_covered_context(self.results, self.eval_data)
    tasks = []
    
    for gcm_obj in self.eval_data.gcm_objectives_small:
        trace_ctx = TraceContext(self.report_id, Stage.STAGE2, FrameworkInfo(Framework.GCM, FrameworkCat.OBJS, gcm_obj["id"]))
        theme = format_gcm_theme(gcm_obj)
        tasks.append(self.analyze_with_context(theme, trace_ctx, semaphore, stage1_context))

    results = await gather(*tasks)
    for result in results: 
        self.results[Stage.STAGE2][result.framework_name][result.framework_category][result.framework_theme_id] = result

# %% ../nbs/08_mappr_lisette.ipynb 107
def get_filtered_srf_output_ids(
    results: PipelineResults, # PipelineResults
    eval_data: EvalData # EvalData
    ) -> list: # list of SRF output IDs
    "Get filtered SRF output IDs based on covered GCM themes."
    covered_gcm = results(Stage.STAGE2, filter_type="covered")
    srf_output_ids = set()
    
    for gcm_theme in covered_gcm:
        gcm_id = gcm_theme.framework_theme_id
        if gcm_id in eval_data.gcm_srf_lut:
            srf_output_ids.update(eval_data.gcm_srf_lut[gcm_id])
    
    return list(srf_output_ids)

# %% ../nbs/08_mappr_lisette.ipynb 110
def get_combined_context(
    results: PipelineResults, # PipelineResults
    eval_data: EvalData, # EvalData
    ) -> str: # combined context
    "Get combined context from previous stages (1 and 2)."
    stage1_context = get_stage1_covered_context(results, eval_data)
    covered_gcm = results(Stage.STAGE2, filter_type="covered")
    
    if not covered_gcm: return stage1_context
    
    gcm_context = "\n".join([f"- **GCM {theme.framework_theme_id}**: {eval_data.gcm_objectives_small[int(theme.framework_theme_id)-1]['title']}" 
                            for theme in covered_gcm])
    
    return f"{stage1_context}\n\n### Covered GCM Objectives\n{gcm_context}"


# %% ../nbs/08_mappr_lisette.ipynb 113
@patch
async def run_stage3(self:PipelineOrchestrator, semaphore):
    "Run stage 3 of the pipeline - Targeted SRF outputs analysis"
    combined_context = get_combined_context(self.results, self.eval_data)
    filtered_output_ids = get_filtered_srf_output_ids(self.results, self.eval_data)
    tasks = []
    
    for output_id in filtered_output_ids:
        output_context = find_srf_output_by_id(self.eval_data, output_id)
        if output_context:
            trace_ctx = TraceContext(self.report_id, Stage.STAGE3, FrameworkInfo(Framework.SRF, FrameworkCat.OUTPUTS, output_id))
            theme = format_srf_output(output_context)
            tasks.append(self.analyze_with_context(theme, trace_ctx, semaphore, combined_context))

    results = await gather(*tasks)
    for result in results: 
        self.results[Stage.STAGE3][result.framework_name][result.framework_category][result.framework_theme_id] = result
