import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from all_arguments import arguments as arguments_dict # Arguments dictionary with sentiment information


# Step 1: Read the CSV file
df = pd.read_csv('speeches.csv', delimiter=';')

# Step 2: Extract relevant columns
# Assuming the arguments start from the 5th column onwards
arguments = df.columns[5:]
df_arguments = df[['_key', 'name'] + list(arguments)]

# Step 3: Create a binary matrix for arguments
# Convert the argument columns to integers
df_arguments[arguments] = df_arguments[arguments].apply(pd.to_numeric, errors='coerce').fillna(0).astype(int)

# Step 4: Sum the arguments for each politician
df_sum = df_arguments.groupby('name')[arguments].sum().reset_index()

# Step 5: Plot the data
plt.figure(figsize=(12, 8))
sns.heatmap(df_sum.set_index('name'), annot=True, cmap='coolwarm', cbar=True)
plt.title('Arguments Used by Politicians')
plt.xlabel('Arguments')
plt.ylabel('Politicians')

# Step 6: Color the x-axis labels based on sentiment
ax = plt.gca()
x_labels = ax.get_xticklabels()
for label in x_labels:
    argument = label.get_text()
    sentiment = arguments_dict.get(argument, {}).get('sentiment', 'neutral')
    if sentiment == 'positive':
        label.set_color('green')
    elif sentiment == 'negative':
        label.set_color('red')
    else:
        label.set_color('black')

plt.xticks(rotation=45, ha='right')
plt.tight_layout()

# Save the plot instead of showing it
plt.savefig('arguments_used_by_politicians.png')
plt.close()