import streamlit as st
from transformers import pipeline
import folium
import numpy as np
from scipy.optimize import linear_sum_assignment
import pandas as pd
from folium import plugins

def calculate_distances(coordinates):
    """Calculate distances between all points using Haversine formula"""
    n = len(coordinates)
    distances = np.zeros((n, n))
    
    for i in range(n):
        for j in range(n):
            if i != j:
                lat1, lon1 = coordinates[i]
                lat2, lon2 = coordinates[j]
                
                R = 6371  # Earth's radius in kilometers
                
                dlat = np.radians(lat2 - lat1)
                dlon = np.radians(lon2 - lon1)
                
                a = np.sin(dlat/2)**2 + np.cos(np.radians(lat1)) * np.cos(np.radians(lat2)) * np.sin(dlon/2)**2
                c = 2 * np.arctan2(np.sqrt(a), np.sqrt(1-a))
                distances[i,j] = R * c
                
    return distances

def optimize_route(coordinates):
    """Optimize route using Hungarian algorithm"""
    distances = calculate_distances(coordinates)
    row_ind, col_ind = linear_sum_assignment(distances)
    return col_ind

def get_place_coordinates(place_name):
    """Get coordinates using Nominatim geocoding"""
    try:
        from geopy.geocoders import Nominatim
        geolocator = Nominatim(user_agent="my_travel_app")
        location = geolocator.geocode(place_name)
        return (location.latitude, location.longitude) if location else None
    except:
        return None

def create_numbered_marker(number):
    """Create a custom marker with a number"""
    return plugins.BeautifyIcon(
        number=number,
        border_color='#b4b4b4',
        border_width=1,
        text_color='black',
        background_color='white',
        inner_icon_style='margin-top:0;'
    )

def main():
    st.title("AI Travel Route Optimizer")
    
    # Sidebar for inputs
    st.sidebar.header("Travel Parameters")
    
    # Input for places
    places = []
    num_places = st.sidebar.number_input("Number of places to visit", min_value=2, max_value=10, value=3)
    
    for i in range(num_places):
        place = st.sidebar.text_input(f"Place {i+1}")
        if place:
            coordinates = get_place_coordinates(place)
            if coordinates:
                places.append((place, coordinates))
    
    # Only proceed if we have all places
    if len(places) == num_places and num_places >= 2:
        
        # Extract coordinates for optimization
        coordinates = [coord for _, coord in places]
        
        # Optimize route
        optimized_indices = optimize_route(coordinates)
        
        # Create map
        center_lat = sum(coord[0] for coord in coordinates) / len(coordinates)
        center_lon = sum(coord[1] for coord in coordinates) / len(coordinates)
        
        m = folium.Map(location=[center_lat, center_lon], zoom_start=4)
        
        # Create optimized itinerary list first
        optimized_places = [(i+1, places[idx][0]) for i, idx in enumerate(optimized_indices)]
        
        # Add markers and route lines
        for i in range(len(optimized_indices)):
            current_idx = optimized_indices[i]
            next_idx = optimized_indices[(i + 1) % len(optimized_indices)]
            
            # Add numbered marker
            place_name, (lat, lon) = places[current_idx]
            folium.Marker(
                [lat, lon],
                popup=f"Stop {i+1}: {place_name}",
                icon=create_numbered_marker(i+1)
            ).add_to(m)
            
            # Add line to next point
            next_lat, next_lon = coordinates[next_idx]
            folium.PolyLine(
                locations=[[lat, lon], [next_lat, next_lon]],
                weight=2,
                color='blue',
                opacity=0.8
            ).add_to(m)
        
        # Display map
        st.components.v1.html(m._repr_html_(), height=600)
        
        # Display optimized itinerary
        st.header("Optimized Itinerary")
        for stop_num, place_name in optimized_places:
            st.write(f"{stop_num}. {place_name}")
        
        # Calculate total distance
        total_distance = 0
        distances = calculate_distances(coordinates)
        for i in range(len(optimized_indices)):
            current_idx = optimized_indices[i]
            next_idx = optimized_indices[(i + 1) % len(optimized_indices)]
            total_distance += distances[current_idx][next_idx]
            
        st.write(f"\nTotal distance: {total_distance:.2f} km")

if __name__ == "__main__":
    main()