import json

import pandas as pd
import numpy as np

from gradiend.setups.emotion import read_vad

vad = read_vad(use_antonyms=True)

terms = vad['term']
antonyms = vad['antonym']

# Example epsilon thresholds
valence_epsilon = 0.25  # max valence difference to consider "similar"
arousal_threshold = 0.2  # min arousal difference to consider "changed"

# Merge the DataFrame with itself to align each term with its antonym
merged = vad.merge(
    vad[['term', 'valence', 'arousal', 'dominance']],
    left_on='antonym', right_on='term',
    suffixes=('', '_antonym')
)

# Compute differences
merged['valence_diff'] = (merged['valence'] - merged['valence_antonym']).abs()
merged['arousal_diff'] = (merged['arousal'] - merged['arousal_antonym']).abs()
merged['dominance_diff'] = (merged['dominance'] - merged['dominance_antonym']).abs()

# Filter: similar valence but different arousal
result_arousal = merged[
    (merged['valence_diff'] <= valence_epsilon) # todo also similar dominance?
]

# Optionally, keep only the original columns of interest
columns = vad.columns
result_arousal = result_arousal[columns]
result_arousal.to_csv('data/emotion/vad_arousal.csv', index=False)
print(f"Saved {len(result_arousal)} terms with similar valence but different arousal.")

result_valence = merged[merged['arousal_diff'] <= valence_epsilon]
result_valence = result_valence[columns]
result_valence.to_csv('data/emotion/vad_valence.csv', index=False)
print(f"Saved {len(result_valence)} terms with similar arousal but different valence.")

result_dominance = merged[(merged['arousal_diff'] <= valence_epsilon) & (merged['dominance_diff'] <= arousal_threshold)]
result_dominance = result_dominance[columns]
result_dominance.to_csv('data/emotion/vad_dominance.csv', index=False)
print(f"Saved {len(result_dominance)} terms with similar valence but different dominance.")



valence_opposite_terms = result_valence['term'].unique().tolist()
print(f"Found {len(valence_opposite_terms)} terms with opposite valence.")
arousal_opposite_terms = result_arousal['term'].unique().tolist()
print(f"Found {len(arousal_opposite_terms)} terms with opposite arousal.")
dominance_opposite_terms = result_dominance['term'].unique().tolist()
print(f"Found {len(dominance_opposite_terms)} terms with opposite dominance.")


# Save the antonyms for valence and arousal
json.dump(valence_opposite_terms, open('data/emotion/valence_opposite_terms.json', 'w'), indent=2)
json.dump(arousal_opposite_terms, open('data/emotion/arousal_opposite_terms.json', 'w'), indent=2)
json.dump(dominance_opposite_terms, open('data/emotion/dominance_opposite_terms.json', 'w'), indent=2)