import random
from agents import User, Country
from network import Network

import networkx as nx
from forceatlas2 import forceatlas2_networkx_layout

import matplotlib.pyplot as plt

class Context:
    def __init__(self, param):

        self.network = Network()
        self.param = param
        self.edge_colors = {}
        self.edge_countdowns = {}
        self.node_colors = {}
        self.node_countdowns = {}
        self.absolute_balkanisation = []
        self.relative_balkanisation = []

        servers = []

        for i in range(self.param.num_users):
            u = User('U' + str(i))
            u.context = self
            u.network = self.network
            u.prob_block_threat = self.param.user_prob_block_threat
            servers.append(u)

        for i in range(self.param.num_countries):
            c = Country('C' + str(i))
            c.context = self
            c.network = self.network
            c.prob_block_threat = self.param.country_prob_block_threat
            servers.append(c)

        self.network.initialise(servers, self.param.edges_per_node)
        self.pos = self.layout(None, 10)

    def packet_made(self, node, packet):
        """ Display a threat packet's path in red and the sender node in purple """
        if packet.threat:
            for a, b in zip(packet.path, packet.path[1:]):
                self.edge_colors[(a,b)] = [1,0,0]
                self.edge_countdowns[(a,b)] = self.param.skip_frames

            self.node_colors[node] = [1,0,1]
            self.node_countdowns[node] = self.param.skip_frames

    def node_blocked_threat(self, node, packet):
        """ Display a blocked packet's intended path and receiver in green, and the blocker in cyan """
        if packet.threat:
            for a, b in zip([node] + packet.path, packet.path):
                self.edge_colors[(a,b)] = [0,1,0]
                self.edge_countdowns[(a,b)] = self.param.skip_frames

            self.node_colors[packet.receiver] = [0,1,0]
            self.node_countdowns[packet.receiver] = self.param.skip_frames

            self.node_colors[node] = [0,1,1]
            self.node_countdowns[node] = self.param.skip_frames

    def node_received_threat(self, node, packet):
        """ Display a threat packet's target in red """
        self.node_colors[node] = [1,0,0]
        self.node_countdowns[node] = self.param.skip_frames

    def run_one_cycle(self, calc_centralities=True, calc_balkanisation=True):
        """ Centralities are used for display to size nodes, balkanisation is used for the final plot """
        random.choice(list(self.network.g.nodes)).run_one_cycle()
        if calc_centralities:
            self.network.update_centralities()
        if calc_balkanisation:
            self.absolute_balkanisation.append(self.network.absolute_balkanisation())
            self.relative_balkanisation.append(self.network.relative_balkanisation([n for n in self.network.g if isinstance(n,Country)]))

    def layout(self, pos, iterations):
        return forceatlas2_networkx_layout(self.network.g, pos=pos, niter=iterations, scalingRatio=0.5)

    def get_colors(self):
        return list(map(lambda x: self.get_node_color(x), self.network.g.nodes()))

    def get_sizes(self):
        return list(map(lambda x: ((self.network.centralities[x]**2)*200 + 10), self.network.g.nodes()))

    def get_edge_colors(self):
        return list(map(lambda x: self.get_edge_color(x), self.network.g.edges()))

    def update(self, i, draw=True, plot=False):
        if i % self.param.skip_frames == 0:
            print(i // self.param.skip_frames)
            self.run_one_cycle(calc_centralities=draw, calc_balkanisation=plot)
        if draw:
            self.pos = self.layout(self.pos, 1)
            nx.draw(self.network.g, pos=self.pos, node_color=self.get_colors(), node_size=self.get_sizes(), edge_color=self.get_edge_colors())

    def plot(self):
        plt.plot(self.absolute_balkanisation, label='Absolute')
        plt.plot(self.relative_balkanisation, label='Relative')
        plt.ylim(0,1)
        plt.xlabel('Time')
        plt.ylabel('Balkanisation')
        plt.legend()


    def get_node_color(self, n):
        """ By default, users are black and countries are blue. Colours can change temporarily in response to packets. """
        if n in self.node_colors:
            color = self.node_colors[n]
            self.node_countdowns[n] -= 1
            if self.node_countdowns[n] <= 0:
                del self.node_colors[n]
            return color
        else:
            return [0,0,0] if isinstance(n, User) else [0,0.5,1]

    def get_edge_color(self, e):
        """ By default, edges are black. Colours can change temporarily in response to packets. """
        a,b = e
        if (a,b) in self.edge_colors:
            color = self.edge_colors[(a,b)]
            self.edge_countdowns[(a,b)] -= 1
            if self.edge_countdowns[(a,b)] <= 0:
                del self.edge_colors[(a,b)]
            return color
        elif (b,a) in self.edge_colors:
            color = self.edge_colors[(b,a)]
            self.edge_countdowns[(b,a)] -= 1
            if self.edge_countdowns[(b,a)] <= 0:
                del self.edge_colors[(b,a)]
            return color
        else:
            return [0,0,0]

import random
from agents import User, Country
from network import Network

import networkx as nx
from forceatlas2 import forceatlas2_networkx_layout

import matplotlib.pyplot as plt

class Context:
    def __init__(self, param):

        self.network = Network()
        self.param = param
        self.edge_colors = {}
        self.edge_countdowns = {}
        self.node_colors = {}
        self.node_countdowns = {}
        self.absolute_balkanisation = []
        self.relative_balkanisation = []

        servers = []

        for i in range(self.param.num_users):
            u = User('U' + str(i))
            u.context = self
            u.network = self.network
            u.prob_block_threat = self.param.user_prob_block_threat
            servers.append(u)

        for i in range(self.param.num_countries):
            c = Country('C' + str(i))
            c.context = self
            c.network = self.network
            c.prob_block_threat = self.param.country_prob_block_threat
            servers.append(c)

        self.network.initialise(servers, self.param.edges_per_node)
        self.pos = self.layout(None, 10)

    def packet_made(self, node, packet):
        """ Display a threat packet's path in red and the sender node in purple """
        if packet.threat:
            for a, b in zip(packet.path, packet.path[1:]):
                self.edge_colors[(a,b)] = [1,0,0]
                self.edge_countdowns[(a,b)] = self.param.skip_frames

            self.node_colors[node] = [1,0,1]
            self.node_countdowns[node] = self.param.skip_frames

    def node_blocked_threat(self, node, packet):
        """ Display a blocked packet's intended path and receiver in green, and the blocker in cyan """
        if packet.threat:
            for a, b in zip([node] + packet.path, packet.path):
                self.edge_colors[(a,b)] = [0,1,0]
                self.edge_countdowns[(a,b)] = self.param.skip_frames

            self.node_colors[packet.receiver] = [0,1,0]
            self.node_countdowns[packet.receiver] = self.param.skip_frames

            self.node_colors[node] = [0,1,1]
            self.node_countdowns[node] = self.param.skip_frames

    def node_received_threat(self, node, packet):
        """ Display a threat packet's target in red """
        self.node_colors[node] = [1,0,0]
        self.node_countdowns[node] = self.param.skip_frames

    def run_one_cycle(self, calc_centralities=True, calc_balkanisation=True):
        """ Centralities are used for display to size nodes, balkanisation is used for the final plot """
        random.choice(list(self.network.g.nodes)).run_one_cycle()
        if calc_centralities:
            self.network.update_centralities()
        if calc_balkanisation:
            self.absolute_balkanisation.append(self.network.absolute_balkanisation())
            self.relative_balkanisation.append(self.network.relative_balkanisation([n for n in self.network.g if isinstance(n,Country)]))

    def layout(self, pos, iterations):
        return forceatlas2_networkx_layout(self.network.g, pos=pos, niter=iterations, scalingRatio=0.5)

    def get_colors(self):
        return list(map(lambda x: self.get_node_color(x), self.network.g.nodes()))

    def get_sizes(self):
        return list(map(lambda x: ((self.network.centralities[x]**2)*200 + 10), self.network.g.nodes()))

    def get_edge_colors(self):
        return list(map(lambda x: self.get_edge_color(x), self.network.g.edges()))

    def update(self, i, draw=True, plot=False):
        if i % self.param.skip_frames == 0:
            print(i // self.param.skip_frames)
            self.run_one_cycle(calc_centralities=draw, calc_balkanisation=plot)
        if draw:
            self.pos = self.layout(self.pos, 1)
            nx.draw(self.network.g, pos=self.pos, node_color=self.get_colors(), node_size=self.get_sizes(), edge_color=self.get_edge_colors())

    def plot(self):
        plt.plot(self.absolute_balkanisation, label='Absolute')
        plt.plot(self.relative_balkanisation, label='Relative')
        plt.ylim(0,1)
        plt.xlabel('Time')
        plt.ylabel('Balkanisation')
        plt.legend()


    def get_node_color(self, n):
        """ By default, users are black and countries are blue. Colours can change temporarily in response to packets. """
        if n in self.node_colors:
            color = self.node_colors[n]
            self.node_countdowns[n] -= 1
            if self.node_countdowns[n] <= 0:
                del self.node_colors[n]
            return color
        else:
            return [0,0,0] if isinstance(n, User) else [0,0.5,1]

    def get_edge_color(self, e):
        """ By default, edges are black. Colours can change temporarily in response to packets. """
        a,b = e
        if (a,b) in self.edge_colors:
            color = self.edge_colors[(a,b)]
            self.edge_countdowns[(a,b)] -= 1
            if self.edge_countdowns[(a,b)] <= 0:
                del self.edge_colors[(a,b)]
            return color
        elif (b,a) in self.edge_colors:
            color = self.edge_colors[(b,a)]
            self.edge_countdowns[(b,a)] -= 1
            if self.edge_countdowns[(b,a)] <= 0:
                del self.edge_colors[(b,a)]
            return color
        else:
            return [0,0,0]

Display a threat packet's path in red and the sender node in purple Display a blocked packet's intended path and receiver in green, and the blocker in cyan Display a threat packet's target in red Centralities are used for display to size nodes, balkanisation is used for the final plot By default, users are black and countries are blue. Colours can change temporarily in response to packets. By default, edges are black. Colours can change temporarily in response to packets.