import pandas as pd

# Load GTFS files
print("Loading GTFS files...")
stop_times = pd.read_csv("stop_times.txt")
stops = pd.read_csv("stops.txt")
print(f"Loaded {len(stop_times)} stop times and {len(stops)} stops.")

# Generate straight-line shapes from stop sequences
shapes_data = []
processed_trips = 0

# Process each trip
print("Generating shapes.txt...")
for trip_id, group in stop_times.groupby("trip_id"):
    group = group.sort_values("stop_sequence")
    
    # Extract coordinates for each stop in the trip
    for seq, row in enumerate(group.itertuples(), start=1):
        stop = stops[stops["stop_id"] == row.stop_id]
        if not stop.empty:
            lat, lon = stop["stop_lat"].values[0], stop["stop_lon"].values[0]
            shapes_data.append([trip_id, lat, lon, seq])

    processed_trips += 1
    if processed_trips % 100 == 0:  # Log every 100 trips
        print(f"Processed {processed_trips} trips...")

# Convert to DataFrame and save as shapes.txt
shapes_df = pd.DataFrame(shapes_data, columns=["shape_id", "shape_pt_lat", "shape_pt_lon", "shape_pt_sequence"])
shapes_df.to_csv("shapes.txt", index=False)

print(f"shapes.txt generated successfully! Processed {processed_trips} trips.")