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

"""A module containing merge_graphs, merge_nodes, merge_edges, merge_attributes, apply_merge_operation and _get_detailed_attribute_merge_operation methods definitions."""

from typing import Any

import networkx as nx
from datashaper import VerbCallbacks, progress_iterable

from graphrag.index.operations.merge_graphs.typing import (
    BasicMergeOperation,
    DetailedAttributeMergeOperation,
    NumericOperation,
    StringOperation,
)

DEFAULT_NODE_OPERATIONS = {
    "*": {
        "operation": BasicMergeOperation.Replace,
    }
}

DEFAULT_EDGE_OPERATIONS = {
    "*": {
        "operation": BasicMergeOperation.Replace,
    },
    "weight": "sum",
}

DEFAULT_CONCAT_SEPARATOR = ","


def merge_graphs(
    graphs: list[nx.Graph],
    callbacks: VerbCallbacks,
    node_operations: dict[str, Any] | None,
    edge_operations: dict[str, Any] | None,
) -> nx.Graph:
    """
    Merge multiple graphs together. The graphs are expected to be in nx.Graph format. The verb outputs a new column containing the merged graph.

    > Note: This will merge all rows into a single graph.

    ## Usage
    ```yaml
    verb: merge_graph
    args:
        nodes: <node operations> # See node operations section below
        edges: <edge operations> # See edge operations section below
    ```

    ## Node Operations
    The merge graph verb can perform operations on the nodes of the graph.

    ### Usage
    ```yaml
    nodes:
        <attribute name>: <operation>
        ... for each attribute or use the special value "*" for all attributes
    ```

    ## Edge Operations
    The merge graph verb can perform operations on the nodes of the graph.

    ### Usage
    ```yaml
    edges:
        <attribute name>: <operation>
        ... for each attribute or use the special value "*" for all attributes
    ```

    ## Operations
    The merge graph verb can perform operations on the nodes and edges of the graph. The following operations are available:

    - __replace__: This operation replaces the attribute with the last value seen.
    - __skip__: This operation skips the attribute, and just uses the first value seen.
    - __concat__: This operation concatenates the attribute with the last value seen.
    - __sum__: This operation sums the attribute with the last value seen.
    - __max__: This operation takes the max of the attribute with the last value seen.
    max
    - __min__: This operation takes the min of the attribute with the last value seen.
    - __average__: This operation takes the mean of the attribute with the last value seen.
    - __multiply__: This operation multiplies the attribute with the last value seen.
    """
    nodes = node_operations or DEFAULT_NODE_OPERATIONS
    edges = edge_operations or DEFAULT_EDGE_OPERATIONS
    node_ops = {
        attrib: _get_detailed_attribute_merge_operation(value)
        for attrib, value in nodes.items()
    }
    edge_ops = {
        attrib: _get_detailed_attribute_merge_operation(value)
        for attrib, value in edges.items()
    }

    mega_graph = nx.Graph()
    num_total = len(graphs)
    for graph in progress_iterable(graphs, callbacks.progress, num_total):
        merge_nodes(mega_graph, graph, node_ops)
        merge_edges(mega_graph, graph, edge_ops)

    return mega_graph


def merge_nodes(
    target: nx.Graph,
    subgraph: nx.Graph,
    node_ops: dict[str, DetailedAttributeMergeOperation],
):
    """Merge nodes from subgraph into target using the operations defined in node_ops."""
    for node in subgraph.nodes:
        if node not in target.nodes:
            target.add_node(node, **(subgraph.nodes[node] or {}))
        else:
            merge_attributes(target.nodes[node], subgraph.nodes[node], node_ops)


def merge_edges(
    target_graph: nx.Graph,
    subgraph: nx.Graph,
    edge_ops: dict[str, DetailedAttributeMergeOperation],
):
    """Merge edges from subgraph into target using the operations defined in edge_ops."""
    for source, target, edge_data in subgraph.edges(data=True):  # type: ignore
        if not target_graph.has_edge(source, target):
            target_graph.add_edge(source, target, **(edge_data or {}))
        else:
            merge_attributes(
                target_graph.edges[(source, target)],  # noqa
                edge_data,
                edge_ops,
            )


def merge_attributes(
    target_item: dict[str, Any] | None,
    source_item: dict[str, Any] | None,
    ops: dict[str, DetailedAttributeMergeOperation],
):
    """Merge attributes from source_item into target_item using the operations defined in ops."""
    source_item = source_item or {}
    target_item = target_item or {}
    for op_attrib, op in ops.items():
        if op_attrib == "*":
            for attrib in source_item:
                # If there is a specific handler for this attribute, use it
                # i.e. * provides a default, but you can override it
                if attrib not in ops:
                    apply_merge_operation(target_item, source_item, attrib, op)
        else:
            if op_attrib in source_item or op_attrib in target_item:
                apply_merge_operation(target_item, source_item, op_attrib, op)


def apply_merge_operation(
    target_item: dict[str, Any] | None,
    source_item: dict[str, Any] | None,
    attrib: str,
    op: DetailedAttributeMergeOperation,
):
    """Apply the merge operation to the attribute."""
    source_item = source_item or {}
    target_item = target_item or {}

    if (
        op.operation == BasicMergeOperation.Replace
        or op.operation == StringOperation.Replace
    ):
        target_item[attrib] = source_item.get(attrib, None) or ""
    elif (
        op.operation == BasicMergeOperation.Skip or op.operation == StringOperation.Skip
    ):
        target_item[attrib] = target_item.get(attrib, None) or ""
    elif op.operation == StringOperation.Concat:
        separator = op.separator or DEFAULT_CONCAT_SEPARATOR
        target_attrib = target_item.get(attrib, "") or ""
        source_attrib = source_item.get(attrib, "") or ""
        target_item[attrib] = f"{target_attrib}{separator}{source_attrib}"
        if op.distinct:
            # TODO: Slow
            target_item[attrib] = separator.join(
                sorted(set(target_item[attrib].split(separator)))
            )

    # We're assuming that the attribute is numeric
    elif op.operation == NumericOperation.Sum:
        target_item[attrib] = (target_item.get(attrib, 0) or 0) + (
            source_item.get(attrib, 0) or 0
        )
    elif op.operation == NumericOperation.Average:
        target_item[attrib] = (
            (target_item.get(attrib, 0) or 0) + (source_item.get(attrib, 0) or 0)
        ) / 2
    elif op.operation == NumericOperation.Max:
        target_item[attrib] = max(
            (target_item.get(attrib, 0) or 0), (source_item.get(attrib, 0) or 0)
        )
    elif op.operation == NumericOperation.Min:
        target_item[attrib] = min(
            (target_item.get(attrib, 0) or 0), (source_item.get(attrib, 0) or 0)
        )
    elif op.operation == NumericOperation.Multiply:
        target_item[attrib] = (target_item.get(attrib, 1) or 1) * (
            source_item.get(attrib, 1) or 1
        )
    else:
        msg = f"Invalid operation {op.operation}"
        raise ValueError(msg)


def _get_detailed_attribute_merge_operation(
    value: str | dict[str, Any],
) -> DetailedAttributeMergeOperation:
    """Normalize the AttributeMergeOperation into a DetailedAttributeMergeOperation."""
    if isinstance(value, str):
        return DetailedAttributeMergeOperation(operation=value)
    return DetailedAttributeMergeOperation(**value)
