import streamlit as st
import folium
from geopy.geocoders import Nominatim
from geopy.distance import geodesic
from itertools import combinations
import numpy as np

# Algoritma Held-Karp untuk TSP
def held_karp_tsp(dist_matrix):
    if len(dist_matrix) < 2:
        return 0, []

    n = len(dist_matrix)
    inf = float('inf')
    memo = {}

    # Initialize memo table for subsets of size 2
    for i in range(1, n):
        memo[(1 << i, i)] = dist_matrix[0][i]

    # Fill memo table for larger subsets
    for subset_size in range(2, n):
        for subset in combinations(range(1, n), subset_size):
            bits = 0
            for bit in subset:
                bits |= 1 << bit
            for next_city in subset:
                prev_bits = bits & ~(1 << next_city)
                if (prev_bits, next_city) in memo:
                    min_dist = memo[(prev_bits, next_city)] + dist_matrix[next_city][0]
                    memo[(bits, next_city)] = min_dist

    # Find optimal tour
    bits = (2 ** n - 1) - 1
    min_tour_cost = inf
    last_city = None
    for k in range(1, n):
        if (bits, k) in memo:
            tour_cost = memo[(bits, k)] + dist_matrix[k][0]
            if tour_cost < min_tour_cost:
                min_tour_cost = tour_cost
                last_city = k

    # Backtrack to find the full tour
    tour = [0]
    bits = (2 ** n - 1) - 1
    for _ in range(n - 1):
        if last_city is not None:
            tour.append(last_city)
            bits &= ~(1 << last_city)
            if bits == 0:
                break
            next_city = min(
                [(memo[(bits, k)] + dist_matrix[k][last_city], k) for k in range(n) if (bits, k) in memo],
                key=lambda x: x[0],
            )[1]
            last_city = next_city
        else:
            break

    tour.append(0)
    return min_tour_cost, tour

# Mendapatkan koordinat kota dari nama
def get_coordinates(city_name):
    geolocator = Nominatim(user_agent="tsp_app")
    location = geolocator.geocode(city_name)
    if location:
        return (location.latitude, location.longitude)
    else:
        return None

# Menghitung jarak antar semua pasangan kota
def create_distance_matrix(coordinates):
    valid_coordinates = [c for c in coordinates if c is not None]
    n = len(valid_coordinates)
    if n < 2:
        return [], []
    dist_matrix = np.zeros((n, n))
    for i in range(n):
        for j in range(i + 1, n):
            dist_matrix[i][j] = geodesic(valid_coordinates[i], valid_coordinates[j]).kilometers
            dist_matrix[j][i] = dist_matrix[i][j]
    return dist_matrix, valid_coordinates

# Fungsi untuk menampilkan peta rute
def plot_route(map_obj, coordinates, route):
    for i in range(len(route) - 1):
        start = coordinates[route[i]]
        end = coordinates[route[i + 1]]
        folium.Marker(location=start, tooltip=f'City {route[i] + 1}').add_to(map_obj)
        folium.Marker(location=end, tooltip=f'City {route[i + 1] + 1}').add_to(map_obj)
        folium.PolyLine([start, end], color="blue", weight=2.5, opacity=1).add_to(map_obj)

    # Add markers for start and end points
    folium.Marker(location=coordinates[0], tooltip='Start', icon=folium.Icon(color="green")).add_to(map_obj)
    folium.Marker(location=coordinates[route[-2]], tooltip='End', icon=folium.Icon(color="red")).add_to(map_obj)
    return map_obj

# Streamlit UI
st.title("Traveling Salesman Problem Solver")

# Create map
map_obj = folium.Map(location=[0, 0], zoom_start=2)

# Input kota
st.subheader("Pilih Kota")
city_count = st.number_input("Jumlah kota", min_value=2, step=1)
city_names = []
city_coords = []
for i in range(int(city_count)):
    city_name = st.text_input(f"Kota {i+1}", key=f"city_{i}")
    city_coords.append(get_coordinates(city_name))
    if city_coords[-1] is None:
        st.warning(f"Kota '{city_name}' tidak ditemukan. Harap periksa ejaan dan coba lagi.")
    else:
        city_names.append(city_name)

if st.button("Optimasi Rute"):
    dist_matrix, valid_coordinates = create_distance_matrix(city_coords)

    if len(valid_coordinates) < 2:
        st.error("Tidak cukup kota yang valid untuk dioptimalkan.")
    else:
        # Hitung rute optimal
        min_cost, optimal_route = held_karp_tsp(dist_matrix)
        
        # Tampilkan hasil
        st.write(f"Biaya total minimum: {min_cost:.2f} km")
        st.write("Rute optimal:", " -> ".join([city_names[i] for i in optimal_route]))

        # Buat peta
        map_obj = folium.Map(location=valid_coordinates[0], zoom_start=5)
        plot_route(map_obj, valid_coordinates, optimal_route)
        
        # Render map
        st.markdown("### Rute Optimal")
        st_folium = st.components.v1.html(folium.Map._repr_html_(map_obj), width=800, height=600)