import networkx as nx
from _arango import arango
import random
from print_color import *
import json
import datetime

# Create a new directed graph
G = nx.DiGraph()

q = "for doc in interrogations return doc"
interrogations = list(arango.db.aql.execute(q))


for interrogation in interrogations:
    if not 'mentioned_persons' in interrogation:
        continue
    person_key = interrogation['person_id'].split('/')[1]
    mentioned_persons = interrogation["mentioned_persons"]
    for mentioned_person in mentioned_persons:
        G.add_edge(
            person_key,
            mentioned_person,
            label=interrogation["_key"],
        )



# Write the graph to a GEXF file
current_time = datetime.datetime.now().strftime("%H-%M-%S")
filename = f"output_files/mentions_{current_time}.gexf"
nx.write_gexf(G, filename)