nyc_taxi_record

"""
nyc_taxi_record.py

Build a compact NYC Yellow Taxi toy dataset.

Source:
- NYC Taxi and Limousine Commission Yellow Taxi Trip Records.

Output:
- assets/toy/transport/nyc_taxi_record.csv
"""

from pathlib import Path

import numpy as np
import pandas as pd


SOURCE_URL = (
    "https://d37ci6vzurychx.cloudfront.net/trip-data/"
    "yellow_tripdata_2024-10.parquet"
)
OUT = Path(__file__).with_name("nyc_taxi_record.csv")
SAMPLE_ROWS = 1000
RANDOM_STATE = 42

RAW_COLUMNS = [
    "tpep_pickup_datetime",
    "tpep_dropoff_datetime",
    "passenger_count",
    "trip_distance",
    "RatecodeID",
    "payment_type",
    "fare_amount",
    "tip_amount",
    "tolls_amount",
    "total_amount",
]

PAYMENT_TYPES = {
    1: "credit_card",
    2: "cash",
    3: "no_charge",
    4: "dispute",
    5: "unknown",
    6: "voided_trip",
}

RATE_CODES = {
    1: "standard",
    2: "jfk",
    3: "newark",
    4: "nassau_westchester",
    5: "negotiated",
    6: "group_ride",
    99: "unknown",
}


def read_source() -> pd.DataFrame:
    print(f"Reading {SOURCE_URL}")
    return pd.read_parquet(SOURCE_URL, columns=RAW_COLUMNS)


def clean_trips(raw: pd.DataFrame) -> pd.DataFrame:
    data = raw.copy()
    data["pickup_datetime"] = pd.to_datetime(data["tpep_pickup_datetime"])
    data["dropoff_datetime"] = pd.to_datetime(data["tpep_dropoff_datetime"])

    data["trip_duration_min"] = (
        data["dropoff_datetime"] - data["pickup_datetime"]
    ).dt.total_seconds() / 60

    data.rename(
        columns={
            "passenger_count": "passenger_count",
            "trip_distance": "trip_distance_miles",
            "RatecodeID": "rate_code_id",
            "fare_amount": "fare_amount_usd",
            "tip_amount": "tip_amount_usd",
            "tolls_amount": "tolls_amount_usd",
            "total_amount": "total_amount_usd",
        },
        inplace=True,
    )

    data = data.loc[
        (data["pickup_datetime"] >= "2024-10-01")
        & (data["pickup_datetime"] < "2024-11-01")
        & (data["trip_duration_min"].between(1, 120))
        & (data["trip_distance_miles"].between(0.1, 50))
        & (data["fare_amount_usd"] > 0)
        & (data["total_amount_usd"] > 0)
        & (data["passenger_count"].between(1, 6))
    ].copy()

    data["avg_speed_mph"] = (
        data["trip_distance_miles"] / (data["trip_duration_min"] / 60)
    )
    data = data.loc[data["avg_speed_mph"].between(1, 70)].copy()

    pretax_tip_base = data["total_amount_usd"] - data["tip_amount_usd"]
    data["tip_pct"] = np.where(
        pretax_tip_base > 0,
        data["tip_amount_usd"] / pretax_tip_base,
        np.nan,
    )
    data = data.loc[data["tip_pct"].between(0, 1.5) | data["tip_pct"].isna()].copy()

    data["pickup_hour"] = data["pickup_datetime"].dt.hour
    data["pickup_weekday"] = data["pickup_datetime"].dt.day_name()
    data["payment_type"] = data["payment_type"].map(PAYMENT_TYPES).fillna("unknown")
    data["rate_code"] = data["rate_code_id"].map(RATE_CODES).fillna("unknown")
    data["is_airport_rate"] = data["rate_code"].isin(["jfk", "newark"])
    data["passenger_count"] = data["passenger_count"].astype(int)

    money_columns = [
        "fare_amount_usd",
        "tip_amount_usd",
        "tolls_amount_usd",
        "total_amount_usd",
    ]
    data[money_columns] = data[money_columns].round(2)
    data["trip_distance_miles"] = data["trip_distance_miles"].round(2)
    data["trip_duration_min"] = data["trip_duration_min"].round(1)
    data["avg_speed_mph"] = data["avg_speed_mph"].round(1)
    data["tip_pct"] = data["tip_pct"].round(3)

    columns = [
        "pickup_hour",
        "pickup_weekday",
        "passenger_count",
        "trip_distance_miles",
        "trip_duration_min",
        "avg_speed_mph",
        "fare_amount_usd",
        "tip_amount_usd",
        "tip_pct",
        "tolls_amount_usd",
        "total_amount_usd",
        "payment_type",
        "rate_code",
        "is_airport_rate",
    ]
    return data[columns].reset_index(drop=True)


def main() -> None:
    raw = read_source()
    cleaned = clean_trips(raw)

    if len(cleaned) < SAMPLE_ROWS:
        raise RuntimeError(
            f"Only {len(cleaned)} rows remain after cleaning; "
            f"need at least {SAMPLE_ROWS}."
        )

    sample = cleaned.sample(n=SAMPLE_ROWS, random_state=RANDOM_STATE).reset_index(
        drop=True
    )
    sample.insert(0, "trip_id", [f"yt_202410_{i:04d}" for i in range(1, len(sample) + 1)])

    sample.to_csv(OUT, index=False)
    print(f"Done: {len(sample)} rows saved to {OUT}")


if __name__ == "__main__":
    main()