import streamlit as st
import numpy as np
import plotly.graph_objects as go
from plotly.subplots import make_subplots
from kinematics_visualizer import Motion1D, Motion2D, KinematicsVisualizer

# Configure Streamlit page
st.set_page_config(
    page_title="Physics Tutorial: Kinematics",
    page_icon="🚀",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS for better styling
st.markdown("""
<style>
    .main-header {
        font-size: 2.5rem;
        font-weight: bold;
        color: #1f77b4;
        text-align: center;
        margin-bottom: 2rem;
    }
    .section-header {
        font-size: 1.5rem;
        font-weight: bold;
        color: #2e7d32;
        margin-top: 2rem;
        margin-bottom: 1rem;
    }
    .physics-equation {
        background-color: #2e2e2e;
        color: white;
        padding: 1rem;
        border-radius: 10px;
        border-left: 5px solid #1f77b4;
        margin: 1rem 0;
    }
    .physics-equation h4 {
        color: #87ceeb;
        margin-bottom: 1rem;
    }
    .physics-equation ul {
        color: white;
    }
    .physics-equation li {
        margin-bottom: 0.5rem;
    }
    .physics-equation strong {
        color: #ffd700;
    }
    .parameter-box {
        background-color: #fff3e0;
        padding: 1rem;
        border-radius: 10px;
        margin: 1rem 0;
    }
</style>
""", unsafe_allow_html=True)

def main():
    st.markdown('<div class="main-header">🚀 Physics Tutorial: Kinematics</div>', unsafe_allow_html=True)
    
    # Sidebar for navigation
    st.sidebar.title("📚 Select Tutorial")
    tutorial_type = st.sidebar.radio(
        "Choose a physics concept:",
        ["1D Motion", "2D Projectile Motion"]
    )
    
    if tutorial_type == "1D Motion":
        show_1d_motion()
    elif tutorial_type == "2D Projectile Motion":
        show_2d_motion()

def create_1d_motion_plot_plotly(motion: Motion1D, duration: float, title: str):
    """Create 1D motion plots using Plotly"""
    # Generate time arrays
    t, x, v, a = motion.time_arrays(duration, dt=0.01)
    
    # Create subplots
    fig = make_subplots(
        rows=3, cols=1,
        subplot_titles=('Position vs Time', 'Velocity vs Time', 'Acceleration vs Time'),
        vertical_spacing=0.08,
        shared_xaxes=True
    )
    
    # Position plot
    fig.add_trace(
        go.Scatter(x=t, y=x, mode='lines', name='Position', 
                  line=dict(color='blue', width=3),
                  hovertemplate='Time: %{x:.1f} s<br>Position: %{y:.1f} m<extra></extra>'),
        row=1, col=1
    )
    
    # Velocity plot
    fig.add_trace(
        go.Scatter(x=t, y=v, mode='lines', name='Velocity', 
                  line=dict(color='red', width=3),
                  hovertemplate='Time: %{x:.1f} s<br>Velocity: %{y:.1f} m/s<extra></extra>'),
        row=2, col=1
    )
    
    # Acceleration plot
    fig.add_trace(
        go.Scatter(x=t, y=a, mode='lines', name='Acceleration', 
                  line=dict(color='green', width=3),
                  hovertemplate='Time: %{x:.1f} s<br>Acceleration: %{y:.1f} m/s²<extra></extra>'),
        row=3, col=1
    )
    
    # Update layout
    fig.update_layout(
        title=dict(text=title, font=dict(size=18)),
        showlegend=False,
        height=800,
        template='plotly_white'
    )
    
    # Update y-axis labels
    fig.update_yaxes(title_text="Position (m)", row=1, col=1)
    fig.update_yaxes(title_text="Velocity (m/s)", row=2, col=1)
    fig.update_yaxes(title_text="Acceleration (m/s²)", row=3, col=1)
    fig.update_xaxes(title_text="Time (s)", row=3, col=1)
    
    # Add grid
    fig.update_xaxes(showgrid=True, gridwidth=1, gridcolor='rgba(128,128,128,0.3)')
    fig.update_yaxes(showgrid=True, gridwidth=1, gridcolor='rgba(128,128,128,0.3)')
    
    return fig

def create_2d_trajectory_plot_plotly(motion: Motion2D, title: str):
    """Create 2D trajectory plot using Plotly"""
    flight_time = motion.calculate_flight_time()
    data = motion.trajectory_data(flight_time)
    
    fig = go.Figure()
    
    # Add trajectory line
    fig.add_trace(go.Scatter(
        x=data['x'], 
        y=data['y'],
        mode='lines',
        name='Trajectory',
        line=dict(color='blue', width=4),
        hovertemplate='X: %{x:.1f} m<br>Y: %{y:.1f} m<extra></extra>'
    ))
    
    # Add launch point
    fig.add_trace(go.Scatter(
        x=[motion.launch_x],
        y=[motion.launch_height],
        mode='markers',
        name='Launch Point',
        marker=dict(color='green', size=12, symbol='circle'),
        hovertemplate='Launch<br>X: %{x:.1f} m<br>Y: %{y:.1f} m<extra></extra>'
    ))
    
    # Add landing point
    if len(data['x']) > 0:
        fig.add_trace(go.Scatter(
            x=[data['x'][-1]],
            y=[data['y'][-1]],
            mode='markers',
            name='Landing Point',
            marker=dict(color='red', size=12, symbol='square'),
            hovertemplate='Landing<br>X: %{x:.1f} m<br>Y: %{y:.1f} m<extra></extra>'
        ))
    
    # Add maximum height point
    if len(data['y']) > 0:
        max_height_idx = np.argmax(data['y'])
        if max_height_idx > 0:
            fig.add_trace(go.Scatter(
                x=[data['x'][max_height_idx]],
                y=[data['y'][max_height_idx]],
                mode='markers',
                name='Max Height',
                marker=dict(color='orange', size=10, symbol='triangle-up'),
                hovertemplate='Max Height<br>X: %{x:.1f} m<br>Y: %{y:.1f} m<extra></extra>'
            ))
    
    # Update layout
    fig.update_layout(
        title=dict(text=title, font=dict(size=18)),
        xaxis_title="Horizontal Position (m)",
        yaxis_title="Vertical Position (m)",
        showlegend=True,
        hovermode='closest',
        template='plotly_white',
        height=600,
        margin=dict(l=50, r=50, t=80, b=50)
    )
    
    # Set axis ranges
    if len(data['x']) > 0 and len(data['y']) > 0:
        fig.update_xaxes(range=[0, max(data['x']) * 1.1])
        fig.update_yaxes(range=[0, max(data['y']) * 1.2])
    else:
        fig.update_xaxes(range=[0, 10])
        fig.update_yaxes(range=[0, 10])
    
    # Add grid
    fig.update_xaxes(showgrid=True, gridwidth=1, gridcolor='rgba(128,128,128,0.3)')
    fig.update_yaxes(showgrid=True, gridwidth=1, gridcolor='rgba(128,128,128,0.3)')
    
    return fig

def show_1d_motion():
    st.markdown('<div class="section-header">📏 One-Dimensional Motion</div>', unsafe_allow_html=True)
    
    # Physics equations display
    with st.expander("📖 Physics Equations (Click to expand)", expanded=False):
        st.markdown("""
        <div class="physics-equation">
        <h4>Kinematic Equations for Constant Acceleration:</h4>
        <ul>
            <li><strong>Position:</strong> x(t) = x₀ + v₀t + ½at²</li>
            <li><strong>Velocity:</strong> v(t) = v₀ + at</li>
            <li><strong>Acceleration:</strong> a(t) = constant</li>
        </ul>
        <p><strong>Where:</strong></p>
        <ul>
            <li>x₀ = initial position (m)</li>
            <li>v₀ = initial velocity (m/s)</li>
            <li>a = acceleration (m/s²)</li>
            <li>t = time (s)</li>
        </ul>
        </div>
        """, unsafe_allow_html=True)
    
    # Create two columns for controls and results
    col1, col2 = st.columns([1, 2])
    
    with col1:
        st.markdown('<div class="section-header">🎛️ Control Parameters</div>', unsafe_allow_html=True)
        
        # Parameter sliders
        initial_pos = st.slider(
            "Initial Position (m)", 
            min_value=-50.0, max_value=50.0, value=0.0, step=1.0,
            help="Starting position of the object"
        )
        
        initial_vel = st.slider(
            "Initial Velocity (m/s)", 
            min_value=-30.0, max_value=30.0, value=5.0, step=1.0,
            help="Starting velocity of the object"
        )
        
        acceleration = st.slider(
            "Acceleration (m/s²)", 
            min_value=-15.0, max_value=15.0, value=2.0, step=0.5,
            help="Constant acceleration (positive = speeding up in positive direction)"
        )
        
        duration = st.slider(
            "Simulation Duration (s)", 
            min_value=1.0, max_value=20.0, value=10.0, step=0.5,
            help="How long to run the simulation"
        )
        
        # Display current parameters
        st.markdown('<div class="parameter-box">', unsafe_allow_html=True)
        st.markdown("**Current Parameters:**")
        st.write(f"• Initial Position: {initial_pos:.1f} m")
        st.write(f"• Initial Velocity: {initial_vel:.1f} m/s")
        st.write(f"• Acceleration: {acceleration:.1f} m/s²")
        st.write(f"• Duration: {duration:.1f} s")
        
        # Calculate final values
        final_pos = initial_pos + initial_vel * duration + 0.5 * acceleration * duration**2
        final_vel = initial_vel + acceleration * duration
        
        st.markdown("**Final Values:**")
        st.write(f"• Final Position: {final_pos:.1f} m")
        st.write(f"• Final Velocity: {final_vel:.1f} m/s")
        st.markdown('</div>', unsafe_allow_html=True)
    
    with col2:
        # Create and display motion
        motion = Motion1D(
            initial_position=initial_pos,
            initial_velocity=initial_vel,
            acceleration=acceleration
        )
        
        # Generate title based on motion type
        if acceleration > 0:
            motion_type = "Accelerating Motion"
        elif acceleration < 0:
            motion_type = "Decelerating Motion"
        else:
            motion_type = "Constant Velocity Motion"
        
        # Create and display Plotly chart
        fig = create_1d_motion_plot_plotly(motion, duration, motion_type)
        st.plotly_chart(fig, use_container_width=True)

def show_2d_motion():
    st.markdown('<div class="section-header">🎯 Two-Dimensional Projectile Motion</div>', unsafe_allow_html=True)
    
    # Physics equations display (updated to include sphere physics)
    with st.expander("📖 Physics Equations (Click to expand)", expanded=False):
        st.markdown("""
        <div class="physics-equation">
        <h4>Projectile Motion Equations:</h4>
        <p><strong>Point Mass Model:</strong></p>
        <ul>
            <li><strong>No Air Resistance:</strong> Standard kinematic equations</li>
            <li><strong>With Air Resistance:</strong> Linear drag model (drag ∝ velocity)</li>
        </ul>
        <p><strong>Sphere Model (more realistic):</strong></p>
        <ul>
            <li><strong>Mass:</strong> m = ρ × (4/3)π × r³</li>
            <li><strong>Cross-sectional Area:</strong> A = π × r²</li>
            <li><strong>Drag Force:</strong> F<sub>drag</sub> = ½ρ<sub>air</sub>C<sub>d</sub>Av²</li>
            <li><strong>Terminal Velocity:</strong> v<sub>t</sub> = √(2mg / ρ<sub>air</sub>C<sub>d</sub>A)</li>
        </ul>
        <p><strong>Where:</strong></p>
        <ul>
            <li>ρ = sphere density (kg/m³)</li>
            <li>r = sphere radius (m)</li>
            <li>C<sub>d</sub> = aerodynamic drag coefficient</li>
            <li>ρ<sub>air</sub> = air density (~1.225 kg/m³)</li>
        </ul>
        </div>
        """, unsafe_allow_html=True)
    
    # Create two columns
    col1, col2 = st.columns([1, 2])
    
    with col1:
        st.markdown('<div class="section-header">🎛️ Launch Parameters</div>', unsafe_allow_html=True)
        
        # Projectile Model Toggle
        st.markdown("### 🎯 Projectile Model")
        is_sphere = st.toggle(
            "Model as Sphere",
            value=False,
            help="Toggle between point mass and realistic sphere with physical dimensions"
        )
        
        # Sphere properties (only show when sphere model is enabled)
        sphere_radius = 0.037  # Default baseball
        sphere_density = 700
        sphere_drag_coeff = 0.47
        
        if is_sphere:
            # Sphere presets
            st.markdown("**Quick Sphere Presets:**")
            presets = Motion2D.get_sphere_presets()
            preset_cols = st.columns(3)
            
            selected_preset = None
            for i, (key, preset) in enumerate(presets.items()):
                with preset_cols[i % 3]:
                    if st.button(preset['name'], key=f"sphere_{key}"):
                        selected_preset = preset
            
            # Apply preset if selected
            if selected_preset:
                sphere_radius = selected_preset['radius']
                sphere_density = selected_preset['density']
                sphere_drag_coeff = selected_preset['drag_coeff']
            
            st.markdown("**Custom Sphere Properties:**")
            
            # Sphere radius slider
            sphere_radius = st.slider(
                "Sphere Radius (mm)", 
                min_value=10.0, max_value=150.0, 
                value=sphere_radius*1000, step=1.0,
                help="Radius of the sphere in millimeters"
            ) / 1000  # Convert back to meters
            
            # Sphere density slider
            sphere_density = st.slider(
                "Sphere Density (kg/m³)", 
                min_value=50, max_value=2000, 
                value=int(sphere_density), step=10,
                help="Material density - affects mass and terminal velocity"
            )
            
            # Drag coefficient slider
            sphere_drag_coeff = st.slider(
                "Aerodynamic Drag Coefficient", 
                min_value=0.1, max_value=1.0, 
                value=sphere_drag_coeff, step=0.01,
                help="0.24 (golf ball), 0.35 (baseball), 0.47 (smooth sphere), 0.51 (tennis ball)"
            )
            
            # Display calculated properties
            temp_motion = Motion2D(
                launch_speed=25, launch_angle=45, 
                is_sphere=True, sphere_radius=sphere_radius,
                sphere_density=sphere_density, sphere_drag_coeff=sphere_drag_coeff
            )
            sphere_info = temp_motion.get_sphere_info()
            
            st.info(f"""
            **Calculated Properties:**
            • Diameter: {sphere_info['diameter_mm']:.1f} mm
            • Mass: {sphere_info['mass_g']:.1f} g
            • Cross-section: {sphere_info['cross_section_cm2']:.1f} cm²
            • Volume: {sphere_info['volume_cm3']:.1f} cm³
            """)
        
        # Air Resistance Toggle
        st.markdown("### 🌬️ Air Resistance")
        air_resistance_enabled = st.toggle(
            "Enable Air Resistance",
            value=False,
            help="Toggle air resistance on/off to see the difference in trajectory"
        )
        
        # Show drag info based on model
        if air_resistance_enabled and not is_sphere:
            drag_coeff = st.slider(
                "Point Mass Drag Coefficient", 
                min_value=0.01, max_value=0.5, value=0.1, step=0.01,
                help="Simple linear drag coefficient for point mass model"
            )
        elif air_resistance_enabled and is_sphere:
            st.info("🌬️ **Sphere Model**: Air resistance calculated from physical properties!")
            drag_coeff = 0.1  # Not used for sphere model
        else:
            drag_coeff = 0.1
        
        st.markdown("---")
        
        # Initialize default values
        default_speed = 25.0
        default_angle = 45.0
        default_height = 0.0
        default_gravity = 9.81
        
        # Handle preset button clicks
        st.markdown("### 🎯 Launch Presets")
        preset_buttons_col1, preset_buttons_col2 = st.columns(2)
        with preset_buttons_col1:
            if st.button("🏀 Basketball Shot"):
                st.session_state.speed_preset = 15.0
                st.session_state.angle_preset = 50.0
                st.session_state.height_preset = 2.0
                st.session_state.gravity_preset = 9.81
        
        with preset_buttons_col2:
            if st.button("🚀 Rocket Launch"):
                st.session_state.speed_preset = 40.0
                st.session_state.angle_preset = 75.0
                st.session_state.height_preset = 0.0
                st.session_state.gravity_preset = 9.81
        
        # Use preset values if they exist, otherwise use defaults
        speed_value = st.session_state.get('speed_preset', default_speed)
        angle_value = st.session_state.get('angle_preset', default_angle)
        height_value = st.session_state.get('height_preset', default_height)
        gravity_value = st.session_state.get('gravity_preset', default_gravity)
        
        st.markdown("### ⚙️ Physics Parameters")
        
        # Parameter sliders
        launch_speed = st.slider(
            "Launch Speed (m/s)", 
            min_value=5.0, max_value=150.0, value=speed_value, step=1.0,
            help="Initial speed of the projectile"
        )
        
        launch_angle = st.slider(
            "Launch Angle (degrees)", 
            min_value=0.0, max_value=90.0, value=angle_value, step=5.0,
            help="Angle above horizontal"
        )
        
        launch_height = st.slider(
            "Launch Height (m)", 
            min_value=0.0, max_value=50.0, value=height_value, step=1.0,
            help="Height above ground level"
        )
        
        gravity = st.slider(
            "Gravity (m/s²)", 
            min_value=0.0, max_value=98.1, value=gravity_value, step=9.81,
            help="Acceleration due to gravity; 1g per step. Earth = 1g = 9.81 m/s"
        )
        
        # Clear preset values when sliders are moved
        if (launch_speed != speed_value or launch_angle != angle_value or 
            launch_height != height_value or gravity != gravity_value):
            for key in ['speed_preset', 'angle_preset', 'height_preset', 'gravity_preset']:
                if key in st.session_state:
                    del st.session_state[key]
        
        # Create motion object with all parameters
        motion = Motion2D(
            launch_speed=launch_speed,
            launch_angle=launch_angle,
            launch_height=launch_height,
            gravity=gravity,
            air_resistance=air_resistance_enabled,
            drag_coefficient=drag_coeff,
            is_sphere=is_sphere,
            sphere_radius=sphere_radius,
            sphere_density=sphere_density,
            sphere_drag_coeff=sphere_drag_coeff
        )
        
        info = motion.get_launch_info()
        
        # Display comprehensive analysis
        st.markdown('<div class="parameter-box">', unsafe_allow_html=True)
        st.markdown("**Launch Analysis:**")
        st.write(f"• Model: {'Sphere' if is_sphere else 'Point Mass'}")
        st.write(f"• Launch Speed: {info['launch_speed']:.1f} m/s")
        st.write(f"• Launch Angle: {info['launch_angle']:.1f}°")
        st.write(f"• Initial Vₓ: {info['initial_velocity_x']:.1f} m/s")
        st.write(f"• Initial Vᵧ: {info['initial_velocity_y']:.1f} m/s")
        
        if is_sphere and 'mass_g' in info:
            st.write(f"• Mass: {info['mass_g']:.1f} g")
            st.write(f"• Diameter: {info['diameter_mm']:.1f} mm")
            if air_resistance_enabled and 'terminal_velocity' in info:
                st.write(f"• Terminal Velocity: {info['terminal_velocity']:.1f} m/s")
        
        st.markdown("**Trajectory Results:**")
        st.write(f"• Flight Time: {info['flight_time']:.2f} s")
        st.write(f"• Range: {info['range']:.1f} m")
        st.write(f"• Max Height: {info['max_height']:.1f} m")
        
        # Physics insights
        if is_sphere and air_resistance_enabled and 'mass_g' in info:
            st.markdown("**Physics Insights:**")
            if info['mass_g'] > 200:
                st.write("🔹 Heavy object: Less affected by air resistance")
            elif info['mass_g'] < 50:
                st.write("🔹 Light object: Significantly affected by air resistance")
            
            if info['diameter_mm'] > 100:
                st.write("🔹 Large cross-section: More air resistance")
            elif info['diameter_mm'] < 50:
                st.write("🔹 Small cross-section: Less air resistance")
        
        st.markdown('</div>', unsafe_allow_html=True)
        
        # Reset button
        if st.button("🔄 Reset to Defaults"):
            for key in ['speed_preset', 'angle_preset', 'height_preset', 'gravity_preset']:
                if key in st.session_state:
                    del st.session_state[key]
            st.rerun()
    
    with col2:
        # Create and display trajectory with model info
        model_info = f"({'Sphere' if is_sphere else 'Point Mass'})"
        air_info = " (with Air Resistance)" if air_resistance_enabled else " (No Air Resistance)"
        trajectory_title = f"Projectile Motion - {launch_angle:.0f}° Launch {model_info}{air_info}"
        
        # Create and display main trajectory plot
        fig = create_2d_trajectory_plot_plotly(motion, trajectory_title)
        st.plotly_chart(fig, use_container_width=True)
        
        # Show model comparison if using sphere
        if is_sphere and 'mass_g' in info:
            st.markdown("### 📊 Sphere vs Point Mass Comparison")
            
            # Create comparison plot using Plotly
            fig_comp = go.Figure()
            
            # Plot sphere model
            data_sphere = motion.trajectory_data(motion.calculate_flight_time())
            fig_comp.add_trace(go.Scatter(
                x=data_sphere['x'], 
                y=data_sphere['y'],
                mode='lines',
                name=f'Sphere Model ({info["mass_g"]:.0f}g)',
                line=dict(color='red', width=3),
                hovertemplate='<b>Sphere</b><br>X: %{x:.1f} m<br>Y: %{y:.1f} m<extra></extra>'
            ))
            
            # Plot equivalent point mass
            motion_point = Motion2D(
                launch_speed=launch_speed, launch_angle=launch_angle, 
                launch_height=launch_height, gravity=gravity, 
                air_resistance=air_resistance_enabled, drag_coefficient=drag_coeff,
                is_sphere=False
            )
            data_point = motion_point.trajectory_data(motion_point.calculate_flight_time())
            fig_comp.add_trace(go.Scatter(
                x=data_point['x'], 
                y=data_point['y'],
                mode='lines',
                name='Point Mass Model',
                line=dict(color='blue', width=2, dash='dash'),
                hovertemplate='<b>Point Mass</b><br>X: %{x:.1f} m<br>Y: %{y:.1f} m<extra></extra>'
            ))
            
            # Update layout for comparison plot
            fig_comp.update_layout(
                title="Sphere Model vs Point Mass Model",
                xaxis_title="Horizontal Position (m)",
                yaxis_title="Vertical Position (m)",
                showlegend=True,
                template='plotly_white',
                height=400
            )
            
            fig_comp.update_yaxes(range=[0, None])
            fig_comp.update_xaxes(showgrid=True, gridwidth=1, gridcolor='rgba(128,128,128,0.3)')
            fig_comp.update_yaxes(showgrid=True, gridwidth=1, gridcolor='rgba(128,128,128,0.3)')
            
            st.plotly_chart(fig_comp, use_container_width=True)

if __name__ == "__main__":
    main()