import networkx as nx
import random
import matplotlib.pyplot as plt
import sys
import numpy as np
from datetime import datetime
import os
from joblib import Parallel, delayed


def plotData(conf):
    for key, value in conf['y'].items():
        plt.plot(conf['x'], value, label=f"{conf['label']}: {key}")
    plt.xlabel(conf['xlabel'])
    plt.ylabel(conf['ylabel'])
    plt.title(conf['title'])
    plt.legend()
    plt.grid(True)
    plt.savefig(conf['plotPath'])
    plt.clf()

def isGroupRecoverable(group, nodes, chi1, chi2, size, sizeK):
    uniqueLines = set()
    for g in group:
        linesSelected = set(random.sample(list(range(1, size + 1)), chi1 if g < nodes[0] else chi2))
        uniqueLines.union(linesSelected)
    return len(uniqueLines) >= sizeK

def isGConnected(deg, nodes, mal, config):
    G = nx.random_regular_graph(deg, sum(nodes))
    malNodes = random.sample(list(G.nodes()), k=mal * sum(nodes) // 100)
    for mn in malNodes:
        G.remove_node(mn)
    
    if nx.is_connected(G): return True
    else:
        isRecoverable = True
        chiR1 = config['custodyR'] * config['validatorPerNode1']
        chiR2 = config['custodyR'] * config['validatorPerNode2']
        rows = config['numberOfRows']
        rowsK = config['numberOfRowsK']
        for group in nx.connected_components(G):
            isRecoverable = isRecoverable and isGroupRecoverable(group, nodes, chiR1, chiR2, rows, rowsK)
        
        return isRecoverable

def getNodeCountPerColumn(config, numOfNodes):
    numberOfCols = config['numberOfColumns']
    chiC1 = config['custodyC'] * config['validatorPerNode1']
    chiC2 = config['custodyC'] * config['validatorPerNode2']
    node1Count = int(numOfNodes * config['class1ratio'])
    nodeCountPerColumn = dict()
    for _ in range(numOfNodes):
        colsSelected = random.sample(list(range(1, numberOfCols + 1)), chiC1 if _ < node1Count else chiC2)
        for col in colsSelected:
            if col in nodeCountPerColumn.keys():
                nodeCountPerColumn[col][0 if _ < numOfNodes else 1] += 1
            else:
                nodeCountPerColumn[col] = [0, 0]
                nodeCountPerColumn[col][0 if _ < numOfNodes else 1] = 1
    
    return nodeCountPerColumn

def runOnce(deg, nodeCountPerCol, malNodesPercentage, config):
    isParted = False
    partCount = 0
    isPartedCount = 0
    for col in nodeCountPerCol.keys():
        nodes = nodeCountPerCol[col]
        if not isGConnected(deg, nodes, malNodesPercentage, config):
            if not isParted: isParted = True
            partCount += 1
    
    if isParted: isPartedCount += 1
    
    return isPartedCount, partCount

def study(config):
    nnPartPercentages = dict()
    nnAvgDisconnectedCols = dict()
    
    for nn in config['numberOfNodes']:
        print(f"\nNumber of Nodes: {nn}")
        
        partPercentages = list()
        avgDisconnectedCols = list()
        for mal in config['mals']:
            isPartedCount = partCount = 0
            nodeCountPerColumn = getNodeCountPerColumn(config, nn)
            results = Parallel(-1)(delayed(runOnce)(config['deg'], nodeCountPerColumn, mal, config) for _run in range(config['runs']))
            isPartedCount = sum([res[0] for res in results])
            partCount = sum([res[1] for res in results])
            partPercentages.append(isPartedCount * 100 / config['runs'])
            avgDisconnectedCols.append(partCount / config['runs'])
            print(f"Malicious Nodes: {mal}%, Partition Percentage: {partPercentages[-1]}, Avg. Partitions: {avgDisconnectedCols[-1]}")
        
        nnPartPercentages[nn] = partPercentages
        nnAvgDisconnectedCols[nn] = avgDisconnectedCols
    
    now = datetime.now()
    execID = now.strftime("%Y-%m-%d_%H-%M-%S_")+str(random.randint(100,999))
    newpath = f"ConnectivityTest/MaliciousNodesVsNumberOfNodes/results/{execID}/"
    if not os.path.exists(newpath): os.makedirs(newpath)
    
    conf1 = {
        'x': config['mals'],
        'y': nnPartPercentages,
        'label': "Nodes",
        'xlabel': "Malicious Node (%)",
        'ylabel': "Partition Possibility (%)",
        "title": "Possibility of Network Graph Get Partitioned for Malicious Nodes",
        "plotPath": f"{newpath}prob.png"
    }
    
    conf2 = {
        'x': config['mals'],
        'y': nnAvgDisconnectedCols,
        'label': "Nodes",
        'xlabel': "Malicious Node (%)",
        'ylabel': "Avg. Disconnected Columns",
        "title": "Malicious Nodes (%) vs. Disconnected Columns",
        "plotPath": f"{newpath}num.png"
    }
    
    plotData(conf1)
    plotData(conf2)


# Configuration
config = {
    'runs': 20,
    'deg': 8,
    'mals': range(5, 100, 5),
    'numberOfColumns': 128,
    'numberOfColumnsK': 64,
    'numberOfRows': 128,
    'numberOfRowsK': 64,
    'custodyC': 4,
    'custodyR': 4,
    'class1ratio': 0.8,
    'validatorPerNode1': 1,
    'validatorPerNode2': 8,
    'numberOfNodes': [int(_) for _ in (np.logspace(2, 4, 5, endpoint=True, base=10) * 5)]
}

if __name__ == "__main__":
    study(config)