ui/map_viewer.py

"""
Map Viewer Component
====================
Handles map display and vehicle path drawing using Folium.
"""

import streamlit as st
import folium
from streamlit_folium import st_folium, folium_static
from folium.plugins import Draw
import json
from typing import List

from config import CASE_STUDY_LOCATION, MAP_CONFIG, COLORS


def create_base_map(location: dict = None) -> folium.Map:
    """
    Create a base Folium map centered on the accident location.
    
    Args:
        location: Dictionary with latitude, longitude, and name
    
    Returns:
        Folium Map object
    """
    if location is None:
        location = CASE_STUDY_LOCATION
    
    # Create map
    m = folium.Map(
        location=[location['latitude'], location['longitude']],
        zoom_start=MAP_CONFIG['default_zoom'],
        tiles='OpenStreetMap'
    )
    
    # Add location marker
    folium.Marker(
        location=[location['latitude'], location['longitude']],
        popup=location.get('name', 'Accident Location'),
        icon=folium.Icon(color='red', icon='info-sign'),
        tooltip="Accident Location"
    ).add_to(m)
    
    # Add circle to show area of interest
    folium.Circle(
        location=[location['latitude'], location['longitude']],
        radius=location.get('radius_meters', 200),
        color='blue',
        fill=True,
        fill_opacity=0.1,
        popup="Analysis Area"
    ).add_to(m)
    
    return m


def add_draw_control(m: folium.Map) -> folium.Map:
    """
    Add drawing controls to the map for path definition.
    
    Args:
        m: Folium Map object
    
    Returns:
        Folium Map with draw controls
    """
    draw = Draw(
        draw_options={
            'polyline': {
                'allowIntersection': True,
                'shapeOptions': {
                    'color': '#FF4B4B',
                    'weight': 4
                }
            },
            'polygon': False,
            'circle': False,
            'rectangle': False,
            'circlemarker': False,
            'marker': True
        },
        edit_options={'edit': True, 'remove': True}
    )
    draw.add_to(m)
    
    return m


def add_vehicle_path(m: folium.Map, path: list, vehicle_id: int) -> folium.Map:
    """
    Add a vehicle path to the map.
    
    Args:
        m: Folium Map object
        path: List of [lat, lng] coordinates
        vehicle_id: 1 or 2 to determine color
    
    Returns:
        Folium Map with vehicle path
    """
    if not path or len(path) < 2:
        return m
    
    color = COLORS['vehicle_1'] if vehicle_id == 1 else COLORS['vehicle_2']
    
    # Add the path line
    folium.PolyLine(
        locations=path,
        color=color,
        weight=4,
        opacity=0.8,
        popup=f"Vehicle {vehicle_id} Path",
        tooltip=f"Vehicle {vehicle_id}"
    ).add_to(m)
    
    # Add start marker
    folium.Marker(
        location=path[0],
        popup=f"Vehicle {vehicle_id} Start",
        icon=folium.Icon(
            color='green' if vehicle_id == 1 else 'blue',
            icon='play'
        )
    ).add_to(m)
    
    # Add end marker
    folium.Marker(
        location=path[-1],
        popup=f"Vehicle {vehicle_id} End",
        icon=folium.Icon(
            color='red' if vehicle_id == 1 else 'darkblue',
            icon='stop'
        )
    ).add_to(m)
    
    # Add direction arrows
    for i in range(len(path) - 1):
        mid_lat = (path[i][0] + path[i+1][0]) / 2
        mid_lng = (path[i][1] + path[i+1][1]) / 2
        
        folium.RegularPolygonMarker(
            location=[mid_lat, mid_lng],
            number_of_sides=3,
            radius=8,
            color=color,
            fill=True,
            fill_color=color,
            fill_opacity=0.7
        ).add_to(m)
    
    return m


def add_collision_point(m: folium.Map, collision_point: list) -> folium.Map:
    """
    Add a collision point marker to the map.
    
    Args:
        m: Folium Map object
        collision_point: [lat, lng] of collision
    
    Returns:
        Folium Map with collision marker
    """
    if not collision_point:
        return m
    
    folium.Marker(
        location=collision_point,
        popup="Estimated Collision Point",
        icon=folium.Icon(color='orange', icon='warning-sign'),
        tooltip="💥 Collision Point"
    ).add_to(m)
    
    # Add impact radius
    folium.Circle(
        location=collision_point,
        radius=10,
        color=COLORS['collision_point'],
        fill=True,
        fill_opacity=0.5,
        popup="Impact Zone"
    ).add_to(m)
    
    return m


def render_map_section(vehicle_id: int = None):
    """
    Render the map section in Streamlit.
    
    Args:
        vehicle_id: If provided, enables path drawing for that vehicle
    """
    location = st.session_state.accident_info['location']
    
    # Create base map
    m = create_base_map(location)
    
    # Add existing vehicle paths
    if st.session_state.vehicle_1.get('path'):
        m = add_vehicle_path(m, st.session_state.vehicle_1['path'], 1)
    
    if st.session_state.vehicle_2.get('path'):
        m = add_vehicle_path(m, st.session_state.vehicle_2['path'], 2)
    
    # Add draw controls if editing a specific vehicle
    if vehicle_id:
        m = add_draw_control(m)
        
        st.info(f"🖊️ Draw the path for Vehicle {vehicle_id} on the map. Click points to create a path, then click 'Finish' to complete.")
    
    # Render map and get drawing data
    map_data = st_folium(
        m,
        width=700,
        height=500,
        returned_objects=["last_active_drawing", "all_drawings"]
    )
    
    # Process drawn paths
    if vehicle_id and map_data and map_data.get('last_active_drawing'):
        drawing = map_data['last_active_drawing']
        
        if drawing.get('geometry', {}).get('type') == 'LineString':
            coords = drawing['geometry']['coordinates']
            # Convert from [lng, lat] to [lat, lng]
            path = [[c[1], c[0]] for c in coords]
            
            if vehicle_id == 1:
                st.session_state.vehicle_1['path'] = path
            else:
                st.session_state.vehicle_2['path'] = path
            
            st.success(f"✅ Path saved for Vehicle {vehicle_id}!")
    
    # Show path info and preset options
    if vehicle_id:
        vehicle_key = f'vehicle_{vehicle_id}'
        current_path = st.session_state[vehicle_key].get('path', [])
        
        if current_path:
            st.success(f"**✅ Path defined:** {len(current_path)} points")
        else:
            st.warning("⚠️ No path drawn yet. Use the drawing tools on the map OR select a preset path below.")
        
        # Preset path options for roundabout
        st.markdown("---")
        st.markdown("**🛣️ Quick Path Selection (Roundabout)**")
        
        preset_col1, preset_col2 = st.columns(2)
        
        with preset_col1:
            entry_direction = st.selectbox(
                f"Entry Direction (V{vehicle_id})",
                options=['north', 'south', 'east', 'west'],
                key=f"entry_dir_{vehicle_id}"
            )
        
        with preset_col2:
            exit_direction = st.selectbox(
                f"Exit Direction (V{vehicle_id})",
                options=['north', 'south', 'east', 'west'],
                index=2 if entry_direction == 'north' else 0,
                key=f"exit_dir_{vehicle_id}"
            )
        
        if st.button(f"🔄 Generate Path for Vehicle {vehicle_id}", key=f"gen_path_{vehicle_id}"):
            generated_path = generate_roundabout_path(
                location['latitude'],
                location['longitude'],
                entry_direction,
                exit_direction
            )
            
            if vehicle_id == 1:
                st.session_state.vehicle_1['path'] = generated_path
            else:
                st.session_state.vehicle_2['path'] = generated_path
            
            st.success(f"✅ Path generated: {entry_direction.title()} → {exit_direction.title()}")
            st.rerun()
        
        # Manual path input as fallback
        with st.expander("📝 Or enter path manually"):
            st.write("Enter coordinates as: lat1,lng1;lat2,lng2;...")
            manual_path = st.text_input(
                "Path coordinates",
                key=f"manual_path_{vehicle_id}",
                placeholder="26.2397,50.5369;26.2400,50.5372"
            )
            
            if st.button(f"Set Manual Path", key=f"set_path_{vehicle_id}"):
                if manual_path:
                    try:
                        path = []
                        for point in manual_path.split(';'):
                            lat, lng = point.strip().split(',')
                            path.append([float(lat), float(lng)])
                        
                        if vehicle_id == 1:
                            st.session_state.vehicle_1['path'] = path
                        else:
                            st.session_state.vehicle_2['path'] = path
                        
                        st.success(f"Path set with {len(path)} points!")
                        st.rerun()
                    except Exception as e:
                        st.error(f"Invalid format: {e}")


def generate_roundabout_path(
    center_lat: float,
    center_lng: float,
    entry_direction: str,
    exit_direction: str
) -> List[List[float]]:
    """
    Generate a realistic path through a roundabout.
    
    Args:
        center_lat: Center latitude of roundabout
        center_lng: Center longitude of roundabout
        entry_direction: Direction vehicle enters from
        exit_direction: Direction vehicle exits to
    
    Returns:
        List of [lat, lng] coordinates forming the path
    """
    import math
    
    # Roundabout parameters
    approach_distance = 0.0015  # ~150 meters
    roundabout_radius = 0.0004  # ~40 meters
    
    # Direction to angle mapping (0 = North, clockwise)
    dir_to_angle = {
        'north': 90,    # Top
        'east': 0,      # Right
        'south': 270,   # Bottom
        'west': 180     # Left
    }
    
    entry_angle = math.radians(dir_to_angle[entry_direction])
    exit_angle = math.radians(dir_to_angle[exit_direction])
    
    path = []
    
    # Entry point (outside roundabout)
    entry_lat = center_lat + approach_distance * math.sin(entry_angle)
    entry_lng = center_lng + approach_distance * math.cos(entry_angle)
    path.append([entry_lat, entry_lng])
    
    # Entry to roundabout edge
    edge_lat = center_lat + roundabout_radius * 1.5 * math.sin(entry_angle)
    edge_lng = center_lng + roundabout_radius * 1.5 * math.cos(entry_angle)
    path.append([edge_lat, edge_lng])
    
    # Points along the roundabout (clockwise)
    entry_deg = dir_to_angle[entry_direction]
    exit_deg = dir_to_angle[exit_direction]
    
    # Calculate arc (always go clockwise in roundabout)
    if exit_deg <= entry_deg:
        exit_deg += 360
    
    # Add intermediate points along the arc
    num_arc_points = max(2, (exit_deg - entry_deg) // 45)
    for i in range(1, num_arc_points + 1):
        angle = entry_deg + (exit_deg - entry_deg) * i / (num_arc_points + 1)
        angle_rad = math.radians(angle)
        arc_lat = center_lat + roundabout_radius * math.sin(angle_rad)
        arc_lng = center_lng + roundabout_radius * math.cos(angle_rad)
        path.append([arc_lat, arc_lng])
    
    # Exit from roundabout edge
    exit_edge_lat = center_lat + roundabout_radius * 1.5 * math.sin(exit_angle)
    exit_edge_lng = center_lng + roundabout_radius * 1.5 * math.cos(exit_angle)
    path.append([exit_edge_lat, exit_edge_lng])
    
    # Exit point (outside roundabout)
    exit_lat = center_lat + approach_distance * math.sin(exit_angle)
    exit_lng = center_lng + approach_distance * math.cos(exit_angle)
    path.append([exit_lat, exit_lng])
    
    return path


def render_results_map(scenarios: list, selected_scenario: int = 0):
    """
    Render a map showing analysis results for a specific scenario.
    
    Args:
        scenarios: List of generated scenarios
        selected_scenario: Index of selected scenario
    """
    if not scenarios:
        st.warning("No scenarios to display")
        return
    
    scenario = scenarios[selected_scenario]
    location = st.session_state.accident_info['location']
    
    # Create map
    m = create_base_map(location)
    
    # Add vehicle paths
    if scenario.get('vehicle_1_path'):
        m = add_vehicle_path(m, scenario['vehicle_1_path'], 1)
    
    if scenario.get('vehicle_2_path'):
        m = add_vehicle_path(m, scenario['vehicle_2_path'], 2)
    
    # Add collision point
    if scenario.get('collision_point'):
        m = add_collision_point(m, scenario['collision_point'])
    
    # Add scenario info popup
    info_html = f"""
    <div style="width: 200px;">
        <h4>Scenario {selected_scenario + 1}</h4>
        <p><b>Probability:</b> {scenario.get('probability', 0)*100:.1f}%</p>
        <p><b>Type:</b> {scenario.get('accident_type', 'Unknown')}</p>
    </div>
    """
    
    folium.Marker(
        location=[location['latitude'], location['longitude']],
        popup=folium.Popup(info_html, max_width=250),
        icon=folium.Icon(color='purple', icon='info-sign')
    ).add_to(m)
    
    # Display map using static version to prevent jumping
    folium_static(m, width=700, height=500)