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):
    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)
                min_dist = inf
                for k in subset:
                    if k == next_city:
                        continue
                    current_dist = memo[(prev_bits, k)] + dist_matrix[k][next_city]
                    if current_dist < min_dist:
                        min_dist = current_dist
                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):
        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):
        tour.append(last_city)
        bits &= ~(1 << last_city)
        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

    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)
    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

# 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")

# Input kota
st.subheader("Pilih Kota")
city_count = st.number_input("Jumlah kota", min_value=2, step=1)
city_names = []
for i in range(int(city_count)):
    city_name = st.text_input(f"Kota {i+1}", key=f"city_{i}")
    city_names.append(city_name)

if st.button("Optimasi Rute"):
    coordinates = [get_coordinates(city) for city in city_names]
    valid_coordinates = [c for c in coordinates if c is not None]

    if len(valid_coordinates) < 2:
        st.error("Ada kota yang tidak ditemukan. Pastikan nama kota benar.")
    else:
        # Buat distance matrix
        dist_matrix = create_distance_matrix(valid_coordinates)
        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_folium = st.components.v1.html(folium.Map._repr_html_(map_obj), width=700, height=500)