# app.py - Final Working Version with Animation
import streamlit as st
import os
import time
import random
import numpy as np
import matplotlib.pyplot as plt
import base64
from PIL import Image
import io
import pandas as pd
import plotly.express as px

# Configure Streamlit page
st.set_page_config(
    page_title="StoryCoder - Learn Python Through Stories",
    page_icon="🧙‍♂️",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS for colorful UI
st.markdown("""
<style>
    @import url('https://fonts.googleapis.com/css2?family=Comic+Neue:wght@700&display=swap');
    
    :root {
        --primary: #FF6B6B;
        --secondary: #4ECDC4;
        --accent: #FFD166;
        --dark: #1A535C;
        --light: #F7FFF7;
    }
    
    body {
        background: linear-gradient(135deg, var(--light) 0%, #E8F4F8 100%);
        font-family: 'Comic Neue', cursive;
    }
    
    .stApp {
        background: url('https://www.transparenttextures.com/patterns/cartographer.png');
    }
    
    .story-box {
        background-color: white;
        border-radius: 20px;
        padding: 25px;
        box-shadow: 0 8px 16px rgba(26, 83, 92, 0.15);
        border: 3px solid var(--accent);
        margin-bottom: 25px;
    }
    
    .header {
        color: var(--dark);
        text-shadow: 2px 2px 4px rgba(0,0,0,0.1);
    }
    
    .concept-card {
        background: linear-gradient(145deg, #ffffff, #f0f0f0);
        border-radius: 15px;
        padding: 15px;
        margin: 10px 0;
        border-left: 5px solid var(--secondary);
        box-shadow: 0 4px 8px rgba(0,0,0,0.05);
    }
    
    .stButton>button {
        background: linear-gradient(45deg, var(--primary), var(--secondary));
        color: white;
        border-radius: 12px;
        padding: 10px 24px;
        font-weight: bold;
        font-size: 18px;
        border: none;
        transition: all 0.3s;
    }
    
    .stButton>button:hover {
        transform: scale(1.05);
        box-shadow: 0 6px 12px rgba(0,0,0,0.15);
    }
    
    .stTextInput>div>div>input {
        border-radius: 12px;
        padding: 12px;
        border: 2px solid var(--accent);
    }
    
    .tabs {
        display: flex;
        gap: 10px;
        margin-bottom: 20px;
        overflow-x: auto;
    }
    
    .tab {
        padding: 10px 20px;
        background-color: var(--accent);
        border-radius: 10px;
        cursor: pointer;
        font-weight: bold;
        white-space: nowrap;
    }
    
    .tab.active {
        background-color: var(--secondary);
        color: white;
    }
    
    @media (max-width: 768px) {
        .tabs {
            flex-wrap: wrap;
        }
    }
    
    .animation-container {
        background-color: #1a1a2e;
        border-radius: 15px;
        padding: 20px;
        box-shadow: 0 8px 32px rgba(0, 0, 0, 0.3);
        margin-bottom: 25px;
        position: relative;
        overflow: hidden;
    }
    
    .animation-canvas {
        border-radius: 10px;
        overflow: hidden;
        margin: 0 auto;
        display: block;
        max-width: 100%;
    }
    
    .character {
        font-size: 48px;
        text-align: center;
        margin: 10px 0;
    }
</style>
""", unsafe_allow_html=True)

# Concept database
CONCEPTS = {
    "loop": {
        "name": "Loop",
        "emoji": "🔄",
        "description": "Loops repeat actions multiple times",
        "example": "for i in range(5):\n    print('Hello!')",
        "color": "#FF9E6D"
    },
    "conditional": {
        "name": "Conditional",
        "emoji": "❓",
        "description": "Conditionals make decisions in code",
        "example": "if sunny:\n    go_outside()\nelse:\n    stay_inside()",
        "color": "#4ECDC4"
    },
    "function": {
        "name": "Function",
        "emoji": "✨",
        "description": "Functions are reusable blocks of code",
        "example": "def greet(name):\n    print(f'Hello {name}!')",
        "color": "#FFD166"
    },
    "variable": {
        "name": "Variable",
        "emoji": "📦",
        "description": "Variables store information",
        "example": "score = 10\nplayer = 'Alex'",
        "color": "#FF6B6B"
    },
    "list": {
        "name": "List",
        "emoji": "📝",
        "description": "Lists store collections of items",
        "example": "fruits = ['apple', 'banana', 'orange']",
        "color": "#1A535C"
    }
}

# Character database
CHARACTERS = [
    {"name": "rabbit", "emoji": "🐰", "color": "#FFB6C1"},
    {"name": "dragon", "emoji": "🐉", "color": "#FF6347"},
    {"name": "cat", "emoji": "🐱", "color": "#DDA0DD"},
    {"name": "dog", "emoji": "🐶", "color": "#FFD700"},
    {"name": "knight", "emoji": "🤺", "color": "#87CEEB"},
    {"name": "wizard", "emoji": "🧙", "color": "#98FB98"},
    {"name": "scientist", "emoji": "🔬", "color": "#20B2AA"},
    {"name": "pirate", "emoji": "🏴‍☠️", "color": "#FFA500"}
]

# Animation templates
ANIMATION_TEMPLATES = {
    "loop": {
        "description": "Character moving to a target multiple times",
        "code": """
# Loop Animation
import matplotlib.pyplot as plt
import numpy as np

fig, ax = plt.subplots(figsize=(10, 6))
ax.set_xlim(0, 10)
ax.set_ylim(0, 10)
ax.axis('off')

character = ax.text(1, 5, "{emoji}", fontsize=48, color='{color}')
target = ax.text(9, 5, "🎯", fontsize=48)
ax.text(5, 9, "{story}", ha='center', fontsize=14)

for i in range({count}):
    for pos in np.linspace(1, 9, 20):
        character.set_position((pos, 5))
        plt.pause(0.05)
plt.show()
        """
    },
    "conditional": {
        "description": "Character making a decision based on a condition",
        "code": """
# Conditional Animation
import matplotlib.pyplot as plt
import numpy as np

fig, ax = plt.subplots(figsize=(10, 6))
ax.set_xlim(0, 10)
ax.set_ylim(0, 10)
ax.axis('off')

character = ax.text(5, 5, "{emoji}", fontsize=48, color='{color}')
ax.text(5, 9, "{story}", ha='center', fontsize=14)

# Condition: {condition}
if {condition}:
    decision = ax.text(7, 7, "✅ Yes", fontsize=24, color='green')
    path = np.linspace(5, 8, 20)
else:
    decision = ax.text(7, 7, "❌ No", fontsize=24, color='red')
    path = np.linspace(5, 2, 20)

for pos in path:
    character.set_position((pos, 5))
    plt.pause(0.05)
plt.show()
        """
    },
    "function": {
        "description": "Character performing an action multiple times",
        "code": """
# Function Animation
import matplotlib.pyplot as plt
import numpy as np

fig, ax = plt.subplots(figsize=(10, 6))
ax.set_xlim(0, 10)
ax.set_ylim(0, 10)
ax.axis('off')

character = ax.text(5, 5, "{emoji}", fontsize=48, color='{color}')
ax.text(5, 9, "{story}", ha='center', fontsize=14)

def perform_action():
    for _ in range(5):
        character.set_fontsize(60)
        plt.pause(0.1)
        character.set_fontsize(48)
        plt.pause(0.1)

for i in range({count}):
    perform_action()
    character.set_position((5 + i, 5))
plt.show()
        """
    }
}

def analyze_story(story):
    """Analyze story and identify programming concepts"""
    story_lower = story.lower()
    detected_concepts = []
    
    # Check for loops
    if any(word in story_lower for word in ["times", "repeat", "again", "multiple"]):
        detected_concepts.append("loop")
    
    # Check for conditionals
    if any(word in story_lower for word in ["if", "when", "unless", "whether"]):
        detected_concepts.append("conditional")
    
    # Check for functions
    if any(word in story_lower for word in ["make", "create", "do", "perform", "cast"]):
        detected_concepts.append("function")
    
    # Check for variables
    if any(word in story_lower for word in ["is", "has", "set to", "value"]):
        detected_concepts.append("variable")
    
    # Check for lists
    if any(word in story_lower for word in ["and", "many", "several", "collection", "items"]):
        detected_concepts.append("list")
    
    return list(set(detected_concepts))

def extract_count_from_story(story):
    """Extract a number from the story to use in animations"""
    for word in story.split():
        if word.isdigit():
            return min(int(word), 10)
    return 3  # Default value

def create_animation(story, concepts):
    """Create an animation visualization based on the story and concepts"""
    try:
        # Choose a random character
        character = random.choice(CHARACTERS)
        count = extract_count_from_story(story)
        
        # Create figure
        fig, ax = plt.subplots(figsize=(10, 6))
        ax.set_facecolor('#f0f8ff')
        ax.set_xlim(0, 10)
        ax.set_ylim(0, 10)
        ax.axis('off')
        
        # Add title
        ax.text(5, 9, '✨ Your Story Animation ✨', 
                fontsize=20, ha='center', color='purple', fontweight='bold')
        
        # Add story text
        story_text = story[:80] + ('...' if len(story) > 80 else '')
        ax.text(5, 8, f'"{story_text}"', 
                fontsize=14, ha='center', color='#333')
        
        # Add character
        ax.text(5, 5, character["emoji"], 
                fontsize=100, ha='center', color=character["color"])
        
        # Add concept visualization
        if "loop" in concepts:
            ax.text(5, 3, f"🔄 Repeating {count} times", 
                    fontsize=18, ha='center', color=CONCEPTS["loop"]["color"])
        elif "conditional" in concepts:
            condition = random.choice(["sunny", "raining", "dark"])
            ax.text(5, 3, f"❓ Checking if it's {condition}", 
                    fontsize=18, ha='center', color=CONCEPTS["conditional"]["color"])
        elif "function" in concepts:
            ax.text(5, 3, f"✨ Performing action {count} times", 
                    fontsize=18, ha='center', color=CONCEPTS["function"]["color"])
        else:
            ax.text(5, 3, "📝 Creating your story visualization", 
                    fontsize=18, ha='center', color=CONCEPTS["variable"]["color"])
        
        # Add footer
        ax.text(5, 1, "Created with StoryCoder", 
                fontsize=12, ha='center', style='italic', color='gray')
        
        # Save to buffer
        buf = io.BytesIO()
        plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0.5, dpi=150)
        buf.seek(0)
        plt.close()
        return buf
    
    except Exception as e:
        st.error(f"Animation error: {str(e)}")
        return None

def create_interactive_animation(story, concepts):
    """Create an interactive animation using Plotly"""
    try:
        # Choose a random character
        character = random.choice(CHARACTERS)
        count = extract_count_from_story(story)
        
        # Create animation data
        frames = []
        for i in range(count):
            frames.append({
                "frame": i,
                "x": np.random.uniform(1, 9),
                "y": np.random.uniform(1, 9),
                "size": np.random.uniform(10, 30),
                "character": character["emoji"]
            })
        
        df = pd.DataFrame(frames)
        
        # Create animated scatter plot
        fig = px.scatter(
            df, 
            x="x", 
            y="y", 
            animation_frame="frame",
            text="character",
            size="size",
            size_max=45,
            color_discrete_sequence=[character["color"]],
            range_x=[0, 10],
            range_y=[0, 10]
        )
        
        # Customize layout
        fig.update_layout(
            title=f'Animation for: "{story[:50]}{"..." if len(story) > 50 else ""}"',
            showlegend=False,
            plot_bgcolor='rgba(240,248,255,1)',
            paper_bgcolor='rgba(240,248,255,1)',
            width=800,
            height=600
        )
        
        fig.update_traces(
            textfont_size=30,
            textposition='middle center'
        )
        
        fig.layout.updatemenus[0].buttons[0].args[1]["frame"]["duration"] = 1000
        fig.layout.updatemenus[0].buttons[0].args[1]["transition"]["duration"] = 500
        
        return fig
    
    except Exception as e:
        st.error(f"Interactive animation error: {str(e)}")
        return None

def generate_animation_code(story, concepts):
    """Generate Python animation code based on story"""
    try:
        # Choose a random character
        character = random.choice(CHARACTERS)
        count = extract_count_from_story(story)
        concept = concepts[0] if concepts else "loop"
        
        # Get the appropriate template
        template = ANIMATION_TEMPLATES.get(concept, ANIMATION_TEMPLATES["loop"])
        
        # Fill in the template
        condition = random.choice(["True", "False"])
        code = template["code"].format(
            emoji=character["emoji"],
            color=character["color"],
            story=story[:80] + ('...' if len(story) > 80 else ''),
            count=count,
            condition=condition
        )
        
        return code, template["description"]
    
    except Exception as e:
        return f"# Error generating code\nprint('Could not generate code: {str(e)}')", ""

def create_story_image(story):
    """Create a story image with Matplotlib"""
    try:
        # Create figure
        fig, ax = plt.subplots(figsize=(10, 6))
        ax.set_facecolor('#f0f8ff')
        ax.set_xlim(0, 10)
        ax.set_ylim(0, 10)
        ax.axis('off')
        
        # Add title
        ax.text(5, 8, '✨ Your Story ✨', fontsize=24, 
                ha='center', color='purple', fontweight='bold')
        
        # Add story text (wrapped)
        words = story.split()
        lines = []
        current_line = []
        
        for word in words:
            if len(' '.join(current_line + [word])) < 60:
                current_line.append(word)
            else:
                lines.append(' '.join(current_line))
                current_line = [word]
        
        if current_line:
            lines.append(' '.join(current_line))
        
        for i, line in enumerate(lines):
            ax.text(5, 6.5 - i*0.7, line, fontsize=14, ha='center', color='#333')
        
        # Add decoration
        ax.text(2, 2, '🐰', fontsize=40, ha='center')
        ax.text(8, 2, '🐉', fontsize=40, ha='center')
        ax.text(5, 1, 'Created with StoryCoder', fontsize=12, 
                ha='center', style='italic', color='gray')
        
        # Save to buffer
        buf = io.BytesIO()
        plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0.5, dpi=150)
        buf.seek(0)
        plt.close()
        return buf
    
    except Exception as e:
        st.error(f"Image generation error: {str(e)}")
        return None

def main():
    """Main application function"""
    st.title("🧙‍♂️ StoryCoder - Learn Python Through Stories!")
    st.subheader("Turn your story into an animation and discover coding secrets!")
    
    # Initialize session state
    if 'story' not in st.session_state:
        st.session_state.story = ""
    if 'concepts' not in st.session_state:
        st.session_state.concepts = []
    if 'animation' not in st.session_state:
        st.session_state.animation = None
    if 'interactive_animation' not in st.session_state:
        st.session_state.interactive_animation = None
    if 'animation_code' not in st.session_state:
        st.session_state.animation_code = ""
    if 'code_description' not in st.session_state:
        st.session_state.code_description = ""
    if 'active_tab' not in st.session_state:
        st.session_state.active_tab = "story"
    
    # Create tabs
    tabs = st.empty()
    tab_cols = st.columns(5)
    with tab_cols[0]:
        if st.button("📖 Create Story"):
            st.session_state.active_tab = "story"
    with tab_cols[1]:
        if st.button("🎬 Animation"):
            st.session_state.active_tab = "animation"
    with tab_cols[2]:
        if st.button("🔍 Concepts"):
            st.session_state.active_tab = "concepts"
    with tab_cols[3]:
        if st.button("💻 Code"):
            st.session_state.active_tab = "code"
    with tab_cols[4]:
        if st.button("🔄 Reset"):
            st.session_state.story = ""
            st.session_state.concepts = []
            st.session_state.animation = None
            st.session_state.interactive_animation = None
            st.session_state.animation_code = ""
            st.session_state.code_description = ""
            st.session_state.active_tab = "story"
    
    # Story creation tab
    if st.session_state.active_tab == "story":
        with st.container():
            st.header("📖 Create Your Story")
            st.write("Write a short story (2-5 sentences) and I'll turn it into an animation!")
            
            story = st.text_area(
                "Your story:",
                height=200,
                placeholder="Once upon a time, a rabbit hopped 3 times to reach a carrot...",
                value=st.session_state.story,
                key="story_input"
            )
            
            if st.button("Create Animation!", use_container_width=True):
                if len(story) < 10:
                    st.error("Your story needs to be at least 10 characters long!")
                else:
                    st.session_state.story = story
                    with st.spinner("🧠 Analyzing your story for coding concepts..."):
                        st.session_state.concepts = analyze_story(story)
                    
                    with st.spinner("🎬 Creating your animation..."):
                        st.session_state.animation = create_animation(
                            story, st.session_state.concepts
                        )
                        
                        st.session_state.interactive_animation = create_interactive_animation(
                            story, st.session_state.concepts
                        )
                        
                        st.session_state.animation_code, st.session_state.code_description = generate_animation_code(
                            story, st.session_state.concepts
                        )
                    
                    st.session_state.active_tab = "animation"
                    st.rerun()
            
            # Show examples
            st.subheader("✨ Story Examples")
            col1, col2, col3 = st.columns(3)
            with col1:
                st.caption("Loop Example")
                st.code('"A dragon breathes fire 5 times at the castle"', language="text")
            with col2:
                st.caption("Conditional Example")
                st.code('"If it rains, the cat stays inside, else it goes out"', language="text")
            with col3:
                st.caption("Function Example")
                st.code('"A wizard casts a spell to make flowers grow"', language="text")
    
    # Animation tab
    elif st.session_state.active_tab == "animation":
        st.header("🎬 Your Story Animation")
        
        if not st.session_state.story:
            st.warning("Please create a story first!")
            st.session_state.active_tab = "story"
            st.rerun()
        
        # Display animation
        st.markdown(f"""
        <div class="animation-container">
            <h3 style="color: white; text-align: center;">"{st.session_state.story[:60]}{'...' if len(st.session_state.story) > 60 else ''}"</h3>
        </div>
        """, unsafe_allow_html=True)
        
        # Display static animation
        if st.session_state.animation:
            st.image(st.session_state.animation, use_container_width=True)
        else:
            st.warning("Static animation couldn't be generated. Showing story image instead.")
            story_img = create_story_image(st.session_state.story)
            if story_img:
                st.image(story_img, use_container_width=True)
        
        # Display interactive animation
        st.subheader("📊 Interactive Animation")
        if st.session_state.interactive_animation:
            st.plotly_chart(st.session_state.interactive_animation, use_container_width=True)
        else:
            st.info("Interactive animation preview. The real animation would run on your computer!")
            st.image("https://i.imgur.com/5X8jYAy.gif", use_container_width=True)
        
        st.success("✨ Animation created successfully!")
        st.caption("This animation was generated with Python code based on your story!")
        
        if st.button("Reveal Coding Secrets!", use_container_width=True):
            st.session_state.active_tab = "concepts"
            st.rerun()
    
    # Concepts tab
    elif st.session_state.active_tab == "concepts":
        st.header("🔍 Coding Concepts in Your Story")
        st.subheader("We secretly used these programming concepts:")
        
        if not st.session_state.concepts:
            st.warning("No concepts detected in your story! Try adding words like '3 times', 'if', or 'make'.")
        else:
            for concept in st.session_state.concepts:
                if concept in CONCEPTS:
                    details = CONCEPTS[concept]
                    st.markdown(f"""
                    <div class="concept-card" style="border-left: 5px solid {details['color']};">
                        <div style="display:flex; align-items:center; gap:15px;">
                            <span style="font-size:36px;">{details['emoji']}</span>
                            <h3 style="color:{details['color']};">{details['name']}</h3>
                        </div>
                        <p>{details['description']}</p>
                        <pre style="background:#f0f0f0; padding:10px; border-radius:8px;">{details['example']}</pre>
                    </div>
                    """, unsafe_allow_html=True)
        
        if st.button("See the Magic Code!", use_container_width=True):
            st.session_state.active_tab = "code"
            st.rerun()
    
    # Code tab
    elif st.session_state.active_tab == "code":
        st.header("💻 The Magic Code Behind Your Animation")
        st.write("Here's the Python code that would create your animation:")
        
        if st.session_state.animation_code:
            st.markdown(f"**{st.session_state.code_description}**")
            st.code(st.session_state.animation_code, language="python")
            
            # Download button
            st.download_button(
                label="Download Animation Code",
                data=st.session_state.animation_code,
                file_name="story_animation.py",
                mime="text/python",
                use_container_width=True
            )
            
            st.info("💡 To run this animation on your computer:")
            st.markdown("""
            1. Install Python from [python.org](https://python.org)
            2. Install required libraries: `pip install matplotlib numpy`
            3. Copy the code above into a file named `animation.py`
            4. Run it with: `python animation.py`
            """)
            
            st.success("🎉 When you run this code, you'll see your story come to life!")
        else:
            st.warning("No code generated yet!")
        
        if st.button("Create Another Story!", use_container_width=True):
            st.session_state.active_tab = "story"
            st.session_state.story = ""
            st.rerun()

if __name__ == "__main__":
    main()