from dataclasses import dataclass
from random import sample
from typing import List

import pandas as pd

from sawaw import SAWAWEntry, analyze_gpt3_5

# Create a DataFrame from ./combined_data_with_aspects.csv
df = pd.read_csv("data/combined_data_with_aspects.csv")
# Keep only text and aspect columns
df = df[["text", "aspect"]]
# Remove rows with NaN values
df = df.dropna()
# Remove rows with empty strings
df = df[df["text"] != ""]
# Create a list of SAWAWEntry objects
entries = []
for index, row in df.iterrows():
    comment = row["text"]
    aspect_words = row["aspect"].split(",")
    aspect_words = [aspect_word.strip() for aspect_word in aspect_words]
    entries.append(SAWAWEntry(comment, aspect_words))

# Pick 10 random entries
entries = sample(entries, 3)
# Query the entries
for entry in entries:
    print(entry)
    analyze_gpt3_5(entry)
    SAWAWEntry.print_legends()
    print(entry)
    input("Press Enter to continue...")