#!/usr/bin/env python3

'''A utility module for knowledge graph operations'''

from typing import Tuple

import networkx as nx

import pandas as pd

def kg_to_df_pandas(kg: nx.DiGraph) -> Tuple[pd.DataFrame, pd.DataFrame]:

    """

    Convert a directed knowledge graph to a pandas DataFrame.

    Args:

        kg: The directed knowledge graph in networkX format.

    Returns:

        df_nodes: A pandas DataFrame of the nodes in the knowledge graph.

        df_edges: A pandas DataFrame of the edges in the knowledge graph.

    """

    # Create a pandas DataFrame of the nodes

    df_nodes = pd.DataFrame.from_dict(kg.nodes, orient='index')

    # Create a pandas DataFrame of the edges

    df_edges = nx.to_pandas_edgelist(kg,

                                    source='node_source',

                                    target='node_target')

    return df_nodes, df_edges

def df_pandas_to_kg(df: pd.DataFrame,

                    df_nodes_attrs: pd.DataFrame,

                    node_source: str,

                    node_target: str

                    ) -> nx.DiGraph:

    """

    Convert a pandas DataFrame to a directed knowledge graph.

    Args:

        df: A pandas DataFrame of the edges in the knowledge graph.

        df_nodes_attrs: A pandas DataFrame of the nodes in the knowledge graph.

        node_source: The column name of the source node in the df.

        node_target: The column name of the target node in the df.

    Returns:

        kg: The directed knowledge graph in networkX format.

    """

    # Assert if the columns node_source and node_target are in the df

    assert node_source in df.columns, f'{node_source} not in df'

    assert node_target in df.columns, f'{node_target} not in df'

    # Assert that the nodes in the index of the df_nodes_attrs

    # are present in the source and target columns of the df

    assert set(df_nodes_attrs.index).issubset(set(df[node_source]).\

                                        union(set(df[node_target]))), \

                                        'Nodes in index of df_nodes not found in df_edges'

    # Create a knowledge graph from the dataframes

    # Add edges and nodes to the knowledge graph

    kg = nx.from_pandas_edgelist(df,

                                source=node_source,

                                target=node_target,

                                create_using=nx.DiGraph,

                                edge_attr=True)

    kg.add_nodes_from(df_nodes_attrs.to_dict('index').items())

    return kg