import ciga as cg
import pandas as pd
import openai

df = pd.DataFrame({
    'season': [1, 1, 1, 1, 1],
    'episode': [1, 1, 1, 1, 1],
    'scene': [1, 1, 1, 2, 2],
    'line': [1, 2, 3, 1, 2],
    'source': ['Alice', 'Bob', 'Charlie', 'Alice', 'Bob'],
    'dialogue': ['Hi Bob', 'Hello Alice', 'Hey everyone', 'Good morning', 'Morning Alice'],
    'action': ['Waive hand', 'Smile', '', 'Smile', 'Waive hand'],
    'scene_description': ['In the room', 'In the room', 'In the room', 'In the room', 'In the room']
})

client = openai.OpenAI()  # Make sure to set your OPENAI_API_KEY in the environment variables
result = cg.infer_listeners(data=df,
                            position=('season', 'episode', 'scene', 'line'),
                            speaker='source',
                            dialogue='dialogue',
                            action='action',
                            scene_description='scene_description',
                            client=client,
                            mode='openai',
                            model='gpt-4o-2024-08-06',
                            max_tokens=200,
                            gap=0.5)
print(result)
# result.to_csv('inferred_listeners.csv', index=False)