import streamlit as st
import pandas as pd
import folium
from streamlit_folium import st_folium


def plot_trip_on_map(trip_id: str, data: dict[str, pd.DataFrame]):
    trips_df = data["trips"]
    shapes_df = data["shapes"]
    stops_times_df = data["stop_times"]
    stops_df = data["stops"]

    # Filter trips and shapes by trip_id
    trip_shapes = trips_df[trips_df["trip_id"] == trip_id]["shape_id"].unique()

    # Get shape points for the trip
    shape_points = shapes_df[shapes_df["shape_id"].isin(trip_shapes)]
    shape_points = shape_points.sort_values(
        by=["shape_pt_sequence"]
    )  # Sort by sequence

    # Extract trips coordinates
    coordinates = shape_points[["shape_pt_lat", "shape_pt_lon"]].values.tolist()

    # Filter stops
    stops_times_trip = stops_times_df[stops_times_df["trip_id"] == trip_id][
        "stop_id"
    ].unique()
    stops_points = stops_df[stops_df["stop_id"].isin(stops_times_trip)]

    # Extract stop coordinates
    stops_coordinates = stops_points[["stop_lat", "stop_lon"]].values.tolist()

    # Create folium map centered at the first point
    if len(coordinates) > 0:
        map_center = coordinates[0]
        map = folium.Map(location=map_center, zoom_start=13)  # type: ignore

        # Plot the shape as a polyline leave for later test because is connecting dots
        #  and I am interesting on the points I receive
        # folium.PolyLine(locations=coordinates, color='blue', weight=5, opacity=0.7).add_to(m)

        # Add markers for each trip point
        for coord in coordinates:
            folium.CircleMarker(location=coord, radius=3, color="red").add_to(map)  # type: ignore

        # Add markers for each stop
        for stop_coord in stops_coordinates:
            folium.CircleMarker(location=stop_coord, radius=3, color="green").add_to(map)  # type: ignore

        # Display map
        st_folium(map, width=2000)
    else:
        st.warning(f"No shape data found for trip {trip_id}.")


@st.cache_data
def load_data() -> dict[str, pd.DataFrame]:
    # read all parquet files
    return {
        "routes": pd.read_parquet("./data/routes.parquet.gzip"),
        "trips": pd.read_parquet("./data/trips.parquet.gzip"),
        "shapes": pd.read_parquet("./data/shapes.parquet.gzip"),
        "stops": pd.read_parquet("./data/stops.parquet.gzip"),
        "stop_times": pd.read_parquet("./data/stop_times.parquet.gzip"),
    }