Update components/animation_engineer.py

components/animation_engineer.py

@@ -1,16 +1,97 @@
+import streamlit as st
+import time
+import random
+from PIL import Image, ImageDraw, ImageFont
+import io
+
 class AnimationEngineer:
+    def __init__(self):
+        self.colors = {
+            'Parkinson': '#FFB3BA',  # Light Pink
+            'Alzheimer': '#BAFFC9',  # Light Green
+            'Treatment A': '#BAE1FF',  # Light Blue
+            'Treatment B': '#FFFFBA',  # Light Yellow
+            'physiotherapy': '#FFD700',  # Gold
+            'cognitive therapy': '#DDA0DD'  # Plum
+        }
+
     def display_care_plan_animation(self, final_plan):
         st.write("Visualizing Care Plan...")
-        progress_bar = st.progress(0)
-        status_text = st.empty()
-
         steps = final_plan.split(", ")
+        
+        # Create a placeholder for our animation
+        animation_placeholder = st.empty()
+        
+        # Progress bar
+        progress_bar = st.progress(0)
+        
         for i, step in enumerate(steps):
-            status_text.text(f"Step {i+1}/{len(steps)}: {step}")
-            for percent_complete in range(0, 101, 10):
-                time.sleep(0.1)  # Simulate some work being done
-                progress_bar.progress(percent_complete)
-            time.sleep(0.5)  # Pause between steps
+            # Update progress
+            progress = int((i + 1) / len(steps) * 100)
+            progress_bar.progress(progress)
+            
+            # Create and display step visualization
+            img = self.create_step_visualization(step, i+1, len(steps))
+            animation_placeholder.image(img)
+            
+            time.sleep(1.5)  # Pause between steps
+        
+        # Final animation
+        self.display_final_animation(animation_placeholder)
+        
+        st.success("Care Plan Visualization Complete!")
+        st.balloons()
+
+    def create_step_visualization(self, step, step_num, total_steps):
+        # Create a new image with a white background
+        img = Image.new('RGB', (400, 200), color='white')
+        d = ImageDraw.Draw(img)
+        
+        # Choose color based on step content
+        color = next((self.colors[key] for key in self.colors if key.lower() in step.lower()), '#FFFFFF')
+        
+        # Draw a colored rectangle
+        d.rectangle([20, 20, 380, 180], fill=color, outline='black')
+        
+        # Add text
+        font = ImageFont.load_default()
+        d.text((30, 30), f"Step {step_num}/{total_steps}", fill='black', font=font)
+        d.text((30, 60), step, fill='black', font=font)
+        
+        # Convert PIL Image to bytes
+        img_byte_arr = io.BytesIO()
+        img.save(img_byte_arr, format='PNG')
+        
+        return img_byte_arr.getvalue()
+
+    def display_final_animation(self, placeholder):
+        # Create a simple animation of a pulsing heart
+        for _ in range(5):  # Pulse 5 times
+            for size in range(50, 100, 10):  # Grow
+                self.display_heart(placeholder, size)
+                time.sleep(0.1)
+            for size in range(100, 50, -10):  # Shrink
+                self.display_heart(placeholder, size)
+                time.sleep(0.1)
 
-        status_text.text("Care Plan Visualization Complete!")
-        st.balloons()
+    def display_heart(self, placeholder, size):
+        # Create a new image with a white background
+        img = Image.new('RGB', (200, 200), color='white')
+        d = ImageDraw.Draw(img)
+        
+        # Draw a red heart
+        d.polygon([
+            (100, 100 - size//2),
+            (100 - size//2, 100 - size//4),
+            (100 - size//2, 100 + size//4),
+            (100, 100 + size//2),
+            (100 + size//2, 100 + size//4),
+            (100 + size//2, 100 - size//4)
+        ], fill='red')
+        
+        # Convert PIL Image to bytes
+        img_byte_arr = io.BytesIO()
+        img.save(img_byte_arr, format='PNG')
+        
+        # Display the image
+        placeholder.image(img_byte_arr.getvalue())