import copy
from typing import Any

import networkx as nx


def merge_states(state_set: list) -> dict:
    """
    Merge states
    :param state_set: list of states
    :return:
    """

    if len(state_set) == 0:
        raise ValueError("No states to merge")

    if len(state_set) == 1:
        return state_set[0]

    merged_state = state_set[0]
    for state in state_set[1:]:
        for key, value in state.items():
            if key in merged_state:
                if isinstance(merged_state[key], nx.Graph):
                    merged_state[key] = merge_graph(merged_state[key], value)
                elif isinstance(merged_state[key], dict):
                    for sub_key, sub_value in value.items():
                        merge_by_key(merged_state[key], sub_key, sub_value)
                elif isinstance(merged_state[key], list):
                    merged_state[key] = merged_state[key] + [x for x in value if x not in merged_state[key]]
                elif merged_state[key] == value:
                    pass
                else:
                    merged_state[key] = merged_state[key] + value
            else:
                merged_state[key] = value

    return merged_state


def merge_by_key(base_dict: dict, key: str | int, value: Any):
    """
    Merge a value into a dictionary by key. If the key is not present, it is added, otherwise the value is updated by addition
    :param base_dict: base dictionary to update
    :param key: key to update
    :param value: value to update with
    :return: nothing, base_dict is updated in place
    """
    if key not in base_dict:
        base_dict[key] = value
    else:
        if isinstance(base_dict[key], dict):
            for sub_key, sub_value in value.items():
                merge_by_key(base_dict[key], sub_key, sub_value)
        elif isinstance(base_dict[key], list):
            if isinstance(value, list):
                base_dict[key] = base_dict[key] + [x for x in value if x not in base_dict[key]]
            else:
                if value not in base_dict[key]:
                    base_dict[key].append(value)
        else:
            if key in base_dict:
                try:
                    base_dict[key] = base_dict[key] + value
                except:
                    pass
            else:
                base_dict[key] = value


def merge_graph(g: nx.Graph, h: nx.Graph) -> nx.Graph:
    """
    Merge two graphs, combines the weights of the edges based on the keys and tries to add their values
    :param g: first graph
    :param h: second graph
    :return: merged graph
    """

    if g is None:
        return h

    if h is None:
        return g

    out = copy.deepcopy(g)

    # iterate over nodes in h
    for node in h.nodes():
        if node not in out.nodes:
            out.add_node(node)

    # iterate over edges in h
    for src, dest, data in h.edges(data=True):
        if out.has_edge(src, dest):
            for key, value in data.items():
                merge_by_key(out[src][dest], key, value)
        else:
            out.add_edge(src, dest, **data)

    return out
