import streamlit as st
from pyproj import CRS
import pandas as pd
import geopandas as gpd
from shapely.geometry import LineString, Point
import folium
from streamlit_folium import st_folium
import networkx as nx
from collections import defaultdict  # Add this line if missing
import os, sys
import traceback

print(">>> Ejecutando archivo:", os.path.abspath(__file__))
print("Python ejecutado:", sys.executable)

# Cargar shapes.txt manteniendo la precisión
shapes = pd.read_csv('shapes.txt')
routes = pd.read_csv('routes.txt')
stops = pd.read_csv('stops.txt')
stop_times = pd.read_csv('stop_times.txt')
trips = pd.read_csv('trips.txt')


# Convertir puntos a geometría
shapes_gdf = (
    shapes.sort_values(["shape_id", "shape_pt_sequence"])
    .groupby("shape_id")
    .apply(lambda x: LineString(zip(x.shape_pt_lon, x.shape_pt_lat)))
    .reset_index()
)

shapes_gdf.columns = ["shape_id", "geometry"]
crs_obj = CRS.from_epsg(4326)  # Crea el objeto CRS correctamente
shapes_gdf = gpd.GeoDataFrame(shapes_gdf, geometry="geometry", crs=crs_obj)

@st.cache_resource
def build_graph(routes, stop_times, trips):
    route_short = routes.set_index('route_id')['route_short_name'].to_dict()
    stops_routes = defaultdict(set)
    G = nx.DiGraph()

    group_cols = ['route_id', 'shape_id']
    trips_grouped = trips.groupby(group_cols)

    transfer_penalty = 20

    for name, group in trips_grouped:
        if group.empty:
            continue
        route_id, shape_id = name
        trip_id_sample = group['trip_id'].iloc[0]

        stops_trip = stop_times[stop_times['trip_id'] == trip_id_sample].sort_values('stop_sequence')
        stop_ids = stops_trip['stop_id'].tolist()

        # Assign stops → route
        for stop_id in stop_ids:
            stops_routes[stop_id].add(route_id)

        # Add ride edges
        for i in range(len(stop_ids) - 1):
            s1 = stop_ids[i]
            s2 = stop_ids[i + 1]
            G.add_edge((s1, route_id), (s2, route_id), weight=1)

    # Add transfer edges
    for stop_id, routes_set in stops_routes.items():
        routes_list = list(routes_set)
        for i in range(len(routes_list)):
            for j in range(len(routes_list)):
                if i != j:
                    G.add_edge((stop_id, routes_list[i]), (stop_id, routes_list[j]), weight=transfer_penalty)

    return G, stops_routes, route_short


G, stops_routes, route_short = build_graph(routes, stop_times, trips)

# ============================
#       STREAMLIT UI
# ============================

st.title("🚍 Planificador inteligente — TransMilenio")
st.write("Esta app detecta si la ruta **en la que ya vas** te sirve para llegar al destino, y sugiere transbordos si es necesario.")

# ----------------------------
# SELECT ROUTE
# ----------------------------
st.header("1️⃣ ¿En qué ruta vas?")
ruta_seleccionada = st.selectbox("Selecciona tu ruta", routes.route_short_name.unique())

route_id = routes.loc[routes.route_short_name == ruta_seleccionada, "route_id"].iloc[0]
trips_ruta = trips[trips.route_id == route_id]

if trips_ruta.empty:
    st.error("No se encontraron viajes para esta ruta.")
    st.stop()

trip_ids = trips_ruta['trip_id'].unique()
all_stops_route = stop_times[stop_times['trip_id'].isin(trip_ids)].merge(stops, on="stop_id")['stop_name'].unique()

# ----------------------------
# ORIGIN
# ----------------------------
st.header("2️⃣ ¿Dónde estás?")
modo_origen = st.radio("Selecciona cómo indicar dónde estás:", ["Parada", "Coordenadas"])

if modo_origen == "Parada":
    parada_origen = st.selectbox("Selecciona tu parada actual", all_stops_route)
    stop_actual = stops[stops.stop_name == parada_origen].iloc[0]
    current_stop_id = stop_actual['stop_id']

else:
    lat = st.number_input("Latitud", value=4.65)
    lon = st.number_input("Longitud", value=-74.1)

    all_stops_route_df = stop_times[stop_times['trip_id'].isin(trip_ids)].merge(stops, on="stop_id").drop_duplicates('stop_id')
    stops_route_gdf = gpd.GeoDataFrame(
        all_stops_route_df,
        geometry=gpd.points_from_xy(all_stops_route_df.stop_lon, all_stops_route_df.stop_lat),
        crs=4326
    )

    user_point = Point(lon, lat)
    distances = stops_route_gdf.geometry.distance(user_point)
    nearest_idx = distances.argmin()

    stop_actual = stops_route_gdf.iloc[nearest_idx]
    current_stop_id = stop_actual['stop_id']

    st.write(f"Parada más cercana detectada: **{stop_actual.stop_name}**")

# ----------------------------
# DESTINATION
# ----------------------------
st.header("3️⃣ ¿A dónde vas?")
destino = st.text_input("Escribe la parada o dirección de destino")
calcular = st.button("Calcular ruta")

# ============================
#  ROUTING ENGINE
# ============================

if calcular:

    # Find trips with this stop
    relevant_trips = stop_times[(stop_times['stop_id'] == current_stop_id) &
                                (stop_times['trip_id'].isin(trip_ids))]['trip_id'].unique()

    if len(relevant_trips) == 0:
        st.error("La parada actual no está en esta ruta.")
        st.stop()

    trip_id = relevant_trips[0]
    stops_trip = stop_times[stop_times.trip_id == trip_id].merge(stops, on="stop_id").sort_values('stop_sequence')

    # Current sequence
    current_seq = stops_trip[stops_trip.stop_id == current_stop_id]['stop_sequence'].iloc[0]

    # Destination match
    destino_results = stops[stops.stop_name.str.contains(destino, case=False, na=False)]
    if destino_results.empty:
        st.error("No se encontró la parada de destino. Intenta con otro nombre.")
        st.stop()

    destino_stop = destino_results.iloc[0]
    destination_id = destino_stop['stop_id']
    st.write(f"Destino interpretado como: **{destino_stop.stop_name}**")

    # ------------
    # DIRECT ROUTE
    # ------------
    destino_in_trip = stops_trip[(stops_trip.stop_id == destination_id) &
                                 (stops_trip.stop_sequence > current_seq)]

    if not destino_in_trip.empty:
        dest_seq = destino_in_trip['stop_sequence'].iloc[0]
        num_paradas = dest_seq - current_seq
        st.success(
            f"Esta ruta **SÍ** te sirve directamente.\n\n"
            f"Debes bajarte en **{num_paradas} paradas** en **{destino_stop.stop_name}**."
        )
        parada_destino = destino_stop

    else:
        # ------------
        # MULTI-ROUTE GRAPH SEARCH
        # ------------
        source = (current_stop_id, route_id)
        targets = [(destination_id, r) for r in stops_routes[destination_id]]

        paths = {}
        for t in targets:
            try:
                path = nx.shortest_path(G, source, t, weight="weight")
                cost = nx.shortest_path_length(G, source, t, weight="weight")
                paths[cost] = path
            except nx.NetworkXNoPath:
                pass

        if not paths:
            st.warning("No se encontró ninguna ruta con transbordos disponibles.")
            parada_destino = None
        else:
            min_cost = min(paths.keys())
            best_path = paths[min_cost]

            legs = []
            current_route = best_path[0][1]
            start_id = best_path[0][0]
            leg_stops = 0

            for i in range(1, len(best_path)):
                nxt_stop, nxt_route = best_path[i]
                if nxt_route != current_route:
                    from_name = stops.loc[stops.stop_id == start_id, "stop_name"].iloc[0]
                    to_name = stops.loc[stops.stop_id == best_path[i-1][0], "stop_name"].iloc[0]
                    legs.append({
                        "route": route_short[current_route],
                        "from": from_name,
                        "to": to_name,
                        "num_paradas": leg_stops
                    })
                    current_route = nxt_route
                    start_id = best_path[i-1][0]
                    leg_stops = 0
                leg_stops += 1

            from_name = stops.loc[stops.stop_id == start_id, "stop_name"].iloc[0]
            to_name = stops.loc[stops.stop_id == best_path[-1][0], "stop_name"].iloc[0]
            legs.append({
                "route": route_short[current_route],
                "from": from_name,
                "to": to_name,
                "num_paradas": leg_stops
            })

            texto = "### 🚏 Ruta recomendada:\n"
            for i, leg in enumerate(legs):
                if i == 0:
                    texto += f"- Vas en **{leg['route']}** desde **{leg['from']}**, bájate en **{leg['num_paradas']}** paradas en **{leg['to']}**.\n"
                else:
                    texto += f"- Luego, en **{leg['from']}**, toma **{leg['route']}** y bájate en **{leg['num_paradas']}** paradas en **{leg['to']}**.\n"

            st.success(texto)
            parada_destino = destino_stop

    # ============================
    #  MAP RENDER
    # ============================

    st.header("🗺️ Mapa de tu ruta y posición")

    m = folium.Map(
        location=[stop_actual["stop_lat"], stop_actual["stop_lon"]],
        zoom_start=13
    )

    folium.Marker(
        [stop_actual["stop_lat"], stop_actual["stop_lon"]],
        tooltip="Estás aquí",
        icon=folium.Icon(color="blue")
    ).add_to(m)

    # Draw shape
    shape_id = trips.loc[trips.trip_id == trip_id, "shape_id"].iloc[0]
    shape_geom = shapes_gdf.loc[shapes_gdf.shape_id == shape_id, "geometry"].iloc[0]

    folium.PolyLine(
        locations=[(lat, lon) for lon, lat in zip(shape_geom.coords.xy[0], shape_geom.coords.xy[1])],
        weight=4,
        color="red"
    ).add_to(m)

    if parada_destino is not None:
        folium.Marker(
            [parada_destino.stop_lat, parada_destino.stop_lon],
            tooltip=f"Destino: {parada_destino.stop_name}",
            icon=folium.Icon(color="green")
        ).add_to(m)

    st_folium(m, width=700, height=500)