src/streamlit_app.py

"""
Traffic Accident Reconstruction System
======================================
Main Streamlit Application
Huawei AI Innovation Challenge 2026

This system uses MindSpore AI to analyze traffic accidents
and generate probable scenarios with 2D simulation.
"""

import streamlit as st
import sys
from pathlib import Path

# Add project root to path
sys.path.insert(0, str(Path(__file__).parent))

from config import STREAMLIT_CONFIG, CASE_STUDY_LOCATION, COLORS
from ui.components import render_sidebar, render_header, render_footer
from ui.map_viewer import render_map_section
from ui.vehicle_input import render_vehicle_input
from ui.party_input import render_party_input, render_evidence_upload, render_party_summary
from ui.results_display import render_results

# ============================================================
# PAGE CONFIGURATION
# ============================================================

st.set_page_config(
    page_title=STREAMLIT_CONFIG["page_title"],
    page_icon=STREAMLIT_CONFIG["page_icon"],
    layout=STREAMLIT_CONFIG["layout"],
    initial_sidebar_state=STREAMLIT_CONFIG["initial_sidebar_state"]
)

# ============================================================
# CUSTOM CSS - CrashLens Modern Design
# ============================================================

st.markdown("""
<style>
    /* Layout - Dark theme with modern spacing */
    .block-container { 
        padding-top: 1.4rem; 
        padding-bottom: 2.5rem; 
        max-width: 1200px; 
    }
    
    [data-testid="stSidebar"] { 
        min-width: 340px; 
        max-width: 340px;
        background: linear-gradient(180deg, #0e1117 0%, #1a1f2e 100%);
    }
    
    /* Main container with gradient background */
    .main {
        background: linear-gradient(135deg, #0a0e27 0%, #1a1f3a 50%, #0f1419 100%);
        color: rgba(255, 255, 255, 0.95);
    }
    
    /* Header styling - Modern minimal */
    .main-header {
        background: linear-gradient(135deg, rgba(30, 58, 95, 0.3) 0%, rgba(45, 90, 135, 0.2) 100%);
        padding: 2rem;
        border-radius: 16px;
        border: 1px solid rgba(255, 255, 255, 0.12);
        color: white;
        margin-bottom: 2rem;
        text-align: center;
        box-shadow: 0 10px 40px rgba(0, 0, 0, 0.3);
    }
    
    .main-header h1 {
        margin: 0;
        font-size: 2.8rem;
        font-weight: 800;
        letter-spacing: 0.5px;
        background: linear-gradient(135deg, #60a5fa 0%, #3b82f6 100%);
        -webkit-background-clip: text;
        -webkit-text-fill-color: transparent;
    }
    
    .main-header p {
        margin: 0.8rem 0 0 0;
        opacity: 0.75;
        font-size: 1.05rem;
        color: rgba(255, 255, 255, 0.72);
    }
    
    /* Card styling - Glass morphism effect */
    .info-card {
        background: rgba(255, 255, 255, 0.04);
        backdrop-filter: blur(10px);
        padding: 1.5rem;
        border-radius: 16px;
        border: 1px solid rgba(255, 255, 255, 0.12);
        margin: 1rem 0;
        box-shadow: 0 10px 30px rgba(0, 0, 0, 0.15);
    }
    
    /* Step indicator - Modern pills */
    .step-indicator {
        display: flex;
        justify-content: space-between;
        margin: 2rem 0;
        gap: 0.5rem;
    }
    
    .step {
        flex: 1;
        text-align: center;
        padding: 1rem;
        background: rgba(255, 255, 255, 0.05);
        border: 1px solid rgba(255, 255, 255, 0.1);
        margin: 0;
        border-radius: 12px;
        position: relative;
        transition: all 0.3s ease;
    }
    
    .step.active {
        background: linear-gradient(135deg, #3b82f6 0%, #2563eb 100%);
        border-color: #3b82f6;
        color: white;
        box-shadow: 0 8px 24px rgba(59, 130, 246, 0.3);
    }
    
    .step.completed {
        background: linear-gradient(135deg, #10b981 0%, #059669 100%);
        border-color: #10b981;
        color: white;
    }
    
    /* Vehicle cards - Dark theme */
    .vehicle-card {
        background: rgba(255, 255, 255, 0.06);
        border: 1px solid rgba(255, 255, 255, 0.15);
        border-radius: 12px;
        padding: 1.5rem;
        margin: 1rem 0;
    }
    
    .vehicle-card.vehicle-1 {
        border-color: rgba(255, 75, 75, 0.4);
        background: rgba(255, 75, 75, 0.08);
    }
    
    .vehicle-card.vehicle-2 {
        border-color: rgba(75, 123, 255, 0.4);
        background: rgba(75, 123, 255, 0.08);
    }
    
    /* Results styling - Modern cards */
    .scenario-card {
        background: rgba(255, 255, 255, 0.06);
        border-radius: 12px;
        padding: 1.5rem;
        margin: 1rem 0;
        box-shadow: 0 4px 12px rgba(0,0,0,0.2);
        border: 1px solid rgba(255, 255, 255, 0.1);
    }
    
    .probability-high {
        color: #10b981;
        font-weight: bold;
    }
    
    .probability-medium {
        color: #f59e0b;
        font-weight: bold;
    }
    
    .probability-low {
        color: #ef4444;
        font-weight: bold;
    }
    
    /* Button styling - Modern gradient */
    .stButton > button {
        width: 100%;
        background: linear-gradient(135deg, #3b82f6 0%, #2563eb 100%);
        color: white;
        border: none;
        padding: 0.85rem 1.5rem;
        border-radius: 12px;
        font-weight: 600;
        font-size: 1rem;
        transition: all 0.3s ease;
        box-shadow: 0 4px 12px rgba(59, 130, 246, 0.3);
    }
    
    .stButton > button:hover {
        background: linear-gradient(135deg, #2563eb 0%, #1d4ed8 100%);
        box-shadow: 0 6px 16px rgba(59, 130, 246, 0.4);
        transform: translateY(-2px);
    }
    
    /* Input fields - Dark theme with better contrast */
    .stTextInput > div > div > input,
    .stTextArea > div > div > textarea,
    .stSelectbox > div > div > select,
    .stNumberInput > div > div > input {
        background: rgba(255, 255, 255, 0.08) !important;
        border: 1px solid rgba(255, 255, 255, 0.15) !important;
        border-radius: 10px !important;
        color: white !important;
        padding: 0.75rem !important;
    }
    
    .stTextInput > div > div > input:focus,
    .stTextArea > div > div > textarea:focus,
    .stSelectbox > div > div > select:focus,
    .stNumberInput > div > div > input:focus {
        border-color: #3b82f6 !important;
        box-shadow: 0 0 0 3px rgba(59, 130, 246, 0.15) !important;
    }
    
    /* Labels - Better visibility */
    label, .stTextInput > label, .stTextArea > label, .stSelectbox > label, .stSlider > label, .stNumberInput > label {
        color: rgba(255, 255, 255, 0.85) !important;
        font-weight: 500 !important;
        font-size: 0.95rem !important;
    }
    
    /* Metric cards - Modern design */
    [data-testid="stMetricValue"] {
        font-size: 2.2rem;
        font-weight: 700;
        color: white;
    }
    
    [data-testid="stMetricLabel"] {
        color: rgba(255, 255, 255, 0.7);
        font-size: 0.95rem;
        font-weight: 500;
    }
    
    /* Tabs styling - Modern */
    .stTabs [data-baseweb="tab-list"] {
        gap: 8px;
        background: rgba(255, 255, 255, 0.03);
        padding: 0.5rem;
        border-radius: 12px;
    }
    
    .stTabs [data-baseweb="tab"] {
        background: transparent;
        border-radius: 8px;
        color: rgba(255, 255, 255, 0.7);
        font-weight: 500;
        padding: 0.75rem 1.5rem;
        border: 1px solid transparent;
    }
    
    .stTabs [aria-selected="true"] {
        background: linear-gradient(135deg, #3b82f6 0%, #2563eb 100%);
        color: white !important;
        border-color: #3b82f6;
    }
    
    /* Messages - Modern alerts */
    .stSuccess, .element-container:has(.stSuccess) {
        background: rgba(16, 185, 129, 0.15) !important;
        border: 1px solid rgba(16, 185, 129, 0.3) !important;
        border-radius: 12px !important;
        padding: 1rem !important;
    }
    
    .stSuccess [data-testid="stMarkdownContainer"] p {
        color: #10b981 !important;
    }
    
    .stError {
        background: rgba(239, 68, 68, 0.15) !important;
        border: 1px solid rgba(239, 68, 68, 0.3) !important;
        border-radius: 12px !important;
    }
    
    .stWarning {
        background: rgba(245, 158, 11, 0.15) !important;
        border: 1px solid rgba(245, 158, 11, 0.3) !important;
        border-radius: 12px !important;
    }
    
    .stInfo {
        background: rgba(59, 130, 246, 0.15) !important;
        border: 1px solid rgba(59, 130, 246, 0.3) !important;
        border-radius: 12px !important;
    }
    
    /* Sidebar - Better styling */
    [data-testid="stSidebar"] h1, [data-testid="stSidebar"] h2, [data-testid="stSidebar"] h3 {
        color: white !important;
    }
    
    [data-testid="stSidebar"] p, [data-testid="stSidebar"] span {
        color: rgba(255, 255, 255, 0.85) !important;
    }
    
    /* Expander - Modern style */
    .streamlit-expanderHeader {
        background: rgba(255, 255, 255, 0.05) !important;
        border-radius: 10px !important;
        border: 1px solid rgba(255, 255, 255, 0.1) !important;
        color: white !important;
        font-weight: 500 !important;
    }
    
    .streamlit-expanderContent {
        background: rgba(255, 255, 255, 0.02) !important;
        border: 1px solid rgba(255, 255, 255, 0.08) !important;
        border-radius: 0 0 10px 10px !important;
    }
    
    /* Checkbox and Radio - Better visibility */
    .stCheckbox label, .stRadio label {
        color: rgba(255, 255, 255, 0.85) !important;
    }
    
    /* Slider - Modern */
    .stSlider [data-baseweb="slider"] {
        background: rgba(59, 130, 246, 0.2);
    }
    
    /* Headings - Better hierarchy */
    h1, h2, h3, h4, h5, h6 {
        color: white !important;
    }
    
    h1 {
        font-weight: 800 !important;
        letter-spacing: -0.5px;
    }
    
    h2 {
        font-weight: 700 !important;
        letter-spacing: -0.3px;
    }
    
    h3 {
        font-weight: 600 !important;
    }
    
    /* Markdown text - Better readability */
    p, span, div {
        color: rgba(255, 255, 255, 0.9);
    }
    
    /* Code blocks - Dark theme */
    code {
        background: rgba(255, 255, 255, 0.08) !important;
        color: #60a5fa !important;
        padding: 0.2rem 0.4rem !important;
        border-radius: 4px !important;
    }
    
    /* Dataframe - Dark theme */
    [data-testid="stDataFrame"] {
        background: rgba(255, 255, 255, 0.05);
        border-radius: 12px;
        border: 1px solid rgba(255, 255, 255, 0.1);
    }
    
    /* Hide Streamlit branding */
    #MainMenu {visibility: hidden;}
    footer {visibility: hidden;}
    
    /* Progress bar */
    .stProgress > div > div {
        background: linear-gradient(90deg, #3b82f6 0%, #8b5cf6 100%);
        border-radius: 8px;
    }
</style>
""", unsafe_allow_html=True)

# ============================================================
# SESSION STATE INITIALIZATION
# ============================================================

def init_session_state():
    """Initialize all session state variables."""
    
    # Application state
    if 'current_step' not in st.session_state:
        st.session_state.current_step = 1
    
    # Accident information
    if 'accident_info' not in st.session_state:
        st.session_state.accident_info = {
            'location': CASE_STUDY_LOCATION.copy(),
            'datetime': None,
            'road_type': 'roundabout',
            'weather': 'clear',
            'road_condition': 'dry',
            'visibility': 1.0,
            'lighting': 'daylight',
            'notes': ''
        }
    
    # Vehicle 1 data (for analysis)
    if 'vehicle_1' not in st.session_state:
        st.session_state.vehicle_1 = {
            'type': 'sedan',
            'speed': 50,
            'direction': 'north',
            'action': 'entering_roundabout',
            'braking': False,
            'signaling': False,
            'lights_on': True,
            'horn_used': False,
            'path': [],
            'description': ''
        }
    
    # Vehicle 2 data (for analysis)
    if 'vehicle_2' not in st.session_state:
        st.session_state.vehicle_2 = {
            'type': 'sedan',
            'speed': 50,
            'direction': 'east',
            'action': 'going_straight',
            'braking': False,
            'signaling': False,
            'lights_on': True,
            'horn_used': False,
            'path': [],
            'description': ''
        }
    
    # Party 1 data (driver details)
    if 'party_1' not in st.session_state:
        st.session_state.party_1 = {
            'full_name': '',
            'id_iqama': '',
            'phone': '',
            'role': 'Driver',
            'vehicle_make_model': '',
            'plate_number': '',
            'insurance': '',
            'damage_notes': '',
            'statement': ''
        }
    
    # Party 2 data (driver details)
    if 'party_2' not in st.session_state:
        st.session_state.party_2 = {
            'full_name': '',
            'id_iqama': '',
            'phone': '',
            'role': 'Driver',
            'vehicle_make_model': '',
            'plate_number': '',
            'insurance': '',
            'damage_notes': '',
            'statement': ''
        }
    
    # Evidence photos
    if 'evidence_photos' not in st.session_state:
        st.session_state.evidence_photos = []
    
    # Analysis results
    if 'analysis_results' not in st.session_state:
        st.session_state.analysis_results = None
    
    # Generated scenarios
    if 'scenarios' not in st.session_state:
        st.session_state.scenarios = []
    
    # Map data
    if 'map_data' not in st.session_state:
        st.session_state.map_data = None
    
    # Analysis ready flag
    if 'analysis_ready' not in st.session_state:
        st.session_state.analysis_ready = False

# ============================================================
# MAIN APPLICATION
# ============================================================

def main():
    """Main application entry point."""
    
    # Initialize session state
    init_session_state()
    
    # Render modern header - clean without logo
    col1, col2, col3 = st.columns([1, 6, 1])
    with col1:
        st.markdown("## 🚗")
    with col2:
        st.markdown("""
        <div style='padding-top: 0.5rem;'>
            <h1 style='margin: 0; font-size: 2.5rem; font-weight: 800; background: linear-gradient(135deg, #60a5fa 0%, #3b82f6 100%); -webkit-background-clip: text; -webkit-text-fill-color: transparent;'>CrashLens AI</h1>
            <p style='margin: 0.3rem 0 0 0; color: rgba(255, 255, 255, 0.7); font-size: 1rem;'>Traffic Accident Scenario Analysis • MindSpore-ready • PDF report</p>
        </div>
        """, unsafe_allow_html=True)
    with col3:
        st.write("")
    
    st.markdown("<br>", unsafe_allow_html=True)
    
    # Render sidebar
    render_sidebar()
    
    # Modern step indicator with pills - Using columns to avoid HTML rendering issues
    steps = ["📍 Location", "🚙 Vehicle 1", "🚗 Vehicle 2", "👥 Parties", "📄 Evidence", "🔍 Analysis", "📊 Results"]
    cols = st.columns(len(steps))
    
    for i, (col, step_name) in enumerate(zip(cols, steps), 1):
        with col:
            if i < st.session_state.current_step:
                # Completed step - Green
                st.markdown(f"""
                <div style='text-align: center; padding: 0.75rem 0.25rem; 
                    background: linear-gradient(135deg, #10b981 0%, #059669 100%); 
                    color: white; border-radius: 10px; font-size: 0.8rem; font-weight: 600;
                    box-shadow: 0 4px 12px rgba(16, 185, 129, 0.3);'>
                    ✓ {step_name.split()[-1]}
                </div>
                """, unsafe_allow_html=True)
            elif i == st.session_state.current_step:
                # Active step - Blue gradient
                st.markdown(f"""
                <div style='text-align: center; padding: 0.75rem 0.25rem; 
                    background: linear-gradient(135deg, #3b82f6 0%, #2563eb 100%); 
                    color: white; border-radius: 10px; font-size: 0.8rem; font-weight: 600;
                    box-shadow: 0 6px 16px rgba(59, 130, 246, 0.4);
                    border: 2px solid rgba(96, 165, 250, 0.5);'>
                    {step_name}
                </div>
                """, unsafe_allow_html=True)
            else:
                # Future step - Dark gray
                st.markdown(f"""
                <div style='text-align: center; padding: 0.75rem 0.25rem; 
                    background: rgba(255, 255, 255, 0.05); 
                    color: rgba(255, 255, 255, 0.4); border-radius: 10px; font-size: 0.8rem;
                    border: 1px solid rgba(255, 255, 255, 0.1);'>
                    {step_name.split()[-1]}
                </div>
                """, unsafe_allow_html=True)
    
    st.markdown("<div style='margin: 1.5rem 0; height: 1px; background: linear-gradient(90deg, transparent, rgba(255,255,255,0.1), transparent);'></div>", unsafe_allow_html=True)
    
    # Main content based on current step
    if st.session_state.current_step == 1:
        render_step_1_location()
    elif st.session_state.current_step == 2:
        render_step_2_vehicle1()
    elif st.session_state.current_step == 3:
        render_step_3_vehicle2()
    elif st.session_state.current_step == 4:
        render_step_4_parties()
    elif st.session_state.current_step == 5:
        render_step_5_evidence()
    elif st.session_state.current_step == 6:
        render_step_6_analysis()
    elif st.session_state.current_step == 7:
        render_step_7_results()
    
    # Render footer
    render_footer()


# ============================================================
# STEP 1: LOCATION SELECTION
# ============================================================

def render_step_1_location():
    """Render the location selection step."""
    
    st.header("📍 Step 1: Accident Location")
    
    st.markdown("""
    <div class="info-card">
        <h4>Accident Details</h4>
        <p>Enter the location and conditions where the accident occurred.</p>
    </div>
    """, unsafe_allow_html=True)
    
    col1, col2 = st.columns([2, 1])
    
    with col1:
        # Map display
        render_map_section()
    
    with col2:
        st.subheader("Location Details")
        
        # Location name
        location_name = st.text_input(
            "Location Name",
            value=st.session_state.accident_info['location']['name']
        )
        
        # Coordinates
        lat = st.number_input(
            "Latitude",
            value=st.session_state.accident_info['location']['latitude'],
            format="%.6f",
            step=0.0001
        )
        
        lon = st.number_input(
            "Longitude",
            value=st.session_state.accident_info['location']['longitude'],
            format="%.6f",
            step=0.0001
        )
        
        # Road type - expanded options
        road_types = [
            'roundabout', 'crossroad', 't_junction', 'highway_merge', 
            'parking', 'highway', 'urban_road', 'other'
        ]
        road_type_labels = {
            'roundabout': 'Roundabout (دوار)',
            'crossroad': 'Crossroad (تقاطع)',
            't_junction': 'T-Junction',
            'highway_merge': 'Highway Merge',
            'parking': 'Parking / Low Speed',
            'highway': 'Highway',
            'urban_road': 'Urban Road',
            'other': 'Other'
        }
        
        road_type = st.selectbox(
            "Road Type",
            options=road_types,
            format_func=lambda x: road_type_labels.get(x, x.title()),
            index=road_types.index(st.session_state.accident_info.get('road_type', 'roundabout'))
        )
        
        # Date and time
        accident_date = st.date_input("Accident Date")
        accident_time = st.time_input("Accident Time")
        
        # Weather conditions
        weather = st.selectbox(
            "Weather Conditions",
            options=['clear', 'cloudy', 'rainy', 'foggy', 'sandstorm'],
            format_func=lambda x: x.title()
        )
        
        # Road condition
        road_condition = st.selectbox(
            "Road Condition",
            options=['dry', 'wet', 'sandy', 'oily'],
            format_func=lambda x: x.title()
        )
        
        # Notes
        notes = st.text_area(
            "Notes (optional)",
            value=st.session_state.accident_info.get('notes', ''),
            placeholder="Any extra context about the accident...",
            height=80
        )
        
        # Update session state
        st.session_state.accident_info.update({
            'location': {
                'name': location_name,
                'latitude': lat,
                'longitude': lon,
                'radius_meters': 200
            },
            'road_type': road_type,
            'datetime': f"{accident_date} {accident_time}",
            'weather': weather,
            'road_condition': road_condition,
            'notes': notes
        })
    
    # Navigation buttons
    col1, col2, col3 = st.columns([1, 2, 1])
    with col3:
        if st.button("Next: Vehicle 1 →", type="primary"):
            st.session_state.current_step = 2
            st.rerun()


# ============================================================
# STEP 2: VEHICLE 1 INPUT
# ============================================================

def render_step_2_vehicle1():
    """Render Vehicle 1 input step."""
    
    st.header("🚙 Step 2: First Vehicle Information")
    
    st.markdown("""
    <div class="info-card" style="border-color: #FF4B4B;">
        <h4>Vehicle 1 (Red)</h4>
        <p>Enter the details of the first vehicle involved in the accident.</p>
    </div>
    """, unsafe_allow_html=True)
    
    col1, col2 = st.columns([2, 1])
    
    with col1:
        # Map with path drawing
        render_map_section(vehicle_id=1)
    
    with col2:
        render_vehicle_input(1)
    
    # Navigation buttons
    col1, col2, col3 = st.columns([1, 2, 1])
    with col1:
        if st.button("← Back"):
            st.session_state.current_step = 1
            st.rerun()
    with col3:
        if st.button("Next: Vehicle 2 →", type="primary"):
            st.session_state.current_step = 3
            st.rerun()


# ============================================================
# STEP 3: VEHICLE 2 INPUT
# ============================================================

def render_step_3_vehicle2():
    """Render Vehicle 2 input step."""
    
    st.header("🚗 Step 3: Second Vehicle Information")
    
    st.markdown("""
    <div class="info-card" style="border-color: #4B7BFF;">
        <h4>Vehicle 2 (Blue)</h4>
        <p>Enter the details of the second vehicle involved in the accident.</p>
    </div>
    """, unsafe_allow_html=True)
    
    col1, col2 = st.columns([2, 1])
    
    with col1:
        # Map with path drawing (showing both vehicles)
        render_map_section(vehicle_id=2)
    
    with col2:
        render_vehicle_input(2)
    
    # Navigation buttons
    col1, col2, col3 = st.columns([1, 2, 1])
    with col1:
        if st.button("← Back"):
            st.session_state.current_step = 2
            st.rerun()
    with col3:
        if st.button("Analyze Accident →", type="primary"):
            st.session_state.current_step = 4
            st.rerun()


# ============================================================
# STEP 4: PARTIES INFORMATION
# ============================================================

def render_step_4_parties():
    """Render the parties (driver details) input step."""
    
    st.header("👥 Step 4: Parties Information")
    
    st.markdown("""
    <div class="info-card">
        <h4>Driver & Vehicle Details</h4>
        <p>Enter information about the parties involved in the accident. This information will be included in the official report.</p>
    </div>
    """, unsafe_allow_html=True)
    
    # Two columns for both parties
    col1, col2 = st.columns(2)
    
    with col1:
        render_party_input(1)
    
    with col2:
        render_party_input(2)
    
    # Navigation buttons
    st.markdown("---")
    col1, col2, col3 = st.columns([1, 2, 1])
    with col1:
        if st.button("← Back to Vehicle 2"):
            st.session_state.current_step = 3
            st.rerun()
    with col3:
        if st.button("Next: Evidence →", type="primary"):
            st.session_state.current_step = 5
            st.rerun()


# ============================================================
# STEP 5: EVIDENCE UPLOAD
# ============================================================

def render_step_5_evidence():
    """Render the evidence upload step."""
    
    st.header("📷 Step 5: Evidence (Optional)")
    
    st.markdown("""
    <div class="info-card">
        <h4>Upload Evidence Photos</h4>
        <p>Upload photos of the accident scene, vehicle damage, or any other relevant evidence. This is optional but recommended.</p>
    </div>
    """, unsafe_allow_html=True)
    
    # Evidence upload
    render_evidence_upload()
    
    # Quick summary of data entered so far
    st.markdown("---")
    st.subheader("📋 Data Summary")
    
    col1, col2, col3 = st.columns(3)
    
    with col1:
        st.markdown("**📍 Location**")
        st.write(f"{st.session_state.accident_info['location'].get('name', 'Not set')}")
        st.write(f"Type: {st.session_state.accident_info.get('road_type', 'N/A').replace('_', ' ').title()}")
    
    with col2:
        st.markdown("**🚙 Vehicle 1**")
        # Defensive check - ensure vehicle_1 is a dict
        v1 = st.session_state.vehicle_1
        if isinstance(v1, dict):
            st.write(f"{v1.get('type', 'sedan').title()} @ {v1.get('speed', 50)} km/h")
        else:
            st.write("Vehicle 1 data")
        party1_name = st.session_state.party_1.get('full_name', '') if isinstance(st.session_state.party_1, dict) else ''
        if party1_name:
            st.write(f"Driver: {party1_name}")
    
    with col3:
        st.markdown("**🚗 Vehicle 2**")
        # Defensive check - ensure vehicle_2 is a dict
        v2 = st.session_state.vehicle_2
        if isinstance(v2, dict):
            st.write(f"{v2.get('type', 'sedan').title()} @ {v2.get('speed', 50)} km/h")
        else:
            st.write("Vehicle 2 data")
        party2_name = st.session_state.party_2.get('full_name', '') if isinstance(st.session_state.party_2, dict) else ''
        if party2_name:
            st.write(f"Driver: {party2_name}")
    
    # Navigation buttons
    st.markdown("---")
    col1, col2, col3 = st.columns([1, 2, 1])
    with col1:
        if st.button("← Back to Parties"):
            st.session_state.current_step = 4
            st.rerun()
    with col3:
        if st.button("Run Analysis →", type="primary"):
            st.session_state.current_step = 6
            st.rerun()


# ============================================================
# STEP 6: AI ANALYSIS
# ============================================================

def render_step_6_analysis():
    """Render the AI analysis step."""
    
    st.header("🤖 Step 4: AI Analysis")
    
    st.markdown("""
    <div class="info-card">
        <h4>MindSpore AI Analysis</h4>
        <p>Review the accident data below and click "Run AI Analysis" to generate possible scenarios.</p>
    </div>
    """, unsafe_allow_html=True)
    
    # Data validation
    validation_passed = True
    validation_messages = []
    
    # Check paths
    if not st.session_state.vehicle_1.get('path'):
        validation_messages.append("⚠️ Vehicle 1 path not defined - will use direction-based estimation")
    if not st.session_state.vehicle_2.get('path'):
        validation_messages.append("⚠️ Vehicle 2 path not defined - will use direction-based estimation")
    
    # Summary of inputs in cards
    st.subheader("📋 Input Data Summary")
    
    col1, col2, col3 = st.columns(3)
    
    with col1:
        st.markdown("""
        <div style='background: rgba(59, 130, 246, 0.15); padding: 1rem; border-radius: 12px; border-left: 4px solid #3b82f6; box-shadow: 0 4px 12px rgba(0,0,0,0.2);'>
        <h4 style='margin:0; color:#60a5fa; font-weight: 700;'>📍 Location</h4>
        </div>
        """, unsafe_allow_html=True)
        st.write(f"**{st.session_state.accident_info['location'].get('name', 'Unknown')}**")
        st.write(f"🛣️ Type: `{st.session_state.accident_info['road_type']}`")
        st.write(f"🌤️ Weather: `{st.session_state.accident_info['weather']}`")
        st.write(f"🛤️ Road: `{st.session_state.accident_info['road_condition']}`")
        st.write(f"📅 {st.session_state.accident_info.get('datetime', 'Not specified')}")
    
    with col2:
        st.markdown("""
        <div style='background: rgba(255, 75, 75, 0.15); padding: 1rem; border-radius: 12px; border-left: 4px solid #FF4B4B; box-shadow: 0 4px 12px rgba(0,0,0,0.2);'>
        <h4 style='margin:0; color:#ff6b6b; font-weight: 700;'>🚙 Vehicle 1 (Red)</h4>
        </div>
        """, unsafe_allow_html=True)
        v1 = st.session_state.vehicle_1 if isinstance(st.session_state.vehicle_1, dict) else {'type': 'sedan', 'speed': 50, 'direction': 'north', 'action': 'going_straight', 'path': []}
        st.write(f"**Type:** {v1.get('type', 'sedan').title()}")
        st.write(f"**Speed:** {v1.get('speed', 50)} km/h")
        st.write(f"**Direction:** {v1.get('direction', 'north').title()}")
        st.write(f"**Action:** {v1.get('action', 'going_straight').replace('_', ' ').title()}")
        path_status = "✅ Defined" if v1.get('path') else "⚠️ Not set"
        st.write(f"**Path:** {path_status}")
    
    with col3:
        st.markdown("""
        <div style='background: rgba(75, 123, 255, 0.15); padding: 1rem; border-radius: 12px; border-left: 4px solid #4B7BFF; box-shadow: 0 4px 12px rgba(0,0,0,0.2);'>
        <h4 style='margin:0; color:#6b9bff; font-weight: 700;'>🚗 Vehicle 2 (Blue)</h4>
        </div>
        """, unsafe_allow_html=True)
        v2 = st.session_state.vehicle_2 if isinstance(st.session_state.vehicle_2, dict) else {'type': 'sedan', 'speed': 50, 'direction': 'east', 'action': 'going_straight', 'path': []}
        st.write(f"**Type:** {v2.get('type', 'sedan').title()}")
        st.write(f"**Speed:** {v2.get('speed', 50)} km/h")
        st.write(f"**Direction:** {v2.get('direction', 'east').title()}")
        st.write(f"**Action:** {v2.get('action', 'going_straight').replace('_', ' ').title()}")
        path_status = "✅ Defined" if v2.get('path') else "⚠️ Not set"
        st.write(f"**Path:** {path_status}")
    
    # Show validation messages
    if validation_messages:
        st.markdown("---")
        for msg in validation_messages:
            st.warning(msg)
    
    st.markdown("---")
    
    # Analysis section
    st.subheader("🚀 Run Analysis")
    
    col1, col2, col3 = st.columns([1, 2, 1])
    
    with col2:
        if st.button("🧠 Run MindSpore AI Analysis", type="primary", use_container_width=True):
            
            # Progress bar
            progress_bar = st.progress(0)
            status_text = st.empty()
            
            try:
                # Step 1: Validate data
                status_text.text("Step 1/4: Validating input data...")
                progress_bar.progress(10)
                import time
                time.sleep(0.5)
                
                # Step 2: Extract features
                status_text.text("Step 2/4: Extracting features...")
                progress_bar.progress(30)
                time.sleep(0.5)
                
                # Step 3: Run AI model
                status_text.text("Step 3/4: Running MindSpore AI model...")
                progress_bar.progress(50)
                
                # Import and run analysis
                from analysis.scenario_analyzer import analyze_accident
                
                results = analyze_accident(
                    accident_info=st.session_state.accident_info,
                    vehicle_1=st.session_state.vehicle_1,
                    vehicle_2=st.session_state.vehicle_2
                )
                
                progress_bar.progress(80)
                time.sleep(0.3)
                
                # Step 4: Generate scenarios
                status_text.text("Step 4/4: Generating scenarios...")
                progress_bar.progress(100)
                
                # Store results
                st.session_state.analysis_results = results
                st.session_state.scenarios = results.get('scenarios', [])
                st.session_state.analysis_ready = True
                
                status_text.empty()
                progress_bar.empty()
                
                # Success message
                st.success(f"""
                ✅ **Analysis Complete!**
                
                - Generated **{len(st.session_state.scenarios)}** possible scenarios
                - Most likely: **{results.get('most_likely_scenario', {}).get('type', 'Unknown').replace('_', ' ').title()}**
                - Confidence: **{results.get('overall_collision_probability', 0)*100:.1f}%**
                """)
                
            except Exception as e:
                st.error(f"❌ Analysis failed: {str(e)}")
                progress_bar.empty()
                status_text.empty()
    
    # Show results preview if available
    if st.session_state.analysis_results:
        st.markdown("---")
        st.subheader("📊 Quick Results Preview")
        
        results = st.session_state.analysis_results
        scenarios = st.session_state.scenarios
        
        # Quick metrics
        metric_cols = st.columns(4)
        
        with metric_cols[0]:
            st.metric(
                "Scenarios Generated",
                len(scenarios)
            )
        
        with metric_cols[1]:
            most_likely = results.get('most_likely_scenario', {})
            st.metric(
                "Most Likely",
                f"#{most_likely.get('id', 1)}"
            )
        
        with metric_cols[2]:
            prob = results.get('overall_collision_probability', 0) * 100
            st.metric(
                "Collision Certainty",
                f"{prob:.1f}%"
            )
        
        with metric_cols[3]:
            fault = results.get('preliminary_fault_assessment', {})
            st.metric(
                "Primary Factor",
                fault.get('primary_factor', 'Unknown').replace('_', ' ').title()[:15]
            )
        
        # Top scenarios
        st.write("**Top Scenarios:**")
        for i, scenario in enumerate(scenarios[:3], 1):
            prob_pct = scenario['probability'] * 100
            st.write(f"{i}. **{scenario['accident_type'].replace('_', ' ').title()}** - {prob_pct:.1f}%")
    
    # Navigation
    st.markdown("---")
    col1, col2, col3 = st.columns([1, 2, 1])
    with col1:
        if st.button("← Back to Evidence"):
            st.session_state.current_step = 5
            st.rerun()
    
    with col3:
        if st.session_state.analysis_results:
            if st.button("View Full Results →", type="primary"):
                st.session_state.current_step = 7
                st.rerun()


# ============================================================
# STEP 7: RESULTS
# ============================================================

def render_step_7_results():
    """Render the results step."""
    
    st.header("📊 Step 7: Analysis Results")
    
    if not st.session_state.analysis_results:
        st.warning("No analysis results available. Please run the analysis first.")
        if st.button("← Go to Analysis"):
            st.session_state.current_step = 6
            st.rerun()
        return
    
    # Render full results
    render_results()
    
    # Navigation
    col1, col2, col3 = st.columns([1, 2, 1])
    with col1:
        if st.button("← Back to Analysis"):
            st.session_state.current_step = 6
            st.rerun()
    
    with col3:
        if st.button("🔄 Start New Analysis"):
            # Reset session state
            for key in list(st.session_state.keys()):
                del st.session_state[key]
            st.rerun()


# ============================================================
# RUN APPLICATION
# ============================================================

if __name__ == "__main__":
    main()