Update app.py

app.py

@@ -1,8 +1,9 @@
 import random
 import numpy as np
+import streamlit as st
 import plotly.graph_objects as go
 from plotly.subplots import make_subplots
-import streamlit as st
+import time
 
 class Organelle:
     def __init__(self, type):
@@ -80,6 +81,12 @@ class Environment:
         self.light_level = 5  # Ambient light level
         self.cells = []
         self.time = 0
+        self.population_history = {
+            "prokaryote": [],
+            "early_eukaryote": [],
+            "advanced_eukaryote": [],
+            "plant_like": []
+        }
 
     def add_cell(self, cell):
         self.cells.append(cell)
@@ -131,10 +138,27 @@ class Environment:
                     cell.color = "darkgreen"
                     cell.size = 4
 
-    def visualize(self):
-        fig = make_subplots(rows=1, cols=2, subplot_titles=("Cell Distribution", "Population Over Time"))
+        # Record population counts
+        for cell_type in self.population_history.keys():
+            count = len([cell for cell in self.cells if cell.cell_type == cell_type])
+            self.population_history[cell_type].append(count)
+
+    def get_visualization_data(self):
+        cell_data = {
+            "prokaryote": {"x": [], "y": [], "size": [], "color": "lightblue"},
+            "early_eukaryote": {"x": [], "y": [], "size": [], "color": "green"},
+            "advanced_eukaryote": {"x": [], "y": [], "size": [], "color": "red"},
+            "plant_like": {"x": [], "y": [], "size": [], "color": "darkgreen"}
+        }
+
+        for cell in self.cells:
+            cell_data[cell.cell_type]["x"].append(cell.x)
+            cell_data[cell.cell_type]["y"].append(cell.y)
+            cell_data[cell.cell_type]["size"].append(cell.size * 3)
+
+        return cell_data, self.population_history
 
-        def setup_figure(env):
+def setup_figure(env):
     fig = make_subplots(rows=1, cols=2, subplot_titles=("Cell Distribution", "Population Over Time"))
 
     # Cell distribution