import streamlit as st
import time
import random
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import numpy as np

# Set page config
st.set_page_config(page_title="Insertion Sort Visualizer", layout="centered")
st.title("🔄 Insertion Sort Algorithm Visualizer")
st.markdown("Watch how Insertion Sort builds a sorted array one element at a time!")

# Sidebar controls
st.sidebar.header("Controls")
array_size = st.sidebar.slider("Array Size", min_value=5, max_value=50, value=20)
speed = st.sidebar.slider("Animation Speed (seconds)", min_value=0.1, max_value=2.0, value=0.5, step=0.1)

# Generate random array
if st.sidebar.button("Generate New Array"):
    st.session_state.array = random.sample(range(1, 100), array_size)
    st.session_state.history = []
    st.session_state.steps = []
    st.rerun()

# Initialize array in session state
if "array" not in st.session_state:
    st.session_state.array = random.sample(range(1, 100), array_size)
    st.session_state.history = []
    st.session_state.steps = []

array = st.session_state.array.copy()

# Insertion Sort with step recording
def insertion_sort_with_steps(arr):
    steps = []
    history = []
    arr = arr.copy()
    
    for i in range(1, len(arr)):
        key = arr[i]
        j = i - 1
        
        # Record current state
        history.append(arr.copy())
        steps.append({
            'comparing': i,
            'key': key,
            'action': f"Pick {key} and insert into sorted portion"
        })
        
        while j >= 0 and arr[j] > key:
            arr[j + 1] = arr[j]
            j -= 1
            
            # Record shift
            history.append(arr.copy())
            steps.append({
                'comparing': j + 1,
                'inserting': i,
                'key': key,
                'action': f"Shift {arr[j+1]} right"
            })
        
        arr[j + 1] = key
        
        # Record after insertion
        history.append(arr.copy())
        steps.append({
            'inserted': j + 1,
            'key': key,
            'action': f"Insert {key} at position {j+1}"
        })
    
    # Final sorted array
    history.append(arr.copy())
    steps.append({'action': "Sorting Complete! 🎉", 'sorted': True})
    
    return history, steps

# Run sorting only when needed
if st.button("Start Sorting") or st.session_state.get("running", False):
    if not st.session_state.history:
        with st.spinner("Running Insertion Sort..."):
            history, steps = insertion_sort_with_steps(st.session_state.array)
            st.session_state.history = history
            st.session_state.steps = steps
            st.session_state.current_step = 0
            st.session_state.running = True
    
    # Animation placeholder
    chart = st.empty()
    info = st.empty()
    progress = st.progress(0)

    # Animation loop
    for idx, (state, step) in enumerate(zip(st.session_state.history, st.session_state.steps)):
        st.session_state.current_step = idx
        
        # Update progress
        progress.progress((idx + 1) / len(st.session_state.history))

        # Create bar plot
        fig, ax = plt.subplots(figsize=(10, 6))
        bars = ax.bar(range(len(state)), state, color='skyblue', edgecolor='black')
        
        # Highlight current key and sorted portion
        if 'comparing' in step:
            bars[step['comparing']].set_color('orange')
        if 'inserted' in step:
            bars[step['inserted']].set_color('green')
        if 'inserting' in step:
            bars[step['inserting']].set_color('red')
        
        # Color the sorted portion (left part)
        sorted_until = 0
        for s in st.session_state.steps[:idx]:
            if 'inserted' in s:
                sorted_until = max(sorted_until, s['inserted'] + 1)
        for i in range(sorted_until):
            bars[i].set_color('#90EE90')  # light green for sorted
        
        ax.set_title("Insertion Sort Visualization", fontsize=16, fontweight='bold')
        ax.set_xlabel("Index")
        ax.set_ylabel("Value")
        ax.set_xlim(-0.5, len(state) - 0.5)
        ax.set_ylim(0, max(state) * 1.1)

        # Add step description
        action_text = step.get('action', 'Processing...')
        info.markdown(f"**Step {idx + 1}/{len(st.session_state.history)}**: {action_text}")

        chart.pyplot(fig)
        plt.close(fig)
        
        time.sleep(speed)
    
    st.success("Insertion Sort Completed!")
    st.balloons()
    st.session_state.running = False

# Display current array
st.subheader("Current Array")
cols = st.columns(len(array))
for i, val in enumerate(array):
    cols[i].metric(label=f"Index {i}", value=val)

# Show algorithm explanation
with st.expander("How Insertion Sort Works"):
    st.markdown("""
    ### Insertion Sort Algorithm
    
    Insertion Sort builds the final sorted array one item at a time.
    
    **Steps:**
    1. Start with the second element (index 1)
    2. Compare it with elements on its left
    3. Shift larger elements one position to the right
    4. Insert the element in its correct position
    5. Repeat until the entire array is sorted
    
    **Time Complexity:**
    - Best Case: O(n) → Already sorted
    - Average/Worst Case: O(n²)
    
    **Great for:** Small datasets or nearly sorted arrays!
    """)

# Reset button
if st.button("Reset"):
    st.session_state.array = random.sample(range(1, 100), array_size)
    st.session_state.history = []
    st.session_state.steps = []
    st.session_state.running = False
    st.rerun()