#!/usr/bin/python

import igraph as ig
import plotly.graph_objs as go

from math import sin, cos, pi, sqrt

__all__ = ['draw_network', 'draw_experiment', 'get_default_test_network']

import rumba.elements.topology

type_to_marker = {
    'eth-dix': 'diamond',
    'eth-llc': 'diamond',
    'eth-udp': 'diamond',
    'unicast': 'circle',
    'broadcast': 'square'
}


def rumba_to_type(_type):
    if _type == rumba.elements.topology.ShimEthIPCP:
        return 'eth-dix'
    if _type == rumba.elements.topology.ShimUDPIPCP:
        return 'eth-udp'
    if _type == rumba.elements.topology.IPCP:
        return 'unicast'


def get_default_test_network():
    return {
        'nodes': list(range(8)),
        'layers': {
            'bottom_layer1': {
                'type': 'eth-dix',
                'nodes': list(range(5)),
                'edges': [(0, 1), (0, 2), (0, 3), (0, 4), (1, 3), (2, 3), (1, 4), (2, 4)]
            },
            'bottom_layer2': {
                'type': 'eth-llc',
                'nodes': list(range(4, 8)),
                'edges': [(4, 5), (5, 6), (5, 7), (6, 7), (4, 7)]
            },
            'medium_layer': {
                'type': 'unicast',
                'nodes': [0, 2, 4, 5, 7],
                'edges': [(0, 2), (0, 4), (4, 5), (5, 7), (2, 4), (4, 7)]
            },
            'top_layer': {
                'type': 'unicast',
                'nodes': [0, 3, 6, 7],
                'edges': [(0, 7), (0, 3), (6, 7)]
            }
        },
        'registrations': {
            'medium_layer': {
                'bottom_layer1': [0, 2, 4],
                'bottom_layer2': [4, 5, 7]
            },
            'top_layer': {
                'medium_layer': [0, 7],
                'bottom_layer1': [0, 3],
                'bottom_layer2': [6, 7]
            }
        }
    }


def _get_nodes_in_dif(exp, dif):
    nodes = []
    n = 0
    for node in exp.nodes:
        if dif in [d.name for d in node.difs]:
            nodes.append(n)
        n += 1

    return nodes


def _get_node_index(exp, node):
    n = 0
    for _node in exp.nodes:
        if _node.name == node.name:
            return n
        n += 1
    return 0


def _get_network_from_rumba_experiment(exp):
    print(exp)
    print(exp.flows)
    _nw = dict()
    _nw['layers'] = dict()

    _nw['nodes'] = list(range(len(exp.nodes)))

    _nw['registrations'] = dict()

    for node in exp.nodes:
        for dif in node.difs:
            if dif.name not in _nw['layers']:
                _nw['layers'][dif.name] = dict()
                _nw['layers'][dif.name]['type'] = rumba_to_type(dif.get_ipcp_class())
                _nw['layers'][dif.name]['nodes'] = _get_nodes_in_dif(exp, dif.name)
                _nw['layers'][dif.name]['edges'] = list()
                if _nw['layers'][dif.name]['type'] != 'unicast':  # shim
                    nodes = _nw['layers'][dif.name]['nodes']
                    _nw['layers'][dif.name]['edges'].append((nodes[0], nodes[1]))
                _nw['registrations'][dif.name] = dict()

    for layer in exp.flows:
        for flow in layer:
            if 'src' in flow and 'dst' in flow:
                src = _get_node_index(exp, flow['src'].node)
                dst = _get_node_index(exp, flow['dst'].node)
                layer = flow['src'].dif.name
                _nw['layers'][layer]['edges'].append((src, dst))
                src_regs = flow['src'].registrations
                dst_regs = flow['dst'].registrations
                for dif in src_regs:
                    if dif.name not in _nw['registrations'][layer]:
                        _nw['registrations'][layer][dif.name] = set()
                    _nw['registrations'][layer][dif.name].add(src)
                for dif in dst_regs:
                    if dif.name not in _nw['registrations'][layer]:
                        _nw['registrations'][layer][dif.name] = set()
                    _nw['registrations'][layer][dif.name].add(dst)

    return _nw


def _get_rank(network, layer):
    rank = 0

    if layer not in network['registrations']:
        return rank

    for _layer in network['registrations'][layer]:
        n_1_rank = _get_rank(network, _layer)
        if n_1_rank >= rank:
            rank = n_1_rank + 1

    return rank


def _get_ipcp_id(network, layer, node):
    for ipcp in network['_ipcps']:
        if ipcp['node'] == node and ipcp['layer'] == layer:
            return ipcp['id']


def _create_ipcp_network(network):
    network['_ipcps'] = list()
    network['_adjs'] = list()
    network['_deps'] = list()
    network['_systems'] = list()

    for node in network['nodes']:
        network['_systems'].append({'node': node})

    ipcp_id = 0
    color = 0

    for layer in network['layers']:
        for node in network['layers'][layer]['nodes']:
            ipcp = {
                'id': ipcp_id,
                'node': node,
                'layer': layer,
                'color': color,
                'type': network['layers'][layer]['type'],
                'rank': _get_rank(network, layer)
            }
            network['_ipcps'].append(ipcp)
            ipcp_id += 1
        for edge in network['layers'][layer]['edges']:
            src = _get_ipcp_id(network, layer, edge[0])
            dst = _get_ipcp_id(network, layer, edge[1])
            adj = {
                'src': src,
                'dst': dst,
                'color': color
            }
            network['_adjs'].append(adj)
        color += 1

    for upper_layer in network['registrations']:
        for lower_layer in network['registrations'][upper_layer]:
            for node in network['registrations'][upper_layer][lower_layer]:
                src = _get_ipcp_id(network, upper_layer, node)
                dst = _get_ipcp_id(network, lower_layer, node)
                network['_deps'].append((src, dst))


def _create_system_graph(network):
    edges = []

    for layer in network['layers']:
        for edge in network['layers'][layer]['edges']:
            if edge not in edges:
                edges.append(edge)

    return edges


def _create_system_coords(network):
    """
    Create coordinates for the systems in the network
    :param network:
    :return:
    """
    edges = _create_system_graph(network)

    g = ig.Graph(edges, directed=False)
    layout = g.layout('fr', dim=2)

    for system in network['_systems']:
        n = system['node']
        system['coords'] = (layout[n][0], layout[n][1])

    _min = None

    for edge in edges:
        x1, y1 = network['_systems'][edge[0]]['coords']
        x2, y2 = network['_systems'][edge[1]]['coords']
        d = sqrt((x2 - x1) * (x2 - x1) + (y2 - y1) * (y2 - y1))
        _min = d if _min is None else min(d, _min)

    return _min


def _get_ipcps_for_system(network, system):
    node = system['node']
    return [ipcp for ipcp in network['_ipcps'] if ipcp['node'] == node]


def _get_ranks_for_ipcps(ipcps):
    return set([ipcp['rank'] for ipcp in ipcps])


def _get_ipcps_by_rank(ipcps, rank):
    return [ipcp for ipcp in ipcps if ipcp['rank'] == rank]


def _get_ipcps_by_layer(network, layer):
    return [ipcp for ipcp in network['_ipcps'] if ipcp['layer'] == layer]


def _get_rank_for_layer(network, layer):
    for ipcp in network['_ipcps']:
        if ipcp['layer'] == layer:
            return ipcp['rank']

    return 0


def _get_layers_by_rank(network, rank):
    return [layer for layer in network['layers']
            if _get_rank_for_layer(network, layer) == rank]


def _create_ipcp_coords(network, radius):
    """
    Create 3d coordinates for IPCPs based on the system layout
    """
    height = 1.0
    max_rank = 0

    for system in network['_systems']:
        ipcps = _get_ipcps_for_system(network, system)
        ranks = _get_ranks_for_ipcps(ipcps)
        n = len(ipcps)
        m = 1
        for ipcp in ipcps:
            rank = ipcp['rank']
            x = system['coords'][0]
            y = system['coords'][1]
            new_x = x + radius * sin((m/n) * pi)
            new_y = y + radius * cos((m/n) * pi)
            ipcp['coords'] = (new_x, new_y, rank * height)
            m += 1
            max_rank = max(max_rank, rank)

    for rank in range(max_rank + 1):
        layers = _get_layers_by_rank(network, rank)
        increment = height / (3 / 2 * len(layers))
        offset = 0
        for layer in layers:
            ipcps = _get_ipcps_by_layer(network, layer)
            for ipcp in ipcps:
                (x, y, z) = ipcp['coords']
                ipcp['coords'] = (x, y, z + offset)
            offset += increment


def _create_ipcp_graph_data(network):
    _create_ipcp_network(network)
    _min = _create_system_coords(network)

    print("_min = %s" % _min)
    _create_ipcp_coords(network, _min / 5)

    print(network)


def _get_ipcp_attributes(network):
    coords = list()
    colors = list()
    labels = list()
    markers = list()

    for ipcp in network['_ipcps']:
        coords.append(ipcp['coords'])
        colors.append(ipcp['color'])
        labels.append(ipcp['layer'] + ' ' + str(ipcp['node']))
        markers.append(type_to_marker[ipcp['type']])

    return coords, colors, labels, markers


def _get_edge_attributes(network):
    x_coords = list()
    y_coords = list()
    z_coords = list()
    colors = list()

    for adj in network['_adjs']:
        src = network['_ipcps'][adj['src']]['coords']
        dst = network['_ipcps'][adj['dst']]['coords']
        x_coords.extend([src[0], dst[0], None])
        y_coords.extend([src[1], dst[1], None])
        z_coords.extend([src[2], dst[2], None])
        colors.append(adj['color'])

    return x_coords, y_coords, z_coords, colors


def _get_deps_attributes(network):
    x_coords = list()
    y_coords = list()
    z_coords = list()
    colors = list()

    for dep in network['_deps']:
        src = network['_ipcps'][dep[0]]['coords']
        dst = network['_ipcps'][dep[1]]['coords']
        x_coords.extend([src[0], dst[0], None])
        y_coords.extend([src[1], dst[1], None])
        z_coords.extend([src[2], dst[2], None])
        colors.append(0)

    return x_coords, y_coords, z_coords, colors


def _extract(coords):
    x = []
    y = []
    z = []

    for coord in coords:
        x.append(coord[0])
        y.append(coord[1])
        z.append(coord[2])

    return x, y, z


def draw_network(network, name='Ouroboros network'):
    _create_ipcp_graph_data(network)
    coords, colors, labels, markers = _get_ipcp_attributes(network)

    x_coords, y_coords, z_coords = _extract(coords)

    nodes = go.Scatter3d(
        x=x_coords,
        y=y_coords,
        z=z_coords,
        mode='markers+text',
        name='actors',
        marker=dict(symbol=markers,
                    size=6,
                    color=colors,
                    colorscale='Viridis',
                    line=dict(color='rgb(50,50,50)', width=0.5)
                    ),
        text=labels,
        hoverinfo='none'
    )

    x_coords, y_coords, z_coords, colors = _get_edge_attributes(network)

    layers = go.Scatter3d(
        x=x_coords,
        y=y_coords,
        z=z_coords,
        mode='lines',
        line=dict(color=colors,
                  colorscale='Viridis',
                  dash='solid',
                  width=4),
        hoverinfo='none'
    )

    x_coords, y_coords, z_coords, colors = _get_deps_attributes(network)

    deps = go.Scatter3d(
        x=x_coords,
        y=y_coords,
        z=z_coords,
        mode='lines',
        line=dict(color=colors,
                  colorscale='Viridis',
                  dash='solid',
                  width=1),
        hoverinfo='none'
    )

    axis = dict(showbackground=False,
                showline=False,
                zeroline=False,
                showgrid=False,
                showticklabels=False,
                title='')

    layout = go.Layout(
        title=name,
        autosize=True,
        width=3600,
        height=1800,
        showlegend=False,
        scene=dict(
            xaxis=dict(axis),
            yaxis=dict(axis),
            zaxis=dict(axis),
        ),
        margin=dict(
            l=0,
            r=0,
            b=0,
            t=100
        ),
        hovermode='closest',
        annotations=[
            dict(
                showarrow=False,
                text="Data source: Ouroboros",
                xref='paper',
                yref='paper',
                xanchor='left',
                yanchor='bottom',
                x=0,
                y=0.1,
                font=dict(
                    size=14
                )
            )
        ],
    )

    data = [nodes, layers, deps]
    fig = go.Figure(data=data, layout=layout)
    fig.update_scenes(aspectmode='data')

    fig.show()


def draw_experiment(exp, name='experiment'):
    _nw = _get_network_from_rumba_experiment(exp)
    draw_network(_nw, name)


if __name__ == '__main__':
    nw = get_default_test_network()
    draw_network(nw)