from typing import Dict
import logging


def _get_y_hat_from_results(results: dict, compare_mode: str) -> dict:
    """
    Iterate through results (indexed by results dirs, and then models),
    and save y_hat indexed in order determined by compare_mode:
    within_dirs: [results_dir][model_name]
    across_dirs: [model_name][results_dir]

    If y_hat is not available, return an empty dictionary.
    """
    y_hat: Dict[str, dict] = {}

    for results_dir, result in results.items():
        if "y_hat_test_outer_folds" in result:
            for model_name in result["y_hat_test_outer_folds"]:
                if compare_mode == "within_dirs":
                    if results_dir not in y_hat:
                        y_hat[results_dir] = {}
                    y_hat[results_dir][model_name] = result["y_hat_test_outer_folds"][
                        model_name
                    ]

                elif compare_mode == "across_dirs":
                    if model_name not in y_hat:
                        y_hat[model_name] = {}
                    y_hat[model_name][results_dir] = result["y_hat_test_outer_folds"][
                        model_name
                    ]
        else:
            logging.warning(
                "y_hat_test_outer_folds not found in"
                f" results[{results_dir}]. Returning y_hat as empty dict.",
            )
    return y_hat