"""Maps IOM evaluation report against evaluation frameworks"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/05_mapper.ipynb.

# %% auto #0
__all__ = ['logger', 'ThemeScore', 'ThemeScores', 'parse_res', 'sort_by_relevance', 'get_top_ids', 'mk_system_blocks',
           'map_themes', 'load_prompts', 'map_all']

# %% ../nbs/05_mapper.ipynb #2054091e
from fastcore.all import *
from pydantic import BaseModel
from lisette.core import completion, mk_msg
from .core import load_prompt, n_tokens
from .themes import load_enbs, load_ccps, load_gcms, load_srf_outs, load_gcms_lut, fmt_enbs, fmt_ccps, fmt_srf_outs, get_srf_outs, load_all_thms
import json
import logging
from litellm.types.utils import ModelResponse

# %% ../nbs/05_mapper.ipynb #e81afd30
logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.WARNING, format='%(name)s - %(levelname)s - %(message)s')
logger.setLevel(logging.DEBUG)

# %% ../nbs/05_mapper.ipynb #fcef7312
class ThemeScore(BaseModel):
    "Single theme's relevance assessment with score, reasoning, and confidence"
    theme_id: str       # Unique identifier for the theme
    theme_title: str    # Human-readable theme name
    relevance_score: float  # 0-1 score indicating how relevant theme is to report
    reasoning: str      # LLM's explanation for the score

# %% ../nbs/05_mapper.ipynb #80bd5b66
class ThemeScores(BaseModel):
    "Collection of theme relevance scores from a single mapping operation"
    scores: list[ThemeScore]  # All theme scores for a single mapping call

# %% ../nbs/05_mapper.ipynb #2f4c1d06
def parse_res(res              # LLM completion response object
             ) -> list[dict]:  # list of score dicts
    "Extract scores list from LLM completion response"
    content = res.choices[0].message.content
    if content.startswith('```'): content = content.split('\n', 1)[1].rsplit('```', 1)[0]
    data = json.loads(content)
    return data.get('parameter', data)['scores']

# %% ../nbs/05_mapper.ipynb #87fe9898
def sort_by_relevance(scores           # list of score dicts with 'relevance_score' key
                     ) -> list[dict]:  # list sorted by relevance descending
    "Sort themes by relevance score descending"
    return sorted(scores, key=lambda x: x['relevance_score'], reverse=True)

# %% ../nbs/05_mapper.ipynb #2d37167e
def get_top_ids(scores,            # list of score dicts with 'theme_id' and 'relevance_score'
                min_score=0.7      # minimum relevance threshold
               ) -> list[str]:     # list of theme IDs above threshold
    "Get IDs of themes with relevance score >= min_score, sorted by relevance"
    return [o['theme_id'] for o in sort_by_relevance(scores) if o['relevance_score'] >= min_score]

# %% ../nbs/05_mapper.ipynb #c1dd3661
def mk_system_blocks(report:str  # Full report text to analyze
                    ) -> list:   # Anthropic-style content blocks with cache control
    "Create cached system message blocks from report text"
    return [{"type": "text", "text": f"## Report to Analyze\n\n{report}", "cache_control": {"type": "ephemeral"}}]

# %% ../nbs/05_mapper.ipynb #76603e54
def map_themes(system_blocks:list,             # Cached system blocks from mk_system_blocks
               themes:str,                     # Formatted themes text to score against
               prompt:str,                     # Mapping instruction prompt
               model:str='claude-haiku-4-5',   # Model to use for completion
               response_format=ThemeScores,    # Pydantic model for structured output
               max_tokens:int=8192,            # Max tokens for completion
               temperature:float=0,            # Temperature for completion
               reasoning_effort:str=None,      # Reasoning effort for completion (low, medium, high)
               **kwargs                        # Additional kwargs for completion
              ) -> ModelResponse:              # LLM completion response with structured output
    "Map report against themes using cached system blocks"
    return completion(model=model, system=system_blocks, messages=[mk_msg(f"{prompt}\n\n## Themes\n\n{themes}")], 
                      response_format=response_format, max_tokens=max_tokens, temperature=temperature, 
                      reasoning_effort=reasoning_effort, **kwargs)

# %% ../nbs/05_mapper.ipynb #ecd0df02
def load_prompts(
    path:Path|str|None=None # Directory containing prompt files, defaults to 'files/prompts'
    ) -> AttrDict:          # Dict with srf_enablers, srf_ccps, gcms, srf_outputs prompts
    "Load all mapping prompts"
    name_map = {'srf_enablers': 'enbs', 'srf_ccps': 'ccps', 'gcms': 'gcms', 'srf_outputs': 'outs'}
    return AttrDict({short: load_prompt(file, path) for file, short in name_map.items()})

# %% ../nbs/05_mapper.ipynb #33c29cf8
@delegates(map_themes)
def map_all(report:str,                      # Full report text to analyze
            path:str='files/themes',         # Directory containing theme JSON files
            prompt_path:str='files/prompts', # Directory containing prompt files
            verbose:bool=True,               # Log progress messages
            **kwargs                         # Additional args passed to map_themes (e.g. model)
           ) -> AttrDict:                    # Dict with enablers, ccp, gcm, outputs results
    "Map report against all theme classes: enablers → CCP → GCM → outputs"
    themes, prompts = load_all_thms(path), load_prompts(prompt_path)
    system_blocks = mk_system_blocks(report)
    
    if verbose: logger.info("Mapping SRF Enablers...")
    enbs_res = parse_res(map_themes(system_blocks, fmt_enbs(themes.enbs), prompts.enbs, **kwargs))
    if verbose: logger.info("Mapping Cross-cutting Priorities...")
    ccps_res = parse_res(map_themes(system_blocks, fmt_ccps(themes.ccps), prompts.ccps, **kwargs))
    if verbose: logger.info("Mapping GCM Objectives...")
    gcms_res = parse_res(map_themes(system_blocks, themes.gcms, prompts.gcms, **kwargs))
    
    top_gcm_ids = get_top_ids(gcms_res)
    if not top_gcm_ids:
        if verbose: logger.info("No GCM objectives scored ≥0.7, skipping SRF Outputs")
        return AttrDict(enbs=enbs_res, ccps=ccps_res, gcms=gcms_res, outs=None)
    
    if verbose: logger.info(f"Top GCM: {top_gcm_ids[0]} (from {len(top_gcm_ids)} candidates)")
    output_ids = get_srf_outs(themes.gcm_lut, [top_gcm_ids[0]])
    if verbose: logger.info(f"Mapping {len(output_ids)} filtered SRF Outputs...")
    outs_res = parse_res(map_themes(system_blocks, fmt_srf_outs(themes.outs, output_ids), prompts.outs, **kwargs))
    return AttrDict(enbs=enbs_res, ccps=ccps_res, gcms=gcms_res, outs=outs_res)

