from __future__ import print_function
from io import open
import pandas
import igraph # Install Igraph with pip install python-igraph
import graphistry

graphistry.register('Email pygraphistry@graphistry.com to get a key.')

"""
Explanation: PyGraphistry Example: Graphing the Marvel Universe
Plots hero social network based on co-appearences between heroes
Install: pip install "graphistry[igraph]"
Note: pip install igraph is the wrong package. if installing manually, use python-igraph

Uses pandas, igraph, and PyGraphistry
Combines comic book and hero data
Near end, computes clusters and to avoid a hairball, weakens the edge weights between  nodes of different clusters
End of explanation
"""


with open('data/characters.txt', encoding="latin-1") as f:
    lines = f.readlines()
heroes = pandas.DataFrame(
    list(map(lambda x: (int(x.split(':')[0].split(' ')[1]), x.split(': ', 1)[1].split('\n')[0]), lines)),
    columns=['hero_id', 'hero_name'])
print('#Heroes:', len(heroes))
heroes[:3]

with open('data/comics.txt', encoding="latin-1") as f:
    lines = f.readlines()
comics = pandas.DataFrame(
    list(map(lambda x: (int(x.split(':')[0].split(' ')[1]), x.split(': ', 1)[1].split('\n')[0]), lines)),
    columns=['comic_id', 'comic_name'])
print('#Comics: ', len(comics))
comics[:3]

with open('data/appearances.txt', encoding="latin-1") as f:
    lines = f.readlines()[len(heroes) + len(comics) + 2:]
def expand (line):
    parts = list(map(int, line.split(' ')))
    return [(parts[0], role) for role in parts[1:]]
appearences = pandas.DataFrame(
    [item for sublist in list(map(expand, lines)) for item in sublist],
    columns=['hero', 'comic'])
appearences[:3]

"""
Explanation: Load heroes, comics, appearences
<font style="color:red">Data should be in data/&lt;file&gt;.txt</font>
End of explanation
"""


# You may need to install numexpr: pip install numexpr
coappearences = \
    appearences\
        .merge(appearences, on='comic')\
        .merge(comics, left_on='comic', right_on='comic_id')\
        [['hero_x', 'hero_y']]\
        .query('hero_x > hero_y')
unique_coappearences = coappearences.drop_duplicates(['hero_x', 'hero_y']).set_index(['hero_x', 'hero_y'])
unique_coappearences['counts'] = coappearences.groupby(['hero_x', 'hero_y']).size()
unique_coappearences = unique_coappearences.reset_index()
print('#edges', len(unique_coappearences))        
unique_coappearences[:5]

"""
Explanation: Link heroes who co-appear
End of explanation
"""


plotter = graphistry.bind(source='hero_x', destination='hero_y', edge_title='counts')

plotter.plot(unique_coappearences)

"""
Explanation: Plot!
End of explanation
"""


# Here we are using two dataframes, one for edges and one for nodes
plotter2 = plotter.bind(node='hero_id', point_title='hero_name')

plotter2.plot(unique_coappearences, heroes)

"""
Explanation: Label Nodes
End of explanation
"""


ig = plotter2.pandas2igraph(unique_coappearences, directed=False)
clusters = ig.community_infomap()
(i_edges, i_nodes) = plotter2.igraph2pandas(ig)
print('#clusters', str(len(list(set(clusters.membership)))))

nodes_colored = pandas.DataFrame({'cluster': clusters.membership})\
    .reset_index().rename(columns={'index': 'denseid'})\
    .merge(i_nodes.reset_index().rename(columns={'index':'denseid'}), on='denseid')\
    .merge(heroes, left_on='hero_id', right_on='hero_id')
print('#colored nodes', str(len(nodes_colored)))
nodes_colored[:3]

nodes_colored['color'] = nodes_colored.apply(lambda x: x['cluster'] % 9, axis=1)
nodes_colored.pivot_table(index=['color'], aggfunc=lambda x: len(x.unique()))

plotter3 = plotter2.bind(point_color='color', edge_weight='counts')

plotter3.plot(unique_coappearences,  nodes_colored)

"""
Explanation: Color using igraph infomap
Infomap Community Detection
End of explanation
"""


big_clusters = nodes_colored\
    .pivot_table(index=['cluster'], aggfunc=lambda x: len(x.unique()))\
    .rename(columns={'hero_id': 'cluster_size'})\
    .query('cluster_size > 100')\
    .reset_index()[['cluster', 'cluster_size']]
print('# big clusters', len(big_clusters))
big_clusters[:3]

good_nodes = nodes_colored.merge(big_clusters, on='cluster')
print('# nodes', len(good_nodes))
good_nodes[:3]

good_edges = unique_coappearences\
    .merge(good_nodes, left_on='hero_x', right_on='hero_id')\
    .merge(good_nodes, left_on='hero_y', right_on='hero_id')\
    [['hero_x', 'hero_y', 'counts']]
print('# edges', len(good_edges))
good_edges[:3]

plotter3.plot(good_edges, good_nodes)

"""
Explanation: Restrict to biggest communities
End of explanation
"""


#label edges whether they stay inside a cluster or connect nodes in different clusters
good_edges2 = good_edges\
        .merge(\
                 good_nodes[['cluster', 'hero_id']].rename(columns={'cluster': 'cluster_x'}),\
                 left_on='hero_x', right_on='hero_id')\
        .merge(\
                 good_nodes[['cluster', 'hero_id']].rename(columns={'cluster': 'cluster_y'}),\
                 left_on='hero_y', right_on='hero_id')
good_edges2['is_inner'] = good_edges2.apply(lambda x: x['cluster_x'] == x['cluster_y'], axis=1)

#bind to edge_weight
good_edges2['weight'] = good_edges2.apply(lambda x: 10 if x['is_inner'] else 8, axis=1)
good_edges2 = good_edges2[['hero_x', 'hero_y', 'counts', 'is_inner', 'weight']]
good_edges2[:3]

"""
Explanation: Seperate communities
Treat intra-community edges as strong edge weights, and inter-community as weak edge weight
End of explanation
"""


plotter3.bind(edge_weight='weight').plot(good_edges2, good_nodes)

"""
Explanation: Plot; control the edge weight in the settings panel
End of explanation
"""


