import streamlit as st
import folium
import numpy as np
from scipy.optimize import linear_sum_assignment
from folium import plugins
import requests
from datetime import datetime
from geopy.geocoders import Nominatim
import polyline
from typing import List, Tuple, Dict, Optional
import time

# Utility Functions
def create_numbered_marker(number: int) -> folium.features.DivIcon:
    """Create a numbered marker icon for the map"""
    return folium.features.DivIcon(
        html=f'<div style="background-color: white; border-radius: 50%; width: 25px; height: 25px; '
             f'display: flex; align-items: center; justify-content: center; border: 2px solid blue; '
             f'font-weight: bold;">{number}</div>'
    )

def get_place_coordinates(place_name: str) -> Optional[Tuple[float, float]]:
    """Get coordinates for a place using Nominatim geocoding"""
    try:
        geolocator = Nominatim(user_agent="travel_optimizer")
        location = geolocator.geocode(place_name)
        if location:
            return (location.latitude, location.longitude)
    except Exception as e:
        st.error(f"Error geocoding {place_name}: {str(e)}")
    return None

def format_instructions(steps: List[Dict]) -> List[str]:
    """Format OSRM route steps into readable instructions"""
    instructions = []
    for step in steps:
        if 'maneuver' in step:
            instruction = step.get('maneuver', {}).get('instruction', '')
            if instruction:
                # Add distance information if available
                if 'distance' in step:
                    distance_km = step['distance'] / 1000
                    if distance_km >= 1:
                        instruction += f" ({distance_km:.1f} km)"
                    else:
                        instruction += f" ({step['distance']:.0f} m)"
                instructions.append(instruction)
    return instructions

def get_route_from_osrm(coord1: Tuple[float, float], coord2: Tuple[float, float]) -> Optional[Dict]:
    """Get route information from OSRM"""
    url = f"http://router.project-osrm.org/route/v1/driving/{coord1[1]},{coord1[0]};{coord2[1]},{coord2[0]}"
    params = {
        "overview": "full",
        "geometries": "polyline",
        "steps": "true",
        "annotations": "true"
    }
    
    # Add rate limiting
    time.sleep(1)  # Respect OSRM usage guidelines
    
    try:
        response = requests.get(url, params=params)
        if response.status_code == 200:
            data = response.json()
            if data["code"] == "Ok" and len(data["routes"]) > 0:
                route = data["routes"][0]
                return {
                    'distance': route['distance'] / 1000,  # Convert to km
                    'duration': route['duration'] / 60,    # Convert to minutes
                    'geometry': route['geometry'],
                    'steps': route['legs'][0]['steps']
                }
    except Exception as e:
        st.warning(f"Error getting route: {str(e)}")
    return None

def calculate_real_distances(places: List[Tuple[str, Tuple[float, float]]]) -> Tuple[np.ndarray, Dict]:
    """Calculate real driving distances using OSRM"""
    n = len(places)
    distances = np.zeros((n, n))
    routing_info = {}
    
    progress_bar = st.progress(0)
    total_calcs = n * (n-1)
    current_calc = 0
    
    for i in range(n):
        for j in range(n):
            if i != j:
                origin = places[i][1]
                destination = places[j][1]
                
                route_data = get_route_from_osrm(origin, destination)
                
                if route_data:
                    distances[i,j] = route_data['distance']
                    routing_info[(i,j)] = {
                        'distance_km': f"{route_data['distance']:.1f} km",
                        'duration_mins': f"{route_data['duration']:.0f} mins",
                        'geometry': route_data['geometry'],
                        'steps': route_data['steps']
                    }
                
                current_calc += 1
                progress_bar.progress(current_calc / total_calcs)
    
    return distances, routing_info

def optimize_route(distances: np.ndarray) -> List[int]:
    """Optimize route using modified Hungarian algorithm for TSP"""
    n = len(distances)
    row_ind, col_ind = linear_sum_assignment(distances)
    
    # Convert assignment to tour
    tour = []
    for i in range(n):
        tour.append(col_ind[i])
    
    return tour

def main():
    st.set_page_config(page_title="AI Travel Route Optimizer", layout="wide")
    
    st.title("🌍 AI Travel Route Optimizer")
    st.markdown("""
    This app helps you find the optimal route between multiple destinations using real driving distances.
    Enter your destinations in the sidebar and get a customized route with turn-by-turn directions.
    """)
    
    # Sidebar for inputs
    with st.sidebar:
        st.header("📍 Destinations")
        
        # Input for places
        num_places = st.number_input("Number of destinations", min_value=2, max_value=10, value=3)
        
        places = []
        for i in range(num_places):
            place = st.text_input(f"Destination {i+1}")
            if place:
                with st.spinner(f"Finding coordinates for {place}..."):
                    coordinates = get_place_coordinates(place)
                    if coordinates:
                        places.append((place, coordinates))
                        st.success(f"✓ Found coordinates for {place}")
                    else:
                        st.error(f"❌ Couldn't find coordinates for {place}")
    
    # Main content area
    if len(places) == num_places and num_places >= 2:
        col1, col2 = st.columns([2, 1])
        
        with col1:
            with st.spinner("🚗 Calculating optimal route..."):
                # Calculate real distances
                distances, routing_info = calculate_real_distances(places)
                
                # Get optimized route order
                optimal_order = optimize_route(distances)
                
                # Create map
                center_lat = sum(coord[0] for _, coord in places) / len(places)
                center_lon = sum(coord[1] for _, coord in places) / len(places)
                m = folium.Map(location=[center_lat, center_lon], zoom_start=4)
                
                # Add markers and route lines
                total_distance = 0
                total_duration = 0
                
                for i in range(len(optimal_order)):
                    current_idx = optimal_order[i]
                    next_idx = optimal_order[(i + 1) % len(optimal_order)]
                    
                    current_place, (lat, lon) = places[current_idx]
                    segment_info = routing_info.get((current_idx, next_idx), {})
                    
                    # Add marker
                    popup_content = f"""
                    <b>Stop {i+1}: {current_place}</b><br>
                    """
                    if i < len(optimal_order) - 1:
                        popup_content += f"""
                        To next stop:<br>
                        Distance: {segment_info.get('distance_km', 'N/A')}<br>
                        Duration: {segment_info.get('duration_mins', 'N/A')}
                        """
                    
                    folium.Marker(
                        [lat, lon],
                        popup=folium.Popup(popup_content, max_width=300),
                        icon=create_numbered_marker(i+1)
                    ).add_to(m)
                    
                    # Draw route line
                    if i < len(optimal_order) - 1:
                        if 'geometry' in segment_info:
                            try:
                                route_coords = polyline.decode(segment_info['geometry'])
                                folium.PolyLine(
                                    route_coords,
                                    weight=2,
                                    color='blue',
                                    opacity=0.8
                                ).add_to(m)
                            except Exception as e:
                                st.warning(f"Error drawing route: {str(e)}")
                                next_place = places[next_idx][1]
                                folium.PolyLine(
                                    [[lat, lon], [next_place[0], next_place[1]]],
                                    weight=2,
                                    color='red',
                                    opacity=0.8
                                ).add_to(m)
                        
                        # Add to totals
                        if 'distance_km' in segment_info:
                            total_distance += float(segment_info['distance_km'].split()[0])
                        if 'duration_mins' in segment_info:
                            total_duration += float(segment_info['duration_mins'].split()[0])
                
                # Display map
                st.components.v1.html(m._repr_html_(), height=600)
        
        with col2:
            st.header("📝 Optimized Itinerary")
            
            # Display summary first
            st.markdown("### 📊 Summary")
            st.info(f"""
            🛣️ Total distance: **{total_distance:.1f} km**  
            ⏱️ Total duration: **{total_duration:.0f} mins** ({total_duration/60:.1f} hours)
            """)
            
            st.markdown("### 🗺️ Route Details")
            for i in range(len(optimal_order)):
                current_idx = optimal_order[i]
                next_idx = optimal_order[(i + 1) % len(optimal_order)]
                
                current_place = places[current_idx][0]
                segment_info = routing_info.get((current_idx, next_idx), {})
                
                st.markdown(f"**Stop {i+1}: {current_place}**")
                
                if i < len(optimal_order) - 1:
                    st.markdown(f"↓ *{segment_info.get('distance_km', 'N/A')}, "
                              f"Duration: {segment_info.get('duration_mins', 'N/A')}*")
                    
                    if 'steps' in segment_info:
                        with st.expander("📍 Turn-by-turn directions"):
                            instructions = format_instructions(segment_info['steps'])
                            for idx, instruction in enumerate(instructions, 1):
                                st.write(f"{idx}. {instruction}")
    else:
        st.info("👈 Please enter all destinations in the sidebar to get started")

if __name__ == "__main__":
    main()