"""Load and format SRF, GCM, and cross-cutting priority data for LLM consumption"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/04_themes.ipynb.

# %% auto 0
__all__ = ['load_enbs', 'load_ccps', 'load_srf_outs', 'load_gcms', 'load_gcms_lut', 'load_thm', 'fmt_enb_ccp', 'get_srf_out',
           'fmt_srf_out', 'fmt_srf_outs', 'get_srf_outs', 'load_all_thms']

# %% ../nbs/04_themes.ipynb 2
from fastcore.all import *
from pathlib import Path
import json

# %% ../nbs/04_themes.ipynb 3
def load_thm(fname:str,  # Filename to load
                path:str=None,  # Directory containing theme files
                md:bool=False  # Return raw text instead of JSON?
               ) -> dict|str:  # Theme data as dict or raw text
    "Load theme file from `path` (JSON by default, or raw text if `md=True`)"
    if path is None:
        try: path = Path(__file__).parent / 'files' / 'themes'
        except NameError: path = Path('files/themes')
    p = Path(path)/fname
    return p.read_text() if md else json.loads(p.read_text())

# %% ../nbs/04_themes.ipynb 4
load_enbs = partial(load_thm, 'srf_enablers.json')
load_ccps = partial(load_thm, 'crosscutting_priorities.json')
load_srf_outs = partial(load_thm, 'srf_objectives.json')
load_gcms = partial(load_thm, 'gcms_long.md', md=True)
load_gcms_lut = partial(load_thm, 'gcm_to_srf_outputs.json')

# %% ../nbs/04_themes.ipynb 4
load_enbs = partial(load_thm, 'srf_enablers.json')
load_ccps = partial(load_thm, 'crosscutting_priorities.json')
load_srf_outs = partial(load_thm, 'srf_objectives.json')
load_gcms = partial(load_thm, 'gcms_long.md', md=True)
load_gcms_lut = partial(load_thm, 'gcm_to_srf_outputs.json')

# %% ../nbs/04_themes.ipynb 16
def fmt_enb_ccp(
    items:list # List of enabler/CCP dicts with id, title, description
    ) -> str:
    "Format enablers or cross-cutting priorities for LLM"
    return '\n\n'.join([f'## Cross-cutting {o["id"]}: {o["title"]}\n### Description\n{o["description"]}' for o in items])

# %% ../nbs/04_themes.ipynb 18
def get_srf_out(objectives:list, # SRF objectives structure
                output_id:str # Output ID to find
               ) -> dict|None:
    "Retrieve single SRF output with hierarchy"
    for obj in objectives:
        for lt_out in obj.get('long_term_outcomes', []):
            for st_out in lt_out.get('short_term_outcomes', []):
                if o := next((o for o in st_out.get('outputs', []) if o['id'] == output_id), None):
                    return dict(obj_id=obj['id'], obj=obj['title'], lt_out_id=lt_out['id'], lt_out=lt_out['title'],
                                st_out_id=st_out['id'], st_out=st_out['title'], output=o)

# %% ../nbs/04_themes.ipynb 20
def fmt_srf_out(output_ctx:dict # Dict with obj, lt_out, st_out, and output fields
               ) -> str:
    "Format SRF output with hierarchical context for LLM"
    o = output_ctx
    return '\n'.join([
        f'## SRF Output {o["output"]["id"]}: {o["output"]["title"]}',
        '### Hierarchical Context',
        f'**Objective {o["obj_id"]}**: {o["obj"]}',
        f'**Long-term Outcome {o["lt_out_id"]}:** {o["lt_out"]}',
        f'**Short-term Outcome {o["st_out_id"]}:** {o["st_out"]}'
    ])

# %% ../nbs/04_themes.ipynb 22
def fmt_srf_outs(objectives:list, # SRF objectives structure
                 output_ids:list # List of output IDs to format
                ) -> str:
    "Format multiple SRF outputs"
    return '\n\n'.join(fmt_srf_out(get_srf_out(objectives, i)) for i in output_ids)


# %% ../nbs/04_themes.ipynb 25
def get_srf_outs(lut:dict, # GCM to SRF lookup dict
                 gcm_ids:list # List of GCM IDs to filter by
                ) -> L:
    "Get SRF output IDs filtered by GCM IDs"
    return L(v for k,v in lut.items() if k in gcm_ids).concat()

# %% ../nbs/04_themes.ipynb 28
def load_all_thms(
    path:str=None  # Directory containing theme JSON files
    ) -> AttrDict: # Dict with enablers, ccp, gcms, srf_outs, gcm_lut
    "Load all theme data from path"
    return AttrDict(enablers=load_enbs(path), ccp=load_ccps(path), gcms=load_gcms(path), 
                    srf_outs=load_srf_outs(path), gcm_lut=load_gcms_lut(path))
