"""
Utility functions common to testing the isl mapper functions.
"""

from pathlib import Path
import islpy as isl

from ruamel.yaml import YAML


def to_isl_maps(obj: str | list | dict) -> dict:
    """
    Given an object, attempt to reduce all strings in tree with isl.Map

    Parameters
    ----------
    obj:
        A DAG which can be explored and contains isl.Map strings within it.

    Returns
    -------
    `obj` but all strings are converted to isl.Map.
    """

    def _to_isl_maps(obj: str | dict | list) -> isl.Map | dict | list:
        """Recursively convert string ISL maps to isl.Map; leave others alone."""
        if isinstance(obj, str):
            return isl.Map.read_from_str(isl.DEFAULT_CONTEXT, obj)
        if isinstance(obj, dict):
            return {k: (_to_isl_maps(v) if k != "type" else v) for k, v in obj.items()}
        if isinstance(obj, list):
            return [_to_isl_maps(v) for v in obj]
        return obj

    return _to_isl_maps(obj)  # type: ignore


def load_solutions(path: Path) -> dict:
    """
    Loads in a dictionary with the isl solutions to a workload problem.

    Parameters
    ----------
    path:
        The path to the solutions.

    Returns
    -------
    A dictionary relating Python-based keys generated by the mapper (e.g.,
    `BufferTensorEinsum` to their corresponding isl.Map.)
    """
    # Load expected solutions (YAML file with string ISL maps)
    yaml: YAML = YAML(typ="safe")

    with open(path, "r", encoding="utf-8") as f:
        return to_isl_maps(yaml.load(f))
