from __future__ import annotations

from pathlib import Path
from typing import Any, Dict

import yaml


def read_system_config(path: str | Path) -> Dict[str, Any]:
    path = Path(path)
    with path.open('r', encoding='utf-8') as f:
        cfg = yaml.safe_load(f) or {}
    return cfg


def flatten_for_bc(system_cfg: Dict[str, Any]) -> Dict[str, Any]:
    paths = system_cfg.get('paths', {})
    files = system_cfg.get('files', {})
    cal = system_cfg.get('calendar', {})
    exp = system_cfg.get('experiment', {})
    stack = system_cfg.get('stacking', {})
    cfg: Dict[str, Any] = {}
    cfg.update(exp)
    cfg.update(stack)
    cfg['data_dir'] = paths.get('hindcast_data_dir')
    cfg['mask_file'] = paths.get('mask_file')
    cfg['output_dir'] = paths.get('output_root', 'outputs')
    cfg['file_pattern'] = files.get('hindcast_pattern', 'lead{lead}_{source_model}_01_BL_Test_VN.nc')
    cfg['year_start'] = int(cal.get('hindcast_start_year', 1982))
    cfg['missing_value'] = float(cal.get('missing_value_hindcast', -99.0))
    cfg['train_years'] = list(cal.get('evaluation_train_years', [1982, 2013]))
    cfg['test_years'] = list(cal.get('evaluation_test_years', [2014, 2021]))
    cfg['fit_years'] = list(cal.get('final_fit_years', [1982, 2021]))
    cfg['climatology_years'] = list(cal.get('climatology_years', [1983, 2020]))
    cfg['random_seed'] = int(exp.get('random_seed', 42))
    cfg['clip_precipitation_to_zero'] = bool(exp.get('clip_precipitation_to_zero', True))
    cfg['include_lat_lon'] = bool(exp.get('include_lat_lon', True))
    cfg['include_region_feature'] = bool(exp.get('include_region_feature', True))
    cfg['valid_regions'] = list(exp.get('valid_regions', [1,2,3,4,5,6,7]))
    # Pass model hyperparameters through to src.models.make_model. Values may be
    # supplied either at top level or under model_params.
    model_params = dict(system_cfg.get('model_params', {}))
    for key in [
        'ridge_alpha', 'elasticnet_alpha', 'elasticnet_l1_ratio', 'lasso_alpha',
        'random_forest_params', 'extratrees_params', 'gbm_params',
        'xgboost_base_params', 'xgboost_variants', 'xgboost_params',
        'mlp_base_params', 'mlp_variants', 'mlp_params', 'cnn1d_small_params',
    ]:
        if key in system_cfg:
            cfg[key] = system_cfg[key]
        elif key in model_params:
            cfg[key] = model_params[key]
    cfg.setdefault('ridge_alpha', 1.0)
    return cfg


def prepare_drought_dirs(output_root: str | Path) -> Dict[str, Path]:
    out = Path(output_root)
    dirs = {
        'output': out,
        'tasks': out / 'tasks',
        'logs': out / 'logs',
        'metrics': out / 'metrics',
        'corrected': out / 'corrected',
        'corrected_full': out / 'corrected_full',
        'corrected_train': out / 'corrected_train',
        'spei': out / 'spei',
        'operational': out / 'operational',
        'figures': out / 'figures',
        'maps': out / 'maps',
        'report': out / 'report',
        'done': out / 'done',
    }
    for p in dirs.values():
        p.mkdir(parents=True, exist_ok=True)
    return dirs