import requests

url = "https://data.calgary.ca/download/npk7-z3bj/application%2Fx-zip-compressed"
r = requests.get(url)
inpath = "example_nb_latest_gtfs.zip"
with open(inpath, 'wb') as f:
    f.write(r.content)
del r

import datetime
import partridge as ptg

# read in the gtfs feed as a partridge object
service_ids_by_date = ptg.read_service_ids_by_date(inpath)
sids = [service_ids for date, service_ids in service_ids_by_date.items() if date == datetime.date(2021, 12, 20)]
gtfs = ptg.load_feed(inpath, {"trips.txt": {"service_id": sids}})

from typing import Union
import numpy as np
import sys
import pyproj
import warnings
import pandas as pd
from pyproj import CRS
import math
import geopandas as gpd
from shapely.geometry import LineString, Point

def extract_the_line(sub: pd.DataFrame) -> LineString:
    sub = sub.sort_values(by="shape_pt_sequence", ascending=True)
    return LineString(zip(sub.shape_pt_lon, sub.shape_pt_lat))


def wgs_to_utm(gdf: Union[gpd.GeoSeries, gpd.GeoDataFrame]) -> Union[gpd.GeoSeries, gpd.GeoDataFrame]:
    """Take a 4236-proj geoseries and return cast to local meter-based."""
    gdf.crs = CRS.from_epsg(4326)
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        xs = gdf.centroid.x
    utm_zone = int(math.floor((np.mean(xs) + 180) / 6) + 1)
    utm_crs = "+proj=utm +zone={} +ellps=WGS84 +datum=WGS84 +units=m +no_defs".format(utm_zone)
    return gdf.to_crs(crs=utm_crs)


# quickly check that it looks like we have a whole bus system here
s = gtfs.shapes.copy()
gds = gpd.GeoSeries(s.groupby("shape_id").apply(extract_the_line))
gds = wgs_to_utm(gds)
del s

ax = gds.plot(figsize=(18,18), lw=1.5, ec="blue", alpha=0.3)
ax.set_axis_off()
ax.set_title("Calgary GTFS system represented via shapes file")


# next, we want to get a popular stop and mark that on the map, too
s = gtfs.stops.copy()
popular_stop = "3392"
p = Point(s.loc[s.stop_id == popular_stop, ["stop_lon", "stop_lat"]].squeeze().tolist())
del s

# mark on the map w/ a 1000 meter-diameter dot
wgs_to_utm(gpd.GeoSeries([p])).buffer(500).plot(ax=ax, color="red")

import time
import pytz

def _as_secs(at: str) -> int:
    """Given a time string (ex: 06:16:00) convert to int of seconds."""
    bkts = list(map(int, at.split(":")))
    h = bkts[0]
    m = bkts[1]
    s = bkts[2]
    return (h * 60 * 60) + (m * 60) + s

curr_time_from_secs = int(time.time())
curr_time_raw = pytz.utc.localize(datetime.datetime.now(), is_dst=None).astimezone(pytz.timezone("America/Denver")).strftime("%H:%M:%S")
print(f"Evaluating at this time: {curr_time_raw}")

current_time = _as_secs(curr_time_raw)  # time in calgary
look_fwd_n_min = 20  # how many minutes ahead you want to include in the query

st = gtfs.stop_times.copy()
st = st.loc[st.stop_id == popular_stop]
st = st.sort_values(by="arrival_time", ascending=True)

# get the next N trips coming to this stop within the next 20 min
n_min_from_current = current_time + (look_fwd_n_min * 60)
st = st[(st.arrival_time > current_time) & (st.arrival_time < n_min_from_current)].reset_index(drop=True)

# get the related trip ids
tids = st.trip_id.unique()
ts = gtfs.trips.copy()
ts = ts.loc[ts.trip_id.isin(tids)]

# get the related route attributes
rs = gtfs.routes.copy()
rs = rs.loc[rs.route_id.isin(ts.route_id.unique())]

# ...and merge those onto the trip deets
ts = pd.merge(ts, rs, on="route_id", how="left").set_index("trip_id")
del rs

print(f"{len(st)} Upcoming scheduled arrivals in the next {look_fwd_n_min} minutes:")
for i, row in st.iterrows():
    # you could imagine returning this object instead to a client that would show user this information
    coming_in = (row.arrival_time - current_time) / 60
    t = ts.loc[row.trip_id]
    print(f"\t{str(i+1).zfill(2)}: {t.route_long_name} arriving in {math.ceil(coming_in)} minutes")

import json
import time
from google.transit import gtfs_realtime_pb2
from google.protobuf.json_format import MessageToJson
import pytz
import requests

def get_ref_stop_sequence(tid: str, stop_id: str) -> int:
    """Get the stop sequence for the popular stop as a reference."""
    st = gtfs.stop_times.copy()
    # note this is an incredibly slow way to do this, just providing quick example
    st = st[st.trip_id == tid]
    if stop_id in st.stop_id.tolist():
        ref_stop_seq = st[st.stop_id == stop_id].squeeze().stop_sequence
        return ref_stop_seq
    return None


# get the latest trip updates from real-time
resp = requests.get("https://data.calgary.ca/download/gs4m-mdc2/application%2Foctet-stream")
feed = gtfs_realtime_pb2.FeedMessage()
feed.ParseFromString(resp.content)

# filter to just the trips we want to review
want_trip_ids = ts.index.tolist()

# we will parse the latest live traffic updates data
loaded = [json.loads(MessageToJson(entity)) for entity in feed.entity]

# first let's toss anything that is not in this trip
loaded = [l for l in loaded if l["tripUpdate"]["trip"]["tripId"] in want_trip_ids]

# extract list of reference stop sequences
ref_stop_seqs = {
    tid: get_ref_stop_sequence(tid, popular_stop)
    for tid in set([l["tripUpdate"]["trip"]["tripId"] for l in loaded])}
    
# further prune to only non-None tids
loaded = [l for l in loaded if ref_stop_seqs[l["tripUpdate"]["trip"]["tripId"]]]
  
parsed_list = []
for as_json in loaded:
    curr_tid = as_json["tripUpdate"]["trip"]["tripId"]
    suds = as_json["tripUpdate"]["stopTimeUpdate"]
    for s in suds:
        s["trip_id"] = curr_tid
    parsed_list.extend(suds)
    
# filter to just those that are at stops before ar at the target stop
parsed_list = [p for p in parsed_list if p["stopSequence"] <= ref_stop_seqs[curr_tid]]

print("Upcoming arrivals from live GTFS-RT feed")
for p in parsed_list:
    # you could imagine merging this with the scheduled times report
    try:
        coming_in = (int(p["departure"]["time"]) - curr_time_from_secs)/60
    except KeyError as e:
        coming_in = 999
    if look_fwd_n_min > coming_in and coming_in > 0:
        t = ts.loc[row.trip_id]
        print(f"\t{str(i+1).zfill(2)}: {t.route_long_name} arriving in {math.ceil(coming_in)} minutes")