from __future__ import annotations
from pathlib import Path
from typing import Any, Dict
import yaml


def project_root() -> Path:
    return Path(__file__).resolve().parents[1]


def load_config(path: str | Path | None = None) -> Dict[str, Any]:
    p = Path(path) if path else project_root() / 'config' / 'system.yaml'
    with open(p, 'r', encoding='utf-8') as f:
        cfg = yaml.safe_load(f)
    cfg['_config_file'] = str(p)
    return cfg


def ensure_dirs(cfg: Dict[str, Any]) -> Dict[str, Path]:
    root = project_root()
    out = root / cfg['paths'].get('output_root', 'outputs')
    art = root / cfg['paths'].get('artifact_root', 'artifacts')
    dirs = {
        'root': root,
        'outputs': out,
        'artifacts': art,
        'tasks': out / 'tasks',
        'logs': out / 'logs',
        'evaluation_members': out / 'evaluation_output' / 'members',
        'evaluation_ensemble': out / 'evaluation_output' / 'ensemble',
        'production_members': out / 'production_output' / 'members',
        'production_ensemble': out / 'production_output' / 'ensemble',
        'clim': out / 'Clim',
        'spei_members': out / 'production_output' / 'spei_members',
        'spei_ai': out / 'evaluation_output' / 'spei_ai',
        'spei_params': art / 'spei_params',
        'models': art / 'models',
        'final_models': art / 'final_models_for_GD2',
        'report': out / 'report',
        'figures': out / 'report' / 'figures',
        'tables': out / 'report' / 'tables',
        'done': out / 'done',
    }
    for d in dirs.values():
        d.mkdir(parents=True, exist_ok=True)
    return dirs