import streamlit as st
import time
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
        }
        self.frame_width = 600
        self.frame_height = 400

    def display_care_plan_animation(self, final_plan):
        st.write("Visualizing Care Plan...")
        steps = final_plan.split(", ")
        
        # Create all frames
        frames = self.create_all_frames(steps)
        
        # Create a placeholder for our animation
        animation_placeholder = st.empty()
        
        # Display frames
        for i, frame in enumerate(frames):
            animation_placeholder.image(frame)
            progress = int((i + 1) / len(frames) * 100)
            st.progress(progress)
            time.sleep(0.5)  # Adjust speed of animation
        
        st.success("Care Plan Visualization Complete!")
        st.balloons()

    def create_all_frames(self, steps):
        frames = []
        
        # Create intro frame
        intro_frame = self.create_intro_frame()
        frames.append(intro_frame)
        
        # Create step frames
        for i, step in enumerate(steps):
            step_frame = self.create_step_frame(step, i+1, len(steps))
            frames.append(step_frame)
        
        # Create transition frames
        transition_frames = self.create_transition_frames()
        frames.extend(transition_frames)
        
        # Create final frame
        final_frame = self.create_final_frame()
        frames.append(final_frame)
        
        return frames

    def create_intro_frame(self):
        img = Image.new('RGB', (self.frame_width, self.frame_height), color='white')
        d = ImageDraw.Draw(img)
        font = ImageFont.load_default()
        d.text((self.frame_width//2, self.frame_height//2), "Care Plan Animation", fill='black', font=font, anchor="mm")
        return img

    def create_step_frame(self, step, step_num, total_steps):
        img = Image.new('RGB', (self.frame_width, self.frame_height), color='white')
        d = ImageDraw.Draw(img)
        
        color = next((self.colors[key] for key in self.colors if key.lower() in step.lower()), '#FFFFFF')
        
        d.rectangle([50, 50, self.frame_width-50, self.frame_height-50], fill=color, outline='black')
        
        font = ImageFont.load_default()
        d.text((self.frame_width//2, 70), f"Step {step_num}/{total_steps}", fill='black', font=font, anchor="mm")
        d.text((self.frame_width//2, self.frame_height//2), step, fill='black', font=font, anchor="mm")
        
        return img

    def create_transition_frames(self):
        frames = []
        for i in range(10):  # 10 transition frames
            img = Image.new('RGB', (self.frame_width, self.frame_height), color='white')
            d = ImageDraw.Draw(img)
            size = 20 + i * 5  # Growing circle
            d.ellipse([self.frame_width//2 - size, self.frame_height//2 - size, 
                       self.frame_width//2 + size, self.frame_height//2 + size], 
                      fill='red')
            frames.append(img)
        return frames

    def create_final_frame(self):
        img = Image.new('RGB', (self.frame_width, self.frame_height), color='white')
        d = ImageDraw.Draw(img)
        font = ImageFont.load_default()
        d.text((self.frame_width//2, self.frame_height//2), "Care Plan Complete!", fill='black', font=font, anchor="mm")
        return img