o rei7@sHdZddlZddlZd ddZddZd ddZd d d Zd d d ZdS)a Shared configuration module for model datasets. This module provides a unified way to access model dataset configurations used across multiple scripts (run_filter.py, run_kmeans.py, etc.). The configuration can be read from Snakefile.original or use a fallback hardcoded mapping. Nc Cs|dur+tjtjt}tj|dtj|dg}|D] }tj|r*|}nq|rtj|rzt|d }|}Wdn1sHwYi}d}t ||}|r| } d} d} d} | | d} | t |kr|| }|d vr| dks|| dd kr| sd } |} n>|| krd} d} n5| s|d kr| d7} n*|d kr| d8} | dkr|| | d}t||d|vrtd||dWSn | d7} | t |ksqWnty}z td|WYd}~nd}~wwtdtS)a7 Get VARIANT_DATASET_CONFIGS from Snakefile or return fallback config. Parameters: ----------- snakefile_path : str, optional Path to Snakefile.original. If None, tries to find it automatically. Returns: -------- dict VARIANT_DATASET_CONFIGS dictionary NzSnakefile.originalZ Snakefilerz VARIANT_DATASET_CONFIGS\s*=\s*\{rF)"'\T{}ZVARIANT_DATASET_CONFIGSz$Loaded VARIANT_DATASET_CONFIGS from z/Warning: Could not parse Snakefile for config: z0Using fallback hardcoded VARIANT_DATASET_CONFIGS)ospathdirnameabspath__file__joinexistsopenreadresearchstartendlenexecprint Exceptionget_fallback_config)snakefile_path script_dirZpotential_pathsr fcontent namespacepatternmatchZ start_posZ brace_countZ in_stringZ string_charicharZ config_strer%/oak/stanford/groups/akundaje/airanman/projects/lab/rare-disease-manuscript/curation/broad/varbook-container/snakemake/model_datasets_config.pyget_model_datasets_configsn            r'cCsdddgddgdgiS)z Get fallback hardcoded VARIANT_DATASET_CONFIGS. Returns: -------- dict Fallback VARIANT_DATASET_CONFIGS z4Broad neurodevelopmental and neuromuscular disordersz Fetal BrainzKUN_FB*z KUN_HDMA*)namemodelsZmodel_supersetr%r%r%r%r&r_s rcCsTt|}||vr gS||}t|trd|vr|dSt|tr"|Std|d)a Get the list of model dataset configs for a variant dataset. Handles both list and dict formats in VARIANT_DATASET_CONFIGS. Parameters: ----------- variant_dataset : str Variant dataset name snakefile_path : str, optional Path to Snakefile.original Returns: -------- list List of model dataset config dicts model_datasetsz+Invalid config format for variant_dataset 'z2': expected list or dict with 'model_datasets' key)r' isinstancedictlist ValueError)variant_datasetrZconfigsconfigr%r%r&get_model_datasets_listss  r1cCsht||}|D]}t|tr!|d|kr!|dg}|r!|Sqtd|d|ddd|D)a Get model patterns for a model_dataset name. Parameters: ----------- variant_dataset : str Variant dataset name model_dataset_name : str Model dataset name (e.g., "Fetal Brain") snakefile_path : str, optional Path to Snakefile.original Returns: -------- list of str Model patterns (e.g., ["KUN_FB*"]) r(r)z1Could not find model patterns for model_dataset 'z' in variant_dataset 'z'. Available model_datasets: cSs g|] }t|tr|dqSr(r+r,get.0cr%r%r& s z3get_model_patterns_from_dataset..)r1r+r,r4r.)r/model_dataset_namerr*r0r)r%r%r&get_model_patterns_from_datasets   r:cCst||}dd|DS)a7 List all available model datasets for a variant dataset. Parameters: ----------- variant_dataset : str Variant dataset name snakefile_path : str, optional Path to Snakefile.original Returns: -------- list of str List of model dataset names cSs(g|]}t|trd|vr|dqSr2r3r5r%r%r&r8s(z'list_model_datasets..)r1)r/rr*r%r%r&list_model_datasetss r;)N)__doc__r rr'rr1r:r;r%r%r%r&s  O  %#