import streamlit as st
import osmnx as ox
import networkx as nx
import pandas as pd
import plotly.graph_objects as go
import heapq
from math import asin, radians, cos, sin, sqrt

# Plotting
DEFAULT_WEIGHT_KG = 70

ELEVATION_SCALE = 5

# --- Location list ---
locations = [
    {'name': 'Grha Sabha Pramana', 'lat': -7.770071488953822, 'lon': 110.3783984379329, 'elevation': 142},
    {'name': 'Boulevard UGM', 'lat': -7.775446818599403, 'lon': 110.37643708634688, 'elevation': 133},
    {'name': 'Fakultas Ilmu Sosial dan Ilmu Politik', 'lat': -7.769082851574989, 'lon': 110.38088752769059, 'elevation': 145},
    {'name': 'Fakultas Teknik', 'lat': -7.765067267048453, 'lon': 110.37252242878728, 'elevation': 139},
    {'name': 'Perpustakaan Pusat', 'lat': -7.769132738499928, 'lon': 110.37819025762326, 'elevation': 142},
    {'name': 'Fakultas Farmasi', 'lat': -7.768920115853235, 'lon': 110.37629125348137, 'elevation': 143},
    {'name': 'Fakultas Kedokteran', 'lat': -7.769241968897335, 'lon': 110.37410241529511, 'elevation': 138},
    {'name': 'Fakultas Hukum', 'lat': -7.770083777982755, 'lon': 110.38117766903606, 'elevation': 143},
    {'name': 'Fakultas Ekonomika dan Bisnis', 'lat': -7.770126237227525, 'lon': 110.37903152864811, 'elevation': 142},
    {'name': 'Fakultas Ilmu Budaya', 'lat': -7.771190276512944, 'lon': 110.37876576685031, 'elevation': 139},
    {'name': 'Fakultas Biologi', 'lat': -7.76442822947009, 'lon': 110.37698323450962, 'elevation': 148},
    {'name': 'Fakultas Kehutanan', 'lat': -7.766739271152903, 'lon': 110.38109875565928, 'elevation': 149},
    {'name': 'Fakultas Peternakan', 'lat': -7.766930486720981, 'lon': 110.38542241877909, 'elevation': 144},
    {'name': 'Fakultas Pertanian', 'lat': -7.767908235085672, 'lon': 110.38100213548348, 'elevation': 148},
    {'name': 'Fakultas Geografi', 'lat': -7.765983577235556, 'lon': 110.37794577655771, 'elevation': 150},
    {'name': 'Fakultas MIPA', 'lat': -7.766770213071046, 'lon': 110.37617551844697, 'elevation': 142},
    {'name': 'Fakultas Psikologi', 'lat': -7.7721010490703355, 'lon': 110.38076316079525, 'elevation': 138},
    {'name': 'Fakultas Filsafat', 'lat': -7.771194580135096, 'lon': 110.38124107602673, 'elevation': 140},
    {'name': 'Rektorat UGM', 'lat': -7.767550457927047, 'lon': 110.37912599375315, 'elevation': 155},
    {'name': 'Masjid Kampus UGM', 'lat': -7.773202436052038, 'lon': 110.38040009798743, 'elevation': 140},
    {'name': 'RSUP Dr. Sardjito', 'lat': -7.76857375077125, 'lon': 110.37380676773756, 'elevation': 138}
]

# Dynamically generate key_locations for quick lookup
key_locations = {loc["name"]: (loc["lat"], loc["lon"], loc["elevation"]) for loc in locations}

def haversine(lat1, lon1, lat2, lon2):
    dlon = radians(lon2 - lon1)
    dlat = radians(lat2 - lat1)
    a = sin(dlat/2)**2 + cos(radians(lat1)) * cos(radians(lat2)) * sin(dlon/2)**2
    c = 2 * asin(sqrt(a))
    return 6371 * c  # km

def elevation_difference(u, v):
    # Get the nearest locations for nodes u and v
    u_coords = (G.nodes[u]['y'], G.nodes[u]['x'])  # lat, lon
    v_coords = (G.nodes[v]['y'], G.nodes[v]['x'])  # lat, lon

    # Find the nearest locations from the locations list
    u_location = min(locations, key=lambda loc: haversine(u_coords[0], u_coords[1], loc['lat'], loc['lon']))
    v_location = min(locations, key=lambda loc: haversine(v_coords[0], v_coords[1], loc['lat'], loc['lon']))

    # Get the elevation values from the nearest locations
    elev_u = u_location['elevation']
    elev_v = v_location['elevation']

    # Return the absolute elevation difference
    return abs(elev_v - elev_u)

def heuristic(u, v):
    u_lat, u_lon = G.nodes[u]['y'], G.nodes[u]['x']
    v_lat, v_lon = G.nodes[v]['y'], G.nodes[v]['x']
    return haversine(u_lat, u_lon, v_lat, v_lon) + (ELEVATION_SCALE * elevation_difference(u, v))

def astar(graph, start, goal):
    open_set = []
    heapq.heappush(open_set, (0, start))
    came_from = {}
    g_score = {node: float('inf') for node in graph.nodes}
    g_score[start] = 0
    f_score = {node: float('inf') for node in graph.nodes}
    f_score[start] = heuristic(start, goal)

    while open_set:
        _, current = heapq.heappop(open_set)
        if current == goal:
            path = []
            while current in came_from:
                path.append(current)
                current = came_from[current]
            path.append(start)
            return path[::-1]

        for neighbor in graph.neighbors(current):
            tentative_g = g_score[current] + graph[current][neighbor][0]['length']
            if tentative_g < g_score[neighbor]:
                came_from[neighbor] = current
                g_score[neighbor] = tentative_g
                f_score[neighbor] = tentative_g + heuristic(neighbor, goal)
                heapq.heappush(open_set, (f_score[neighbor], neighbor))

    raise nx.NetworkXNoPath("No path found")

def get_coordinates(locations, target_name):
    for loc in locations:
        if loc["name"].lower() == target_name.lower():
            return (loc["lon"], loc["lat"])  # Return in (lon, lat) format
    return None

def adjust_map_zoom_and_center(lats, lons, real_distance, total_distance):
    # Calculate the center of the route for better zoom
    center_lat = sum(lats) / len(lats)
    center_lon = sum(lons) / len(lons)

    # Adjust zoom level based on the route length
    route_length_km = real_distance / 1000 if real_distance else total_distance
    zoom_level = max(10, min(16, 16 - int(route_length_km / 2)))

    return center_lat, center_lon, zoom_level

# Streamlit App
st.set_page_config(page_title="UGM Running Route", layout="wide")
st.title("🚶‍♂️ Running Route Through UGM")
st.markdown("---")

# Layout: Input, Map, and Output
col1, col2, col3 = st.columns([1, 2, 1])

# Input Section
with col1:
    st.header("📍 Input")
    st.markdown("Select your starting and ending points for the route.")
    with st.form("route_form"):
        st.write("### Route Options")
        location_names = [loc["name"] for loc in locations]
        start_point = st.selectbox("Start Point", location_names)
        end_point = st.selectbox("End Point", location_names)
        submitted = st.form_submit_button('🚀 Submit Route')
        
start_loc = next(loc for loc in locations if loc['name'] == start_point)
end_loc = next(loc for loc in locations if loc['name'] == end_point)
center_lat = (start_loc['lat'] + end_loc['lat']) / 2
center_lon = (start_loc['lon'] + end_loc['lon']) / 2
G = ox.graph_from_point((center_lat, center_lon), dist=1000, network_type='walk')

start_node = ox.distance.nearest_nodes(G, start_loc['lon'], start_loc['lat'])
end_node = ox.distance.nearest_nodes(G, end_loc['lon'], end_loc['lat'])

#calculations 
route = astar(G, start_node, end_node)
route_coords = [(G.nodes[n]['y'], G.nodes[n]['x']) for n in route]  # lat, lon
real_distance = sum(G[u][v][0]['length'] for u, v in zip(route[:-1], route[1:]))
elevation_gain = sum(elevation_difference(u, v) for u, v in zip(route[:-1], route[1:]))

if elevation_gain < 5:
    difficulty = "Easy"
elif elevation_gain < 15:
    difficulty = "Moderate"
else:
    difficulty = "Difficult"

calories_burned = {}
for speed in [4, 6, 8]:
    time_h = real_distance / 1000 / speed
    met = 3.5 if speed == 4 else 7 if speed == 6 else 10
    cal = met * DEFAULT_WEIGHT_KG * time_h
    calories_burned[f"{speed} km/h"] = round(cal, 2)

lats, lons = zip(*route_coords)

fig = go.Figure()

for loc in locations:
    fig.add_scattermapbox(
        lat=[loc['lat']],
        lon=[loc['lon']],
        mode='markers+text',
        marker=dict(size=8, color='red'),
        text=[loc['name']],
        textposition="top center",
        name=loc['name'],
        showlegend=False,
        hoverinfo='text'  
    )

# Route line
fig.add_trace(go.Scattermapbox(
    lat=lats,
    lon=lons,
    mode='lines+markers',
    line=dict(width=4, color='blue'),
    marker=dict(size=6, color='red'),
    name='A* Route'
))

# Start point
fig.add_trace(go.Scattermapbox(
    lat=[start_loc['lat']],
    lon=[start_loc['lon']],
    mode='markers+text',
    marker=dict(size=12, color='green'),
    text=[start_point],
    textposition="top center",
    name='Start'
))

# End point
fig.add_trace(go.Scattermapbox(
    lat=[end_loc['lat']],
    lon=[end_loc['lon']],
    mode='markers+text',
    marker=dict(size=12, color='green'),
    text=[end_point],
    textposition="top center",
    name='End'
))

# Map layout
fig.update_layout(
    mapbox_style="open-street-map",
    mapbox_zoom=16,
    mapbox_center={"lat": center_lat, "lon": center_lon},
    margin={"r":0,"t":0,"l":0,"b":0},
    title="A* Route Between UGM Locations"
)

# Use the new method to adjust map zoom and center
# center_lat, center_lon, zoom_level = adjust_map_zoom_and_center(lats, lons, real_distance, total_distance)
fig.update_layout(
    mapbox=dict(
        center=dict(lat=center_lat, lon=center_lon),
        zoom=14
    )
)

fig.update_layout(
    mapbox_style="open-street-map",
    margin={"r":0, "t":0, "l":0, "b":0}
)

# Map Section
with col2:
    st.header("🗺️ Map")
    if submitted:
        st.plotly_chart(fig, use_container_width=True)
    else:
        st.info("Submit the form to see the map.")

# Output Section
with col3:
    st.header("📊 Output")
    if submitted:
        st.success("### Route Summary")
        st.metric(label="Real Distance (km)", value="{:.2f}".format(real_distance/1000) if real_distance else "N/A")
        st.metric(label="Elevation Gain (m)", value=round(elevation_gain, 2))
        st.metric(label="Difficulty", value=difficulty)
        st.write("### Calories Burned")
        for speed, calories in calories_burned.items():
            st.write(f"- **{speed}:** {calories} kcal")
    else:
        st.info("Submit the form to see the output.")

st.markdown("---")
st.caption("Developed with ❤️ for UGM runners.")