# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

from collections import OrderedDict
from typing import Any, Callable, Dict, Union

from archai.discrete_search.search_spaces.config.discrete_choice import DiscreteChoice


def flatten_dict(odict: Dict[str, Any]) -> dict:
    """Flatten a nested dictionary into a single level dictionary.

    Args:
        odict: Nested dictionary.

    Returns:
        Flattened dictionary.

    """

    fdict = dict()

    def _flatten(prefix: str, d: Dict[str, Any]) -> Dict[str, Any]:
        prefix = prefix + "." if prefix else prefix

        if isinstance(d, dict):
            for k, v in d.items():
                flat_v = _flatten(prefix + k, v)

                if flat_v is not None:
                    fdict[prefix + k] = flat_v
        else:
            return d

    _flatten("", odict)

    return fdict


def order_dict_keys(base_dict: OrderedDict, target_dict: Dict[str, Any]) -> OrderedDict:
    """Order the keys of a target dictionary based on a base dictionary.

    Args:
        base_dict (OrderedDict[str, Any]): Dictionary with the desired key order.
        target_dict (Dict[str, Any]): Dictionary to be ordered.

    Returns:
        OrderedDict[str, Any]: Ordered version of `target_dict` dictionary.
    """
    ordered_dict = OrderedDict()

    for k in base_dict:
        if k in target_dict:
            ordered_dict[k] = target_dict[k]

    return ordered_dict


def replace_ptree_choices(
    config_tree: Union[Dict, DiscreteChoice], repl_fn: Callable[[DiscreteChoice], Any]
) -> OrderedDict:
    """Replace all DiscreteChoice nodes in a tree with the output of a function.

    Args:
        config_tree: Tree with DiscreteChoice nodes.
        repl_fn: Function to replace DiscreteChoice nodes.

    Returns:
        Replaced tree.

    """

    def _replace_tree_nodes(node, repl_fn, ref_map):
        if isinstance(node, dict):
            output_tree = OrderedDict()

            for param_name, param in node.items():
                output_tree[param_name] = _replace_tree_nodes(param, repl_fn, ref_map)

        elif isinstance(node, DiscreteChoice):
            if id(node) not in ref_map:
                ref_map[id(node)] = repl_fn(node)
            return ref_map[id(node)]

        else:
            return node

        return output_tree

    return _replace_tree_nodes(config_tree, repl_fn, {})


def replace_ptree_pair_choices(
    query_tree: Union[Dict, DiscreteChoice], aux_tree: Union[Dict, Any], repl_fn: Callable[[DiscreteChoice, Any], Any]
) -> OrderedDict:
    """Replace all DiscreteChoice nodes in a tree with the output of a function and an auxilary tree.

    Args:
        query_tree: Tree with DiscreteChoice nodes.
        aux_tree: Auxiliary tree with DiscreteChoice nodes.
        repl_fn: Function that takes a `query_node` and an `aux_node` and returns a replacement for `query_node`.

    Returns:
        Replaced tree.

    """

    def _replace_tree_nodes(query_node, aux_node, repl_fn, ref_map):
        if isinstance(query_node, dict):
            output_tree = OrderedDict()

            for param_name, param in query_node.items():
                assert param_name in aux_node, "`aux_tree` must be identical to `query_tree` apart from terminal nodes"

                output_tree[param_name] = _replace_tree_nodes(param, aux_node[param_name], repl_fn, ref_map)

        elif isinstance(query_node, DiscreteChoice):
            if id(query_node) not in ref_map:
                ref_map[id(query_node)] = repl_fn(query_node, aux_node)

            return ref_map[id(query_node)]

        else:
            return query_node

        return output_tree

    return _replace_tree_nodes(query_tree, aux_tree, repl_fn, {})
