import ciga as cg
import pandas as pd
import matplotlib.pyplot as plt
import os

from viztracer import VizTracer

data_path = os.path.join(os.path.dirname(__file__), '..', '..', 'm_data.csv')
df = pd.read_csv(data_path)

# custom weight function
# input: interaction (str)
# output: weight (float)
def weight_func(interaction):
    return 1

# load data
positions = ('Season', 'Episode', 'Scene', 'Line')
interactions = cg.prepare_data(df, ('Season', 'Episode', 'Scene', 'Line'),
                                source='Speaker', target='Listener', interaction='Words')

sub_interactions = cg.segment(interactions, start=(1, 1, 1, 1), end=(2, 1, 1, 1), positions=positions)

weights = cg.calculate_weights(sub_interactions, weight_func=weight_func)
print(weights.head(20))
agg_weights = cg.agg_weights(weights, positions, agg_func=lambda x: sum(x))

# create network
tg = cg.TGraph(data=agg_weights, positions=positions, directed=False)

# centrality analysis
res = cg.tgraph_degree(tg, weighted=True, accumulate=True, normalized=True)
res.head()

g = tg.get_graph((2, 1, 1, 1))
cg.graph_viz(g, output_file='graph.html', port=None, open_browser=True)
print(g.vcount(), g.ecount())

res.to_csv('centrality_res.csv', index=False)
