import pandas as pd
import os
import random
import argparse
import numpy as np

from seed import set_seed

set_seed()

parser = argparse.ArgumentParser()
parser.add_argument(
    "--motif_number", nargs='?', type=int, default=1000,
    help="number of motifs"
)
parser.add_argument("--dir", nargs='?', default='./motifs', help="output dir")
args = parser.parse_args()

motif_dir = args.dir
motif_number = args.motif_number

distibution = [0, 1, 3, 5]


def read_real_dblp():
    nodes = pd.read_csv('./ref/node.csv')
    edges = pd.read_csv('./ref/link.csv')
    nodes = nodes.iloc[:, [0, 1]]
    # First filter out the edge types
    types = pd.unique(edges['2']).tolist()
    # Construct a matrix that = node numbers * edge types
    edges_counts = pd.DataFrame(
        index=range(edges.max()[0] + 1), columns=range(len(types))
    ).reset_index()
    # Count each type by each node
    for type in types:
        type_edges = edges[edges['2'] == type]
        edges_counts[type] = edges_counts['index'].map(
            type_edges['0'].value_counts(sort=False)
        )
    edges_counts = edges_counts.drop('index', axis=1)
    edges_counts['id'] = nodes['1']

    return edges_counts


edges_counts = read_real_dblp()

for motif in range(motif_number):
    nodes = pd.read_csv(
        os.path.join(motif_dir, 'nodes_' + str(motif) + '.csv')
    )
    edges = pd.read_csv(
        os.path.join(motif_dir, 'edges_' + str(motif) + '.csv')
    )
    nodes.columns = nodes.columns.astype(int)
    edges.columns = edges.columns.astype(int)
    counts = pd.DataFrame({
        'type': range(4)
    }).join(nodes[1].value_counts()).fillna(0)
    counts = counts.drop([0])
    # add paper to author first
    paper_number = pd.Series(edges_counts[2]).dropna().sample(n=1
                                                              ).tolist()[0] - 1
    # if paper_number > 49:
    #     paper_number = 49
    for _ in range(int(paper_number)):
        distibution = [0, 0, 1, int(np.random.normal(8, 0.4, 1)[0])]
        new_paper_id = nodes[0].max() + 1
        nodes = pd.concat(
            [nodes, pd.DataFrame([[new_paper_id, 1]])], ignore_index=True
        )
        edges = pd.concat(
            [
                edges,
                pd.DataFrame([[0, new_paper_id, 1], [new_paper_id, 0, 0]])
            ], ignore_index=True
        )
        for type, d in enumerate(distibution):
            for bonus in range(d):
                new_id = nodes[0].max() + 1

                #### Add the node to csv
                nodes = pd.concat(
                    [nodes, pd.DataFrame([[new_id, type]])], ignore_index=True
                )
                #### Add edges
                edge_type = (
                    type - 1
                ) * 2  ##### Here are same fix edge types define
                edges = pd.concat(
                    [
                        edges,
                        pd.DataFrame(
                            [
                                [new_paper_id, new_id, edge_type],
                                [new_id, new_paper_id, edge_type + 1]
                            ]
                        )
                    ], ignore_index=True
                )

    paper_number = pd.Series(edges_counts[2]).dropna().sample(n=1
                                                              ).tolist()[0] - 1
    for _ in range(int(paper_number)):
        distibution = [0, 0, 1, int(np.random.normal(8, 0.4, 1)[0])]
        new_paper_id = nodes[0].max() + 1
        nodes = pd.concat(
            [nodes, pd.DataFrame([[new_paper_id, 1]])], ignore_index=True
        )
        edges = pd.concat(
            [
                edges,
                pd.DataFrame([[1, new_paper_id, 1], [new_paper_id, 1, 0]])
            ], ignore_index=True
        )
        for type, d in enumerate(distibution):
            for bonus in range(d):
                new_id = nodes[0].max() + 1
                #### Add the node to csv
                nodes = pd.concat(
                    [nodes, pd.DataFrame([[new_id, type]])], ignore_index=True
                )
                #### Add edges
                edge_type = (
                    type - 1
                ) * 2  ##### Here are same fix edge types define
                edges = pd.concat(
                    [
                        edges,
                        pd.DataFrame(
                            [
                                [new_paper_id, new_id, edge_type],
                                [new_id, new_paper_id, edge_type + 1]
                            ]
                        )
                    ], ignore_index=True
                )

    distibution = [0, 0, 1, int(np.random.normal(8, 0.4, 1)[0])]
    distibution = [i if i >= 0 else 0 for i in distibution]
    # print(distibution)
    for i, row in counts.iterrows():
        for bonus in range(int(distibution[i] - row.iloc[1])):
            if (int(distibution[i] - row.iloc[1])) <= 0:
                print(int(distibution[i] - row.iloc[1]))
            #### Select the connect node type and location
            node_type = row['type']
            # legal_connect_types = legal_connects[node_type] # find the legal connect edges of the bonus node
            # ### The differnt between random choice the edge type first or just random choice the connect node
            # connect_node_type = random.choice(legal_connect_types) # find the legal connect nodes
            # legal_nodes = nodes[nodes[1]==connect_node_type][0].to_list() # find the
            # connect_location = random.choice(legal_nodes)
            new_id = nodes[0].max() + 1

            #### Add the node to csv
            nodes = pd.concat(
                [nodes, pd.DataFrame([[new_id, node_type]])], ignore_index=True
            )
            #### Add edges
            edge_type = (
                node_type - 1
            ) * 2  ##### Here are same fix edge types define
            edges = pd.concat(
                [
                    edges,
                    pd.DataFrame(
                        [[2, new_id, edge_type], [new_id, 2, edge_type + 1]]
                    )
                ], ignore_index=True
            )

            new_id = nodes[0].max() + 1

            #### Add the node to csv
            nodes = pd.concat(
                [nodes, pd.DataFrame([[new_id, node_type]])], ignore_index=True
            )
            #### Add edges
            edge_type = (
                node_type - 1
            ) * 2  ##### Here are same fix edge types define
            edges = pd.concat(
                [
                    edges,
                    pd.DataFrame(
                        [[3, new_id, edge_type], [new_id, 3, edge_type + 1]]
                    )
                ], ignore_index=True
            )
            # print(counts['type'],counts[1])

    edges = edges.astype(int)
    nodes = nodes.astype(int)
    # print(nodes)
    # print(edges)
    # for bonus in range(10): # deal with the random part if the  bonus node is not 3
    #     #### Select the connect node type and location
    #     node_type =  random.choice([1,2,3]) # choice the node type that will add
    #     legal_connect_types = legal_connects[node_type] # find the legal connect edges of the bonus node
    #     ### The differnt between random choice the edge type first or just random choice the connect node
    #     connect_node_type = random.choice(legal_connect_types) # find the legal connect nodes
    #     legal_nodes = nodes[nodes[1]==connect_node_type][0].to_list() # find the
    #     connect_location = random.choice(legal_nodes)
    #     new_id = nodes[0].max()+1
    #
    #     #### Add the node to csv
    #     nodes = pd.concat([nodes, pd.DataFrame([[new_id, node_type]])], ignore_index=True)
    #     #### Add edges
    #     edge_type = (node_type-1)*2 ##### Here are same fix edge types define
    #     edges = pd.concat([edges, pd.DataFrame([[connect_location, new_id, edge_type], [new_id, connect_location, edge_type+1]])], ignore_index=True)

    nodes.to_csv(
        os.path.join(motif_dir, 'nodes_b_' + str(motif) + '.csv'), index=False
    )
    edges.to_csv(
        os.path.join(motif_dir, 'edges_b_' + str(motif) + '.csv'), index=False
    )
