import pandas as pd
import re
import numpy as np

tweets_path = "../raw_data/tweets.csv"
accounts_path = "../raw_data/accounts.txt"

output_path = "../data/cooccurence_matrix.csv"


tweets = pd.read_csv(tweets_path)

with open(accounts_path, "r") as f:
    accounts_regex = f.read().replace("\n","")

# extract mentions from tweets
tweets["mentions"] = tweets.tweet_text.str.findall(accounts_regex, flags=re.IGNORECASE)

# some tweets have multiple mentions - we're gonna turn one row into n rows if the tweets has n mentions
tweets = tweets.explode("mentions")

# make sure all handles are lowercase
tweets.handle = tweets.handle.str.lower()
tweets.mentions = tweets.mentions.str.lower()

# TODO IDEA
# we could generate multiple different cocc matrices
# e.g.: weigh mentions by interaction counts, exclude RTs


# generate cooccurence matrix
mention_counts = tweets[["handle", "mentions"]].value_counts().reset_index()
cooc_matrix = pd.DataFrame(mention_counts).pivot(index="handle", columns="mentions",
        values=0).reset_index()

cooc_matrix.set_index("handle",inplace=True)

# networkx wants the matrix to be square, so make sure it is suare
# add columns
for r in cooc_matrix.index:
    if r not in cooc_matrix.columns:
        print(f"{r} is not in colums")
        cooc_matrix[r] = None

not_in_index = list(filter(lambda x: x not in cooc_matrix.index, cooc_matrix.columns))
new_rows = pd.DataFrame(index=not_in_index)
cooc_matrix = pd.concat([cooc_matrix, new_rows])


#save cooccurence matrix
cooc_matrix.to_csv(output_path, sep=";")
