import networkx as nx

import numpy as np
import pandas as pd


def get_info_for_node(graph: nx.Graph, node: str):
    """
    Get info for a node
    :param graph: graph to analyze
    :param node: node to analyze
    :return: info for the node
    """
    edges = list(graph.edges(node, data=True))
    return edges


def extract_connections(graph: nx.Graph) -> list:
    """
    Extract connections from a graph
    :param graph: graph to analyze
    :return: list of connections
    """
    connections = []
    for src, dest, data in graph.edges(data=True):
        connections.append((src, dest, data))
    return connections


def create_dataset_for_connection(connection_edge_data: dict,
                                  per_second_attributes: list[str] = None,
                                  per_connection_attributes: list[str] = None,
                                  first_second: int = None,
                                  last_second: int = None) -> pd.DataFrame:
    """
    Create a dataset for a connection based on the graph edge data
    :param connection_edge_data: graph edge data for the connection
    :param per_second_attributes: attributes to extract per second
    :param per_connection_attributes: attributes to extract per connection
    :param first_second: first second to include in the dataset
    :param last_second: last second to include in the dataset
    :return: dataset for the connection
    """


def create_dataset_for_node(graph: nx.Graph,
                            node: str,
                            per_second_attributes: list[str] = None,
                            per_connection_attributes: list[str] = None,
                            first_second: int = None,
                            last_second: int = None) -> pd.DataFrame:
    """
    Create a dataset for a connection
    :param graph: graph to analyze
    :param node: name of the node
    :param per_second_attributes: attributes to extract per second
    :param per_connection_attributes: attributes to extract per connection
    :param first_second: first second to include in the dataset
    :param last_second: last second to include in the dataset
    :return: dataset for the connection
    """

    # get incoming and outgoing edges for node
    incoming = list(graph.in_edges(node, data=True))
    outgoing = list(graph.out_edges(node, data=True))

    out_packet_seconds = list()
    for edge in outgoing:
        out_packet_seconds.extend(list(edge[2]["packets_per_second"].keys()))

    out_packet_seconds = sorted(list(set(out_packet_seconds)))

    in_packet_seconds = list()
    for edge in incoming:
        in_packet_seconds.extend(list(edge[2]["packets_per_second"].keys()))

    in_packet_seconds = sorted(list(set(in_packet_seconds)))

    # filter out seconds, if first_second or last_second is given
    if first_second is not None or last_second is not None:
        if first_second is None:
            first_second = min(out_packet_seconds + in_packet_seconds)
        if last_second is None:
            last_second = max(out_packet_seconds + in_packet_seconds)

        in_packet_seconds = [second for second in in_packet_seconds if first_second <= second <= last_second]
        out_packet_seconds = [second for second in out_packet_seconds if first_second <= second <= last_second]

    # extract per connection values
    per_connection_values_in = dict()
    for src, dest, data in incoming:
        for attribute in per_connection_attributes:
            if attribute in data:
                if isinstance(data[attribute], list):
                    # when list, use length
                    per_connection_values_in[f"{attribute}_in"] = [len(data[attribute])] * len(in_packet_seconds)
                else:
                    per_connection_values_in[f"{attribute}_in"] = [data[attribute]] * len(in_packet_seconds)
            else:
                per_connection_values_in[f"{attribute}_in"] = [0] * len(in_packet_seconds)

    if len(incoming) == 0:
        per_connection_values_in = {f"{attribute}_in": [0] * len(in_packet_seconds) for attribute in
                                    per_connection_attributes}

    per_connection_values_out = dict()
    for src, dest, data in outgoing:
        for attribute in per_connection_attributes:
            if attribute in data:
                if isinstance(data[attribute], list):
                    # when list, use length
                    per_connection_values_out[f"{attribute}_out"] = [len(data[attribute])] * len(out_packet_seconds)
                else:
                    per_connection_values_out[f"{attribute}_out"] = [data[attribute]] * len(out_packet_seconds)
            else:
                per_connection_values_out[f"{attribute}_out"] = [0] * len(out_packet_seconds)

    if len(outgoing) == 0:
        per_connection_values_out = {f"{attribute}_out": [0] * len(out_packet_seconds) for attribute in
                                     per_connection_attributes}

    # extract per second value for incoming connections
    per_second_values_in = list()
    for second in in_packet_seconds:
        values_this_second = dict()
        for src, dest, data in incoming:
            for attribute in per_second_attributes:
                # extract attribute values for the second, if they exist, else set to 0
                if attribute in data:
                    attribute_values = data[attribute]
                    if second in attribute_values:
                        if attribute not in values_this_second:
                            values_this_second[f"{attribute}_in"] = attribute_values[second]
                        else:
                            values_this_second[f"{attribute}_in"] += attribute_values[second]
                    else:
                        values_this_second[f"{attribute}_in"] = 0
                else:
                    values_this_second[f"{attribute}_in"] = 0

        # add to per second values
        per_second_values_in.append(values_this_second)

    # still add attributes to dataset, in case of no incoming packets
    if len(in_packet_seconds) == 0:
        per_second_values_in = [{f"{attribute}_in": 0 for attribute in per_second_attributes}]

    # extract per second values for outgoing connections
    per_second_values_out = list()
    for second in out_packet_seconds:
        values_this_second = dict()
        for src, dest, data in outgoing:
            for attribute in per_second_attributes:
                # extract attribute values for the second, if they exist, else set to 0
                if attribute in data:
                    attribute_values = data[attribute]
                    if second in attribute_values:
                        if attribute not in values_this_second:
                            values_this_second[f"{attribute}_out"] = attribute_values[second]
                        else:
                            values_this_second[f"{attribute}_out"] += attribute_values[second]
                    else:
                        values_this_second[f"{attribute}_out"] = 0
                else:
                    values_this_second[f"{attribute}_out"] = 0

        # add to per second values
        per_second_values_out.append(values_this_second)

    # still add attributes to dataset, in case of no outgoing packets
    if len(out_packet_seconds) == 0:
        per_second_values_out = [{f"{attribute}_out": 0 for attribute in per_second_attributes}]

    """
    Process meta data for the node, such as time offsets between packets
    """

    # get time diff to last outgoing packet
    diff_to_last_outgoing = [0] + np.diff(out_packet_seconds).tolist()
    # get time diff to last incoming packet
    diff_to_last_incoming = [0] + np.diff(in_packet_seconds).tolist()

    # get the second where each incoming connection was first seen by extracting it from the first packet second
    incoming_first_seen = [list(data["packets_per_second"].keys())[0] for src, dest, data in incoming]

    # get new incoming connections per second by adding the first_seen occurrences for each second
    incoming_new_per_second = [sum([1 for first_seen in incoming_first_seen if first_seen == second])
                               for second in in_packet_seconds]

    # get the total number of incoming connection made per second, is sum of new connections per second until current index
    total_incoming_per_second = [sum(incoming_new_per_second[:index + 1]) for index in range(len(in_packet_seconds))]

    # get the second where each outgoing connection was first seen by extracting it from the first packet second
    outgoing_first_seen = [list(data["packets_per_second"].keys())[0] for src, dest, data in outgoing]

    # get new outgoing connections per second by adding the first_seen occurrences for each second
    outgoing_new_per_second = [sum([1 for first_seen in outgoing_first_seen if first_seen == second])
                               for second in out_packet_seconds]

    # get the total number of outgoing connection made per second
    total_outgoing_per_second = [sum(outgoing_new_per_second[:index + 1]) for index in range(len(out_packet_seconds))]

    # get total incoming ports
    # incoming_ports = list()
    # for src, dest, data in incoming:
    #     incoming_ports = [port for port in data["src_ports"] if port not in incoming_ports]
    #
    # if len(incoming_ports) == 0:
    #     incoming_ports = [0]
    #
    # # get total outgoing ports
    # outgoing_ports = list()
    # for src, dest, data in outgoing:
    #     outgoing_ports = [port for port in data["dst_ports"] if port not in outgoing_ports]
    #
    # if len(outgoing_ports) == 0:
    #     outgoing_ports = [0]

    per_second_incoming_df = pd.DataFrame(per_second_values_in, index=in_packet_seconds)
    incoming_data = {
                        "diff_to_last_incoming": diff_to_last_incoming,
                        "new_incoming_at_second": incoming_new_per_second,
                        "total_incoming_at_second": total_incoming_per_second,
                    } | per_connection_values_in
    # set index to the packet seconds
    per_node_incoming_df = pd.DataFrame(incoming_data, index=in_packet_seconds)
    incoming_df = pd.merge(per_node_incoming_df, per_second_incoming_df, how="outer", left_index=True,
                           right_index=True)

    per_second_outgoing_df = pd.DataFrame(per_second_values_out, index=out_packet_seconds)
    outgoing_data = {
                        "diff_to_last_outgoing": diff_to_last_outgoing,
                        "new_outgoing_at_second": outgoing_new_per_second,
                        "total_outgoing_at_second": total_outgoing_per_second,
                    } | per_connection_values_out
    # set index to the packet seconds
    per_node_outgoing_df = pd.DataFrame(outgoing_data, index=out_packet_seconds)
    outgoing_df = pd.merge(per_node_outgoing_df, per_second_outgoing_df, how="outer", left_index=True,
                           right_index=True)

    # merge incoming and outgoing data on packet seconds
    data = pd.merge(incoming_df, outgoing_df, how="outer", left_index=True, right_index=True)

    # fill missing values with last set value (forward fill)
    data = data.ffill()

    # fill missing values with 0, there can still be missing value in the first rows
    data = data.fillna(0)

    # parse all columns as int
    data = data.astype(int)

    # add identifier for the node
    data["node"] = node
    return data
