from gradiend.setups.emotion import read_vad

import numpy as np
from nltk.corpus import wordnet as wn
from scipy.spatial.distance import cosine
import pandas as pd
import numpy as np
from nltk.corpus import wordnet as wn
from scipy.spatial.distance import cosine


def get_antonym_candidates_from_wordnet(terms):
    antonym_dict = {}
    for term in terms:
        antonym_set = set()
        for syn in wn.synsets(term):
            for lemma in syn.lemmas():
                if lemma.name().lower() != term:
                    continue  # Only consider exact lemma match

                for ant in lemma.antonyms():
                    ant_term = ant.name().lower()
                    if ant_term in terms and ant_term != term:
                        antonym_set.add(ant_term)
        if antonym_set:
            antonym_dict[term] = antonym_set
    return antonym_dict


def create_antonym_mapping(vad_threshold=0.5):
    vad = read_vad()
    vad = vad.dropna()

    # Lowercase all terms for consistency
    vad["term"] = vad["term"].str.lower()
    vad_dict = vad.set_index("term")[["valence", "arousal", "dominance"]].to_dict(orient="index")
    terms_in_vad = set(vad["term"])

    # Get antonym candidates from WordNet
    antonym_candidates = get_antonym_candidates_from_wordnet(terms_in_vad)

    antonym_mapping = {}

    for term, candidates in antonym_candidates.items():
        v1 = np.array(list(vad_dict[term].values()))
        best_antonym = None
        best_similarity = 1  # cosine similarity ranges from -1 (opposite) to 1 (identical)

        for ant in candidates:
            v2 = np.array(list(vad_dict[ant].values()))
            similarity = 1 - cosine(v1, v2)

            if similarity < best_similarity:
                best_similarity = similarity
                best_antonym = ant

        if best_similarity < (1 - vad_threshold):
            antonym_mapping[term] = best_antonym

    return antonym_mapping


if __name__ == "__main__":
    antonym_mapping = create_antonym_mapping(vad_threshold=0.5)
    print(f"Found {len(antonym_mapping)} antonym pairs:")
    for term, antonym in antonym_mapping.items():
        print(f"{term} ↔ {antonym}")

    output_file = 'data/emotion/antonyms.json'
    import json
    with open(output_file, 'w') as f:
        json.dump(antonym_mapping, f, indent=2)