import pickle
from sklearn.ensemble import ExtraTreesClassifier, RandomForestClassifier
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
from sklearn.model_selection import cross_val_score
from fastsklearnfeature.interactiveAutoML.new_bench.multiobjective.metalearning.analyse.time_measure import get_recall
from fastsklearnfeature.interactiveAutoML.new_bench.multiobjective.metalearning.analyse.time_measure import time_score2
from fastsklearnfeature.interactiveAutoML.new_bench.multiobjective.metalearning.analyse.time_measure import get_avg_runtime
from fastsklearnfeature.interactiveAutoML.new_bench.multiobjective.metalearning.analyse.time_measure import get_optimum_avg_runtime

from sklearn.metrics import make_scorer
from sklearn.dummy import DummyClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.tree import DecisionTreeClassifier
from sklearn import tree
from sklearn.tree import export_graphviz
from subprocess import call
from sklearn.model_selection import LeaveOneGroupOut
from sklearn.model_selection import GroupKFold
from sklearn.model_selection import RandomizedSearchCV
import copy
import glob
import matplotlib.pyplot as plt

def is_pareto_efficient_simple(costs):
    """
    Find the pareto-efficient points
    :param costs: An (n_points, n_costs) array
    :return: A (n_points, ) boolean array, indicating whether each point is Pareto efficient
    """
    is_efficient = np.ones(costs.shape[0], dtype = bool)
    for i, c in enumerate(costs):
        if is_efficient[i]:
            is_efficient[is_efficient] = np.any(costs[is_efficient]<c, axis=1)  # Keep any point with a lower cost
            is_efficient[i] = True  # And keep self
    return is_efficient

'''
Exhaustive Search & $x \pm y$ && x\\
Forward Selection & $x \pm y$ && x\\
Backward Selection & $x \pm y$ && x\\
Forward Floating Selection & $x \pm y$ && x\\
Backward Floating Selection & $x \pm y$ && x\\
Recursive Feature Elimination & $x \pm y$ && x\\
Hyperopt(KBest(Fisher Score)) & $x \pm y$ && x\\
Hyperopt(KBest(ReliefF)) & $x \pm y$ && x\\
Hyperopt(KBest(Mutual Information)) & $x \pm y$ && x\\
Hyperopt(KBest(FCBF)) & $x \pm y$ && x\\
Hyperopt(KBest(MCFS)) & $x \pm y$ && x\\
Hyperopt(KBest(Variance)) & $x \pm y$ && x\\
Hyperopt(KBest($\chi^2$)) & $x \pm y$ && x\\
Ranking-free Hyperopt & $x \pm y$ && x\\
Ranking-free Simulated Annealing & $x \pm y$ && x\\
Ranking-free NSGA-II & $x \pm y$ && x\\ \midrule
Meta-learned Strategy Choice & $x \pm y$ && x\\
'''

map_dataset2name = {}
map_dataset2name['31'] = 'German Credit'
map_dataset2name['802'] = 'Primary Biliary Cirrhosis'
map_dataset2name['1590'] = 'Adult'
map_dataset2name['1461'] = 'Bank Marketing'
map_dataset2name['42193'] = 'COMPAS'
map_dataset2name['1480'] = 'Indian Liver Patient'
#map_dataset2name['804'] = 'hutsof99_logis'
map_dataset2name['42178'] = 'Telco Customer Churn'
map_dataset2name['981'] = 'KDD Internet Usage'
map_dataset2name['40536'] = 'Speed Dating'
map_dataset2name['40945'] = 'Titanic'
map_dataset2name['451'] = 'Irish Educational Transitions'
#map_dataset2name['945'] = 'Kidney'
map_dataset2name['446'] = 'Leptograpsus crabs'
map_dataset2name['1017'] = 'Arrhythmia'
map_dataset2name['957'] = 'Brazil Tourism'
map_dataset2name['41430'] = 'Diabetic Mellitus'
map_dataset2name['1240'] = 'AirlinesCodrnaAdult'
map_dataset2name['1018'] = 'IPUMS Census'
#map_dataset2name['55'] = 'Hepatitis'
map_dataset2name['38'] = 'Thyroid Disease'
map_dataset2name['1003'] = 'Primary Tumor'
map_dataset2name['934'] ='Social Mobility'

mappnames = {1:'TPE(Variance)',
			 2: 'TPE($\chi^2$))',
			 3:'TPE(FCBF))',
			 4: 'TPE(Fisher Score))',
			 5: 'TPE(Mutual Information))',
			 6: 'TPE(MCFS))',
			 7: 'TPE(ReliefF))',
			 8: 'TPE(no ranking)',
             9: 'Simulated Annealing(no ranking)',
			 10: 'NSGA-II(no ranking)',
			 11: 'Exhaustive Search(no ranking)',
			 12: 'Forward Selection(no ranking)',
			 13: 'Backward Selection(no ranking)',
			 14: 'Forward Floating Selection(no ranking)',
			 15: 'Backward Floating Selection(no ranking)',
			 16: 'RFE(Logistic Regression)'
			 }

names = ['accuracy',
	 'fairness',
	 'k_rel',
	 'k',
	 'robustness',
	 'privacy',
	 'search_time',
	 'cv_acc - acc',
	 'cv_fair - fair',
	 'cv_k - k rel',
	 'cv_k - k',
	 'cv_robust - robust',
     'cv time',
	 'rows',
	 'columns']

def print_constraints_2(features):

	my_str = ''
	for i in range(len(names)):
		my_str += names[i] + ': ' + str(features[i]) + ' '
	print(my_str)

#logs_adult = pickle.load(open('/home/felix/phd/meta_learn/classification/metalearning_data_adult.pickle', 'rb'))
#logs_heart = pickle.load(open('/home/felix/phd/meta_learn/classification/metalearning_data_heart.pickle', 'rb'))

#get all files from folder

all_files = glob.glob("/home/felix/phd/meta_learn/new_bugfree/*.pickle")

dataset = {}
for afile in all_files:
	data = pickle.load(open(afile, 'rb'))
	for key in data.keys():
		if not key in dataset:
			dataset[key] = []
		dataset[key].extend(data[key])

assert len(dataset['success_value']) == len(dataset['best_strategy'])

joined_strategies = []

map_data_2_coverage = {}

for run in range(len(dataset['best_strategy'])):

	if dataset['best_strategy'][run] != 0:
		for s in range(1, len(mappnames) + 1):

			if not dataset['dataset_id'][run] in map_data_2_coverage:
				map_data_2_coverage[dataset['dataset_id'][run]] = {}
			if not s in map_data_2_coverage[dataset['dataset_id'][run]]:
				map_data_2_coverage[dataset['dataset_id'][run]][s] = []

			if s in dataset['success_value'][run] and len(dataset['success_value'][run][s]) > 0 and dataset['success_value'][run][s][0]==True:
				map_data_2_coverage[dataset['dataset_id'][run]][s].append(True)
			else:
				map_data_2_coverage[dataset['dataset_id'][run]][s].append(False)

from matplotlib import pyplot as plt

def bar_plot(ax, data, colors=None, total_width=0.8, single_width=1, legend=True):
	# Check if colors where provided, otherwhise use the default color cycle
	if colors is None:
		colors = plt.rcParams['axes.prop_cycle'].by_key()['color']

	# Number of bars per group
	n_bars = len(data)

	# The width of a single bar
	bar_width = total_width / n_bars

	# List containing handles for the drawn bars, used for the legend
	bars = []

	# Iterate over all data
	for i, (name, values) in enumerate(data.items()):
		# The offset in x direction of that bar
		x_offset = (i - n_bars / 2) * bar_width + bar_width / 2

		# Draw a bar for every value of that type
		for x, y in enumerate(values):
			bar = ax.bar(x + x_offset, y, width=bar_width * single_width, color=colors[i % len(colors)])

		# Add a handle to the last drawn bar, which we'll need for the legend
		bars.append(bar[0])

	# Draw legend if we need
	if legend:
		ax.legend(bars, data.keys())

data = {}

for key_data, value_data in map_data_2_coverage.items():
	for i in range(1, 17):
		if not mappnames[i] in data:
			data[mappnames[i]] = []
		data[mappnames[i]].append(np.sum(value_data[i]) / float(len(value_data[i])))

fig, ax = plt.subplots()
bar_plot(ax, data, total_width=.8, single_width=.9)
plt.show()


from example import example

from qtpy import QtCore
import pytest
from pytest_bdd import (
    given,
    scenario,
    then,
    when,
    parsers,
    scenarios,
)

@pytest.fixture
def window():
    return example.ExampleWindow()

scenarios('example.feature')

@given('The window')
def the_window_is_open(qtbot, window):
    """The window is open."""
    qtbot.addWidget(window)

@when('I press the button')
def i_press_the_button(qtbot, window):
    """I press the button."""
    qtbot.mouseClick(window.btn, QtCore.Qt.LeftButton)

@when(parsers.parse('I type "{line}"'))
def type_in_line(line, window, qtbot):
    """I type "line"."""
    window.line_edit.clear()
    qtbot.keyClicks(window.line_edit, line)

@then(parsers.parse('The text is "{line}"'))
def the_text_is(window, line):
    """The text is "line"."""
    assert window.line_edit.text() == line


from collections import deque, defaultdict

def part1(players, marbles):
    board = [0, 4, 2, 5, 1, 6, 3, 7]
    player = 8
    curr = 7
    points = dict()

    for marble in range(8, marbles+1):
        if (marble % 23 == 0):
            curr = curr - 7
            if (curr < 0):
                curr = curr + len(board)
            left = board[curr]
            board.pop(curr)
            points[player] = points.setdefault(player, 0) + marble + left
        else:
            curr = (curr + 2) % len(board)
            board.insert(curr, marble)
        player = (player + 1) % players
    print(max(points.values()))
    #print(board)

def part2(players, marbles):
    points = defaultdict(int)
    board = deque([0])

    for marble in range(1, marbles+1):
        if marble % 23 == 0:
            board.rotate(7)
            points[marble % players] += marble + board.pop()
            board.rotate(-1)
        else:
            board.rotate(-1)
            board.append(marble)
    print(max(points.values()))

if __name__ == "__main__":
    part2(9, 25)
    part2(10, 1618)
    part2(13, 7999)
    part2(17, 1104)
    part2(412, 71646)
    part2(412, 7164600)


from typing import List
import json
import logging
import pathlib
import random
import time

import numpy as np
import torch

from fairdiplomacy.pydipcc import Game
from fairdiplomacy.models.consts import POWERS
from fairdiplomacy.agents import build_agent_from_cfg

def load_games(game_jsons) -> List[Game]:
    if not game_jsons:
        return [Game()]

    games = []
    for path in game_jsons:
        with open(path) as f:
            game_json = f.read()
        games.append(Game.from_json(game_json))
    return games

def run(cfg: "conf.conf_pb2.BenchmarkAgent"):

    logger = logging.getLogger("timing")
    logger.setLevel(logging.DEBUG)
    timing_outpath = pathlib.Path(".").resolve() / "timing.log"
    logger.info("Will write timing logs to %s", timing_outpath)
    logger.addHandler(logging.FileHandler(timing_outpath))

    random.seed(cfg.seed)
    np.random.seed(cfg.seed)
    torch.manual_seed(cfg.seed)

    agent = build_agent_from_cfg(cfg.agent)
    games = load_games(cfg.game_jsons)

    logger.info("Warmup")
    agent.get_orders(games[0], POWERS[0])

    logger.info("Running benchmark")
    times = []
    for _ in range(cfg.repeats):
        times.append([])
        for game in games:
            times[-1].append([])
            for power in POWERS:
                start = time.time()
                agent.get_orders(game, power)
                times[-1][-1].append(time.time() - start)
    times = np.array(times)
    logger.info("Total time: %s", times.sum())
    per_repeat = times.sum(-1).sum(-1)
    logger.info("Total time per repeat: %s +- %s", per_repeat.mean(), per_repeat.std())
    per_call = times.flatten()
    logger.info("Total time per call: %s +- %s", per_call.mean(), per_call.std())


from __future__ import division
import numpy as np
from matplotlib import pyplot as plt
from matplotlib import rc
from numpy.linalg import norm
import time

from math import sqrt
from numba import jit,vectorize,int64
rc('mathtext', default='regular')

L = 5.173403878  #cubic box side length rou*sigma^3=0.78
def LJ(r):  # r is a variable of distance between 2 particles
    if r == 0:
        P = 0
    elif r > L / 2:  # potential cutoff
        P = 0
    else:
        P = 4 * (r ** -12 - r ** -6)  # P is potential between 2 particles
    return P

def Force(a):
    r = norm(a)
    if r == 0:
        F = 0
    elif r > L / 2: #Force cutoff
        F = 0
    else:
        F = -24 * (-2 * r ** -13 + r ** -7)
    return F

def Z(r):
    if r == 0:
        a = 1
    else:
        a = r
    return a

def main():
    tb=time.time()
    N = 108
    Steps = 1600

    from Config import R #initial config position
    from Config import V #velocity
    RDF=np.zeros(500)
    norm=np.linalg.norm
    M = 1
    dt = 0.0025
    S = []  # steps
    P = []  # positional
    ff = []  # force constants
    ke = []  # kinetics
    p = []  # potential
    e = []  # energy
    rrr=[]
    vvv=[]

    V_vcc=np.zeros((Steps//4+1,324))
    Coeff=np.zeros(100)
    F = np.zeros(N * 3)
# Velocity Verlet Algorithm
    for step in range(Steps):
        V = V + F * dt / (2 * M)
        R = R + V * dt

    # in this loop time mooo forward
        for val in range(N):  # Velocity and position update according to acc
            for j in range(3):
                # positional pbc begins
                while R[val + j * N] > L:
                    R[val + j * N] = R[val + j * N] - L
                while R[val + j * N] < 0:
                    R[val + j * N] = R[val + j * N] + L
                pass
                # PBC of position ends
    #this loop updated velocity and position

        F = np.zeros(N * 3)  # Force update loop
        P = 0
        # reconstruct force loop
        for val in range(N):  # we gonna process force on val^{th} particle
            for i in range(N):  # it is the i^{th} particle's force on the val^{th} particle
                a = np.zeros(3)
                for j in range(3):
                    #force pbc: pseudo vector 'a' generation
                    if R[val + j * N] - R[i + j * N] > L / 2:
                        a[j] = R[val + j * N] - R[i + j * N] - L
                    elif R[val + j * N] - R[i + j * N] < -L / 2:
                        a[j] = R[val + j * N] - R[i + j * N] + L
                    else:
                        a[j] = R[val + j * N] - R[i + j * N]
                # a pbc precess of i^{th} particle 's position to val
                d = norm(a)
                P = P + LJ(d)
                F[val + 0 * N] = F[val + 0 * N] + a[0] * Force(d) / Z(d)
                F[val + 1 * N] = F[val + 1 * N] + a[1] * Force(d) / Z(d)
                F[val + 2 * N] = F[val + 2 * N] + a[2] * Force(d) / Z(d)
                if step % 4==1 and d<0.5*L and d !=0:
                    b=((d/L)*1000) // 1
                    b=int(b)
                    RDF[b]=RDF[b]+1  #rdf modules
        A = P / 2
        p.append(A)

        V= V+ F * dt / (2 * M)
        # velocity 2nd update
        KE=0.5*sum(V*V)
        ke.append(KE)

        if step%4==0:
            s = step // 4
            V_vcc[s] = V
        else:pass

        E = KE + P / 2
        print(KE)
        e.append(E)
        S.append(step)
        print(step)

    for j in range(100):
        for i in range(Steps//4+1-j):
            Coeff[j]=Coeff[j]+np.sum(V_vcc[i]*V_vcc[j+i])/np.sum(V_vcc[i]*V_vcc[i])

#DATA WRITE MODULE

    ama = max(ke)
    ami = min(ke)
    sig = np.var(ke)
    ava=np.average(ke)

    print(ama, ami, sig, ava)

    for i in range(100):
        Coeff[i]=Coeff[i]/(Steps//4-i)

    c=np.arange(-0.01,1,0.01)
    y=np.zeros(100)
    fig = plt.figure()
    ax0 = fig.add_subplot(111)
    lns0 = ax0.plot(c,Coeff,'red', label='Cvv')
    lns1 =ax0.plot(c,y, 'blue')
    lns = lns1
    ax0.set_xlabel('steps 0.0025 tao per step')
    ax0.set_ylabel("energy emma=118.7K*kb")
    plt.legend()
    plt.show()

    te=time.time()
    print(te-tb,'s')

main()


# https://github.com/aladdinpersson/Machine-Learning-Collection/blob/master/ML/Pytorch/image_segmentation/semantic_segmentation_unet/model.py
# https://www.youtube.com/watch?v=IHq1t7NxS8k&t=1008s
from itertools import zip_longest

import torch
import torch.nn as nn
import torchvision.transforms.functional as TF
import torch.nn.functional as F

from torchsummary import summary

GROWTH_K = 6
DROPOUT_RATE = 0.3

class DilationBlock(nn.Module):
    def __init__(self,  dilation_sz, in_channels_sz=294, filters=1024):
        super(DilationBlock, self).__init__()

        self.dilation_block = nn.Sequential(
            nn.Conv2d(in_channels_sz, filters, kernel_size=3, stride=1, padding=dilation_sz, dilation=dilation_sz, bias=False),
            nn.ReLU(inplace=True),
            nn.BatchNorm2d(filters)
        )

    def forward(self, x):
        return self.dilation_block(x)

class Dilation(nn.Module):
    def __init__(self, in_channels_sz=294, dilation_sz=[6, 12, 18, 24]):
        super(Dilation, self).__init__()

        self.dilations = nn.ModuleList()

        for d in dilation_sz:
            self.dilations.append(DilationBlock(dilation_sz=d, in_channels_sz=in_channels_sz))

    def forward(self, x):
        # create 4 different dilation blocks and then concat
        y = []
        for dilation in self.dilations:
            y.append(dilation(x))

        x = torch.cat([*y], dim=1)
        return x # default: torch.Size([1, 4096, 8, 8])

class UpConv2d(nn.Module):
    def __init__(self, input_channels_sz, out_channels_sz ):
        super(UpConv2d, self).__init__()

        self.up_conv = nn.Sequential(
            nn.Conv2d(input_channels_sz, out_channels_sz, kernel_size=3, padding=1, bias=False),
            nn.ReLU(inplace=True),
            nn.BatchNorm2d(out_channels_sz),
            nn.Conv2d(out_channels_sz, out_channels_sz, kernel_size=3, padding=1, bias=False),
            nn.ReLU(inplace=True),
            nn.BatchNorm2d(out_channels_sz),
        )

    def forward(self, x):
        return self.up_conv(x)

class BottleneckBlock(nn.Module):
    def __init__(self, in_planes, out_planes, dropRate=DROPOUT_RATE):
        super(BottleneckBlock, self).__init__()
        inter_planes = out_planes * 4

        self.bn1 = nn.BatchNorm2d(in_planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv1 = nn.Conv2d(in_planes, inter_planes, kernel_size=3, stride=1,
                               padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(inter_planes)
        self.conv2 = nn.Conv2d(inter_planes, out_planes, kernel_size=3, stride=1,
                               padding=1, bias=False)
        self.droprate = dropRate

    def forward(self, x):
        out = self.conv1(self.relu(self.bn1(x)))
        if self.droprate > 0:
            out = F.dropout(out, p=self.droprate, inplace=False, training=self.training)
        out = self.conv2(self.relu(self.bn2(out)))
        if self.droprate > 0:
            out = F.dropout(out, p=self.droprate, inplace=False, training=self.training)
        return torch.cat([x, out], 1)

class TransitionBlock(nn.Module):
    def __init__(self, in_planes, out_planes, dropRate=DROPOUT_RATE):
        super(TransitionBlock, self).__init__()
        self.bn1 = nn.BatchNorm2d(in_planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv1 = nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=1,
                               padding=0, bias=False)
        self.droprate = dropRate

    def forward(self, x):
        out = self.conv1(self.relu(self.bn1(x)))
        if self.droprate > 0:
            out = F.dropout(out, p=self.droprate, inplace=False, training=self.training)
        return F.avg_pool2d(out, 2)

class DenseBlock(nn.Module):
    def __init__(self,  nb_layers, in_planes, block=BottleneckBlock, growth_rate=GROWTH_K, dropRate=DROPOUT_RATE):
        super(DenseBlock, self).__init__()
        self.layer = self._make_layer(block, in_planes, growth_rate, nb_layers, dropRate)

    def _make_layer(self, block, in_planes, growth_rate, nb_layers, dropRate):
        layers = []
        for i in range(nb_layers):
            layers.append(block(in_planes+i*growth_rate, growth_rate, dropRate))
        return nn.Sequential(*layers)

    def forward(self, x):
        return self.layer(x)

# ========= START FROM HERE =======

class UNET(nn.Module):
    def __init__(
            self, dense_down_in_channels=[64, 6, 6, 6],
            trans_down_in_channels=[136, 150, 438],
            nb_layers=[3, 6, 18, 12],
            up_list=[4096, 512, 256, 128, 64],
            up_conv_in=[950, 406, 264, 128],

    ):
        super(UNET, self).__init__()
        self.ups_trans = nn.ModuleList()
        self.ups_conv = nn.ModuleList()
        self.downs_dense = nn.ModuleList()
        self.downs_trans = nn.ModuleList()
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)

        # fist part
        self.conv_start = nn.Sequential(
            nn.BatchNorm2d(1),
            nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False),
            nn.ReLU(inplace=True),
        )

        # Down part of UNet
        for dense_chan, nb_layer in zip(dense_down_in_channels, nb_layers):
            self.downs_dense.append(DenseBlock(in_planes=dense_chan, nb_layers=nb_layer, growth_rate=24))

        for trans_chan in trans_down_in_channels:
            self.downs_trans.append(TransitionBlock(in_planes=trans_chan, out_planes=6))

        self.dilation = Dilation()

        # upwards
        self.convT1 = nn.ConvTranspose2d(4096, 512, kernel_size=2, stride=2)

        for in_chan, out_chan in zip(up_list, up_list[1:]):
            self.ups_trans.append(nn.ConvTranspose2d(in_chan, out_chan, kernel_size=2, stride=2))

        up_conv_out = 512
        for up_in in up_conv_in:
            self.ups_conv.append(UpConv2d(input_channels_sz=up_in, out_channels_sz=up_conv_out))
            up_conv_out //= 2

        # downwards (end of UNet)
        self.end_of_UNet = nn.Sequential(
            nn.ConvTranspose2d(64, 64, kernel_size=2, stride=2),
            nn.BatchNorm2d(64),
            nn.Conv2d(64, 1, kernel_size=1, stride=1, bias=False),
            nn.Sigmoid()
        )

    def forward(self, x):
        skip_connections = []

        x = self.conv_start(x)
        skip_connections.append(x)

        x = self.pool(x)

        #  TODO fix, hack of None
        for down_dense, down_trans in zip_longest(self.downs_dense, self.downs_trans, fillvalue=None):
            x = down_dense(x)

            if down_trans!=None:
                skip_connections.append(x)
                x = down_trans(x)

        x = self.dilation(x)
        print("DDUnet shape:", x.shape)

        skip_connections = skip_connections[::-1]  # reverse list

        # upwards
        out_conv = 512
        for skip, up_trans, up_conv in zip(skip_connections, self.ups_trans, self.ups_conv):
            x = up_trans(x)
            x = torch.cat([x, skip], dim=1)
            x = up_conv(x)
            print(x.shape)
            out_conv /= 2

        # down again (end)
        x = self.end_of_UNet(x)
        print(x.shape)

        # TODO return lambdas
        map_lambda1 = torch.exp((2.0 - x)/(1.0 + x))
        map_lambda2 = torch.exp((1.0 + x)/(2.0 - x))
        return x

def test():
    x = torch.randn((1, 1, 256, 256))
    model = UNET()
    out_seg = model(x)
    # print(out_seg.shape)
    # print(x.shape)
    # print(model)
    # image = torch.rand((1, 572, 572))
    # summary(model, (1, 256, 256))
    # assert out_seg.shape == x.shape

    map_lambda1 = torch.exp((2.0 - out_seg) / (1.0 + out_seg))
    map_lambda2 = torch.exp((1.0 + out_seg) / (2.0 - out_seg))

if __name__ == "__main__":
    test()


from asyncio.windows_events import NULL
from Link import *
from Node import *
import random,math
import numpy as np
import networkx as nx
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
class Topo(object):#定义拓扑
    def __init__(self):
        self.node=[]
        self.edge=[]
        self.topo=[]

    def CreatTopo(self,node_edge_set):#输出邻接矩阵
        self.node=node_edge_set[0]
        self.edge=node_edge_set[1]
        #初始化拓扑,拓扑包含邻接表
        self.topo=[]
        for i in range(len(self.node)):#初始化邻接表
            t=[]
            for j in range(len(self.node)):
                t.append(Link())
            self.topo.append(t)

        #节点与数字映射
        i=0
        dictnode={}
        for node in self.node:
            dictnode[node]=i
            i+=1

        for edge in self.edge:
            self.topo[dictnode.get(edge.fr)][dictnode.get(edge.to)]=Link(dictnode.get(edge.fr),dictnode.get(edge.to),edge.c,edge.rate)

        return self.topo

    def CreatNodeEdgeSet(self,basic_node_edge_set,c=200,rate=100,option=0 ):#option=0默认为无向图 将点边集合增加边的属性
        newedge=[] #存储Link结构的边
        newnode=[] #存储节点

        newnode=basic_node_edge_set[0]

        for edge in basic_node_edge_set[1]:

            newedge.append(Link(edge[0],edge[1],c,rate))

        return [newnode,newedge]
    def Changelrate(self,topo,a=0,b=100):#随机生成链路的密钥速率
        for i in range(len(topo)):
           for j in range(i):
               if topo[i][j].Is_connected is True:
                   c=random.randint(a,b)
                   topo[i][j].rate=c
                   topo[j][i].rate=c

    def Changelc(self,topo,a=0,b=100):#随机生成链路的密钥量
        for i in range(len(topo)):
           for j in range(i):
               if topo[i][j].Is_connected is True:
                   n=random.randint(a,b)
                   topo[i][j].c=n
                   topo[j][i].c=n

    def TopoFilter(self,g,req,flag):#删去不符合需求的链路req=(s,d,keyvo,keyrate) flag:0表示筛选密钥量，1表示密钥速率，否则为智能筛选
        if flag==0:
            for i in range(len(g)):
                for j in g[i]:
                    if j.Is_connected==True and j.c<req[2] :
                        j.dellink()
        elif flag==1:
            for i in range(len(g)):
                for j in g[i]:
                    if j.Is_connected==True and j.rate<req[3]:
                        j.dellink()
        else:
             if req[2] is NULL: #说明只对密钥速率有需求
                for i in range(len(g)):
                    for j in g[i]:
                        if j.Is_connected==True and j.rate<req[3]:
                            j.dellink()
             elif req[3] is NULL: #说明只对密钥量有需求
                 pass
             else:  #说明对两者都有需求,密钥池能满足或者密钥池跟链路配和
                 for i in range(len(g)):
                    for j in g[i]:
                        if j.Is_connected==True:
                           if j.c<req[2] and j.rate+j.c*req[3]/req[2]<req[3]:
                                j.dellink()

    def creattr(self,a,n,sd):
        count=0
        for i in a:
            if i==1:
                count+=1
        while count<n:
            t=random.randint(0,38)
            while a[t]==1 or t in sd:
                t=random.randint(0,38)
            a[t]=1
            count+=1
        return a

    def TopoUpdater(self,g,path):#删去路径上的所有节点 对比算法。
        for i in range(len(path)-1):
            g[0][path[i]][path[i+1]].dellink()
            g[0][path[i+1]][path[i]].dellink()
        nodelist=path[1:-1]
        for i in range(len(g[0])):
            for j in range(len(g[0][i])):
                if i in  nodelist or j in nodelist:
                    if g[0][i][j].Is_connected==True:
                        g[0][i][j].dellink()
                        g[0][j][i].dellink()
    def Toporeduce(self,g):#将邻接矩阵化为邻接表
        tmptopo=[]
        for i in g:
            t=[]
            for j in range(len(i)):
                if i[j].fr != None:
                    t.append([i[j].to,1])
            tmptopo.append(t)
        return tmptopo

    def Toporeducehop(self,g):#将邻接矩阵化为hop邻接表
        tmptopo=[]
        for i in g[0]:
            t=[]
            for j in range(len(i)):
                if i[j].fr != None:
                    t.append([i[j].to,1])
            tmptopo.append(t)
        return [tmptopo,g[1]]

    def Changenodesp(self,topo,v):
        for i in range(len(topo[1])):
            if topo[1][i]!=1:
                topo[1][i]=v

    def create_random_topology(self,nodes_num=50, a=0.3, b=3,flag=0):##a = alpha, b = beta
        nodes = []
        edges = []
        positions = []
        edges_cost = []
        degree=[0]*nodes_num
        topology = [[0 for i in range(110)] for j in range(110)]
        distance = [[0 for i in range(nodes_num)] for j in range(nodes_num)]
        for n in range(0, nodes_num):
            x = random.randint(0, 100)
            y = random.randint(0, 100)
            while topology[x][y] == 1:
                x = random.randint(0, 100)
                y = random.randint(0, 100)
            topology[x][y] = 1
            nodes.append(f"N{n}")
            positions.append([x, y])
        i = 0
        max_length = -1
        while i < nodes_num:
            j = i + 1
            while j < nodes_num:
                distance[i][j] = (abs(positions[i][0] - positions[j][0]) ** 2 + abs(
                    positions[i][1] - positions[j][1])) ** 0.5
                max_length = max(max_length, distance[i][j])
                j += 1
            i += 1
        i = 0
        while i < nodes_num:
            j = i + 1
            while j < nodes_num:
                p =5* b * math.exp(-distance[i][j] / (max_length * a))/nodes_num  ##a = alpha, b = beta
                if random.random() < p:
                    edges_cost.append(round(distance[i][j], 1))
                    degree[i]+=1
                    degree[j]+=1
                    edges.append((nodes[i], nodes[j]))
                    edges.append((nodes[j], nodes[i]))
                j += 1
            i += 1
        for i in range(nodes_num):
            if degree[i]<2:
                while degree[i]<2:
                    j=random.randint(0,nodes_num-1)
                    while i==j:
                        j=random.randint(0,nodes_num-1)
                    p =5* b * math.exp(-distance[i][j] / (max_length * a))/nodes_num  ##a = alpha, b = beta
                    if random.random() < p:
                        edges_cost.append(round(distance[i][j], 1))
                        degree[i]+=1
                        degree[j]+=1
                        edges.append((nodes[i], nodes[j]))
                        edges.append((nodes[j], nodes[i]))

        if flag==1:
            g = nx.Graph()  # create graph
            g.add_nodes_from(nodes)  # add nodes
            g.add_edges_from(edges)  # add edges
            nodes_position = dict(zip(nodes, positions))  #
            node_labels = dict(zip(nodes, nodes))  # label of nodes
            edge_labels = dict(zip(edges, edges_cost))  # label of edges
            nx.draw_networkx_nodes(g, nodes_position, node_size=100, node_color="#6CB6FF")  # draw nodes
            nx.draw_networkx_edges(g, nodes_position, edges)  # draw edges
            nx.draw_networkx_labels(g, nodes_position, node_labels)  # draw label of nodes

            plt.axis('off')
            plt.show()
        return [nodes, edges]