import numpy as np
import h5py
import seaborn as sns
import os
import matplotlib.pyplot as plt
import matplotlib.colors as clr
import networkx as nx
import statsmodels.api as sm

from mpl_toolkits.axes_grid1.inset_locator import inset_axes

import params
import utils

CULTURE_LABELS_TOP_GAP = 0.225
CULTURE_LABELS_GAP = 0.19

# OUTPUT_FILE_NAMES = [
#     "layout_net_degrees",
#     "layout_net_TEs",
#     "layout_net_centrality"
# ]

options = {
    #'node_color': 'black',
    #'node_size': 2,
    #'edge_color': 'black',
    #'width': 3.0,
    #'alpha': 0.5,
    'connectionstyle': "arc3,rad=0.1",
    'arrowstyle': '-|>',
    'arrowsize': 4,
    #'edge_cmap': plt.cm.Blues,
}


def make_big_legend(axs, position):
    legend_graph = nx.grid_2d_graph(4, 4)
    legend_graph.remove_edges_from(legend_graph.edges())
    legend_graph = legend_graph.to_directed()
    pos = {(x,y):(x,-y) for x,y in legend_graph.nodes()}
    legend_graph.add_edges_from([((3, 0), (2, 0)), ((3, 1), (2, 1)), ((3, 2), (2, 2)), ((3, 3), (2, 3))])
    node_sizes = [50, 50, 50, 50, 120, 80, 40, 10, 0, 0, 0, 0, 0, 0, 0, 0]
    node_colours = [1.0, 0.66, 0.33, 0.1, 0.6, 0.6, 0.6, 0.6, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5]
    edge_widths = [3, 2, 1.5, 1]
    edge_colors = [1.0, 0.8, 0.6, 0.3]

    nx.draw(legend_graph, ax = axs[position[0], position[1]], pos = pos,
        width = edge_widths,
        edge_color = edge_colors, edge_cmap = plt.cm.GnBu, #arrowsize = arrow_sizes,
        edge_vmin = 0, edge_vmax = 1,
        node_size = node_sizes,
        node_color = node_colours, cmap = plt.cm.YlOrRd, vmin = 0,
        vmax = 1,
        **options)

    axs[position[0], position[1]].text(0.15, -0.4, "high\ninward\nTE", fontsize = 8)
    axs[position[0], position[1]].text(0.15, -3.0, "low\ninward\nTE", fontsize = 8)
    axs[position[0], position[1]].text(1.15, -0.4, "high\noutward\nTE", fontsize = 8)
    axs[position[0], position[1]].text(1.15, -3.0, "low\noutward\nTE", fontsize = 8)
    axs[position[0], position[1]].text(3.15, -0.3, "high\nTE", fontsize = 8)
    axs[position[0], position[1]].text(3.15, -3.0, "low\nTE", fontsize = 8)


def make_little_legend(axs):
    legend_graph = nx.grid_2d_graph(4, 4)
    legend_graph.remove_edges_from(legend_graph.edges())
    legend_graph = legend_graph.to_directed()
    pos = {(x,y):(x,-y) for x,y in legend_graph.nodes()}
    node_sizes = [0, 0, 0, 0, 120, 80, 40, 10, 0, 0, 0, 0, 0, 0, 0, 0]
    node_colours = [0.5, 0.5, 0.5, 0.5, 1.0, 0.66, 0.33, 0.1, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5]
    edge_widths = [3, 2, 1.5, 1]
    nx.draw(legend_graph, ax = axs[0, 3], pos = pos,
        node_size = node_sizes,
        node_color = node_colours, cmap = plt.cm.YlOrRd, vmin = 0,
        vmax = 1,
        **options)
    axs[0, 3].text(1.15, -0.4, "high\ndegree", fontsize = 8)
    axs[0, 3].text(1.15, -3.0, "low\ndegree", fontsize = 8)


def make_empty_network(G, ax, run_index, day_index):
    nx.draw(G, nx.get_node_attributes(G, 'pos'), ax = ax,
            edgelist = [], node_size = 2,
            node_color = list(0.3 * np.ones(59)), cmap = plt.cm.YlOrRd, vmin = 0,
            vmax = 1, **options)
    offset = 0.01
    bbox = ax.get_position()
    ax.set_position([bbox.x0 + offset, bbox.y0 + offset, bbox.x1-bbox.x0 - 2 * offset, bbox.y1 - bbox.y0 - 2 * offset])
    ax.set_title("day " + str(params.DAYS[run_index][day_index]), y = -0.2)


def make_network(G, non_zero_vals, ax, run_index, day_index):
# Edge colours
    edge_color_max = np.mean(non_zero_vals) + 5 * np.std(non_zero_vals)
    edge_color_min = np.mean(non_zero_vals) - 2 * np.std(non_zero_vals)
    edges, edge_TEs = zip(*nx.get_edge_attributes(G, 'TE').items())

    edge_TEs, edges = zip(*sorted(zip(edge_TEs, edges)))
    edge_colors = edge_TEs
    # Edge widths
    edge_widths = 1000 * edge_TEs
    if np.count_nonzero(edge_TEs) > 1:
        edge_widths = 0.01 + 0.2 * (np.maximum(0, np.array(edge_TEs) - np.mean(edge_TEs))/np.std(edge_TEs))
    else:
        edge_widths = 0.01 + np.array(edge_TEs)
    arrow_sizes = 2 * edge_widths
    node_sizes = np.array([degree for (node, degree) in G.out_degree()])
    node_colours = np.array([degree for (node, degree) in G.in_degree()])

    node_sizes = []
    node_colours = []
    for node in G.nodes():
        if len(G.in_edges(node)) > 0:
            inward_TEs = [data['TE'] for source, target, data in G.in_edges(node, data = True)]
            node_colours.append(sum(inward_TEs))
        else:
            node_colours.append(0)
    for node in G.nodes():
        if len(G.out_edges(node)) > 0:
            outward_TEs = [data['TE'] for source, target, data in G.out_edges(node, data = True)]
            node_sizes.append(sum(outward_TEs))
        else:
            node_sizes.append(0)

    #if np.count_nonzero(node_sizes) > 1:
    node_sizes = 6 + 33 * np.maximum(0, node_sizes - np.mean(node_sizes))/np.std(node_sizes)
    node_sizes = list(node_sizes)
    node_colours = list(node_colours)
    node_colour_max = max(node_colours)
    node_colour_min = min(node_colours) - np.std(node_colours)


    nx.draw(G, nx.get_node_attributes(G, 'pos'), ax = ax,
        edgelist = edges, width = edge_widths,
        edge_color = edge_colors, edge_cmap = plt.cm.GnBu, #arrowsize = arrow_sizes,
        edge_vmin = edge_color_min, edge_vmax = edge_color_max,
        node_size = node_sizes,
        node_color = node_colours, cmap = plt.cm.YlOrRd, vmin = node_colour_min,
        vmax = node_colour_max,
        **options)

    ax.set_title("day " + str(params.DAYS[run_index][day_index]), y = -0.1)


def construct_spatial_nets(list_of_networks, sub_base_folder, run_group_index):

    plt.rc('axes', labelsize = params.LABEL_SIZES[run_group_index])
    plt.rc('font', size = params.LABEL_SIZES[run_group_index])
    plt.rc('xtick', labelsize = params.TICK_SIZES[run_group_index])
    plt.rc('ytick', labelsize = params.TICK_SIZES[run_group_index])
    plt.rc('axes', titlesize = params.TITLE_SIZES[run_group_index])

    fig, axs = utils.make_basic(True, run_group_index)

    #if plot_type in [1, 2]:
    make_big_legend(axs, params.NETWORK_LEGEND_POSITIONS[run_group_index])

    #elif plot_type in [3]:
    #    make_little_legend(axs, plot_type)


    #offset = 0.02
    #bbox = axs[0, 3].get_position()
    #axs[0, 3].set_position([bbox.x0 + offset, bbox.y0 + offset, bbox.x1-bbox.x0 - 2 * offset, bbox.y1 - bbox.y0 - 2 * offset])

    for run_index_in_group in range(len(params.RUN_GROUPS[run_group_index])):
        run_index = params.RUN_GROUPS[run_group_index][run_index_in_group]
        for day_index in range(params.MAX_NUM_DAYS_BY_GROUP[run_group_index]):

            print(run_group_index, run_index, day_index)

            if day_index >= (len(params.DAYS[run_index])):
                axs[run_index_in_group, day_index].axis('off')

            else:
                G = list_of_networks[run_index][day_index]

                if len(G.edges) > 0:
                    edges, non_zero_vals = zip(*nx.get_edge_attributes(G, 'TE').items())
                    make_network(G, non_zero_vals, axs[run_index_in_group, day_index], run_index, day_index)
                else:
                    make_empty_network(G, axs[run_index_in_group, day_index], run_index, day_index)

    plt.savefig(params.BASE_FIGURE_FOLDER  + sub_base_folder + params.NETS_AND_MAPS_FIGURE_FOLDER
                + params.RUN_GROUP_PREFIXES[run_group_index] + "layout_net_TEs.pdf")