%matplotlib inline

import glob
import matplotlib.pyplot as plt
import seaborn as sns
import networkx as nx
from nxviz import MatrixPlot, CircosPlot, ArcPlot
import pandas as pd
import tqdm

import sys
sys.path.append('../')
import termite as trmt

"""
Explanation: Loading Libraries
End of explanation
"""


nest = trmt.Experiment('/media/dmrib/tdata/Syntermes/N11HHS2018-8W4S/N11HHS2018-8W4S-1/Expanded/')

"""
Explanation: Load Dataset
The Experiment class in the termite module provides tools for preprocessing a tracking result dataset, we'll assume that the data was generated using this process from now on.
Let's proceed to load the data then...
End of explanation
"""


def build_frame_encounters_graph(nest, frame_number):
    G = nx.Graph()
    color_map = []
    for termite in nest.termites:
        G.add_node(termite.label)
        if termite.caste == 'S':
            color_map.append('gray')
        else:
            color_map.append('black')
    for termite in nest.termites:
        for other in nest.termites:
            if termite != other:
                if termite.trail.loc[frame_number, f'encountering_{other.label}']:
                    G.add_edge(termite.label, other.label)
                    
    return G, color_map 

"""
Explanation: This gives us an object containing both the nests' termites trails and their metadata so we can further investigate the underlying structure of termite self organization.
Encounters Graphs
Let's build a graph representing the encounters in a given frame...
End of explanation
"""


frame_encounters, c_map = build_frame_encounters_graph(nest, 1000)
nx.draw_circular(frame_encounters, node_color=c_map, node_size=1000, alpha=0.9, with_labels=True, edge_color='red', linewidths=3, width=5, font_color='white')

"""
Explanation: And here's how to visualize it:
End of explanation
"""


def build_encounters_graph(nest):
    G = nx.Graph()
    color_map = []
    for termite in nest.termites:
        G.add_node(termite.label)
        if termite.caste == 'S':
            color_map.append('gray')
        else:
            color_map.append('black')
    for termite in nest.termites:
        for other in nest.termites:
            if termite != other:
                if termite.trail[f'encountering_{other.label}'].sum() > 0:
                    G.add_edge(termite.label, other.label)
                    
    return G
        

"""
Explanation: And we can also create the graph for the entire experiment...
End of explanation
"""


all_encounters = build_encounters_graph(nest)
nx.draw_circular(all_encounters, node_color=c_map, node_size=1000, alpha=0.9, with_labels=True, edge_color='red', linewidths=3, width=5, font_color='white')

"""
Explanation: And see how it looks like:
End of explanation
"""


plot = MatrixPlot(all_encounters)
plot.draw()

"""
Explanation: This one is better visualized with a matrix plot!
End of explanation
"""


