import streamlit as st
import pandas as pd
import plotly.express as px
import json
import time

st.set_page_config(layout="wide")

# Load the positions data
positions_df = pd.read_csv('data/positions.csv')

# Preprocess the 'Positions' column to convert from string to dictionary
positions_df['Positions'] = positions_df['Positions'].apply(json.loads)

# Create two columns in the Streamlit app
left_column, right_column = st.columns(2)

# In the left column, allow the user to input a common goal for all agents
common_goal = left_column.text_area("Common Goal for the Agents", "Enter a goal")

# Initialize lists to hold agent names, tasks, and backstories from the dataframe
agent_names = list(positions_df['Agent'].unique())
tasks = list(positions_df['Task'].unique())
backstories = list(positions_df['Backstory'].unique())

# Placeholder for the Plotly chart in the right column
chart_placeholder = right_column.empty()

# Index for accessing rows in positions_df, starting with the first row after the header
current_index = 0

while True:  # This loop will simulate real-time data updates
    # Fetch the current row's position data
    current_positions = positions_df.iloc[current_index % len(positions_df)]['Positions']
    
    # Convert current_positions into a list of dicts suitable for DataFrame creation
    new_data_points = [{'Agent': agent, 'x': pos[0], 'y': pos[1]} for agent, pos in current_positions.items()]
    
    # Create a new DataFrame from the list of new data points
    new_data = pd.DataFrame(new_data_points)
    
    # Create a new Plotly figure for the scatter plot, using Agent names for colors
    fig = px.scatter(new_data, x='x', y='y', color='Agent', title="Random Dots with Assigned Agents")
    
    # Update the figure layout to better visualize the dots
    fig.update_layout(xaxis_title='X Axis',
                      yaxis_title='Y Axis',
                      autosize=True,
                      xaxis=dict(range=[-5, 5]),
                      yaxis=dict(range=[-5, 5]))
    
    # Update traces to adjust the appearance of the dots, setting radius to 10
    fig.update_traces(marker=dict(size=20))  # Size adjusted for radius of 10

    # Display the figure in the right column
    chart_placeholder.plotly_chart(fig, use_container_width=True)

    # Increment the index to move to the next row for the next update
    current_index += 1

    # Pause for a moment before updating the chart with new dots
    time.sleep(1)