trainvel/sncfgtfs/models.py

from django.db import models
from django.utils.translation import gettext_lazy as _


class Country(models.TextChoices):
    """
    Country list by ISO 3166-1 alpha-2 code.
    Only countries that are member of the Council of Europe
    are listed for now.
    """
    ALBANIA = "AL", _("Albania")
    ANDORRA = "AD", _("Andorra")
    ARMENIA = "AM", _("Armenia")
    AUSTRIA = "AT", _("Austria")
    AZERBAIJAN = "AZ", _("Azerbaijan")
    BELGIUM = "BE", _("Belgium")
    BOSNIA_AND_HERZEGOVINA = "BA", _("Bosnia and Herzegovina")
    BULGARIA = "BG", _("Bulgaria")
    CROATIA = "HR", _("Croatia")
    CYPRUS = "CY", _("Cyprus")
    CZECH_REPUBLIC = "CZ", _("Czech Republic")
    DENMARK = "DK", _("Denmark")
    ESTONIA = "EE", _("Estonia")
    FINLAND = "FI", _("Finland")
    FRANCE = "FR", _("France")
    GEORGIA = "GE", _("Georgia")
    GERMANY = "DE", _("Germany")
    GREECE = "GR", _("Greece")
    HUNGARY = "HU", _("Hungary")
    ICELAND = "IS", _("Iceland")
    IRELAND = "IE", _("Ireland")
    ITALY = "IT", _("Italy")
    LATVIA = "LV", _("Latvia")
    LIECHTENSTEIN = "LI", _("Liechtenstein")
    LITHUANIA = "LT", _("Lithuania")
    LUXEMBOURG = "LU", _("Luxembourg")
    MALTA = "MT", _("Malta")
    MOLDOVA = "MD", _("Moldova")
    MONACO = "MC", _("Monaco")
    MONTENEGRO = "ME", _("Montenegro")
    NETHERLANDS = "NL", _("Netherlands")
    NORTH_MACEDONIA = "MK", _("North Macedonia")
    NORWAY = "NO", _("Norway")
    POLAND = "PL", _("Poland")
    PORTUGAL = "PT", _("Portugal")
    ROMANIA = "RO", _("Romania")
    SAN_MARINO = "SM", _("San Marino")
    SERBIA = "RS", _("Serbia")
    SLOVAKIA = "SK", _("Slovakia")
    SLOVENIA = "SI", _("Slovenia")
    SPAIN = "ES", _("Spain")
    SWEDEN = "SE", _("Sweden")
    SWITZERLAND = "CH", _("Switzerland")
    TURKEY = "TR", _("Turkey")
    UNITED_KINGDOM = "GB", _("United Kingdom")
    UKRAINE = "UA", _("Ukraine")


class LocationType(models.IntegerChoices):
    STOP_PLATFORM = 0, _("Stop/platform")
    STATION = 1, _("Station")
    ENTRANCE_EXIT = 2, _("Entrance/exit")
    GENERIC_NODE = 3, _("Generic node")
    BOARDING_AREA = 4, _("Boarding area")


class AccessInformation(models.IntegerChoices):
    NO_INFORMATION = 0, _("No information")
    POSSIBLE = 1, _("Possible")
    NOT_POSSIBLE = 2, _("Not possible")


class PickupType(models.IntegerChoices):
    REGULAR = 0, _("Regular")
    NONE = 1, _("None")
    MUST_PHONE_AGENCY = 2, _("Must phone agency")
    MUST_COORDINATE_WITH_DRIVER = 3, _("Must coordinate with driver")


class RouteType(models.IntegerChoices):
    TRAM = 0, _("Tram")
    METRO = 1, _("Metro")
    RAIL = 2, _("Rail")
    BUS = 3, _("Bus")
    FERRY = 4, _("Ferry")
    CABLE_CAR = 5, _("Cable car")
    GONDOLA = 6, _("Gondola")
    FUNICULAR = 7, _("Funicular")


class Direction(models.IntegerChoices):
    OUTBOUND = 0, _("Outbound")
    INBOUND = 1, _("Inbound")


class TransferType(models.IntegerChoices):
    RECOMMENDED = 0, _("Recommended")
    TIMED = 1, _("Timed")
    MINIMUM_TIME = 2, _("Minimum time")
    NOT_POSSIBLE = 3, _("Not possible")


class ExceptionType(models.IntegerChoices):
    ADDED = 1, _("Added")
    REMOVED = 2, _("Removed")


class TripScheduleRelationship(models.IntegerChoices):
    SCHEDULED = 0, _("Scheduled")
    ADDED = 1, _("Added")
    UNSCHEDULED = 2, _("Unscheduled")
    CANCELED = 3, _("Canceled")
    REPLACEMENT = 5, _("Replacement")
    DUPLICATED = 6, _("Duplicated")
    DELETED = 7, _("Deleted")


class StopScheduleRelationship(models.IntegerChoices):
    SCHEDULED = 0, _("Scheduled")
    SKIPPED = 1, _("Skipped")
    NO_DATA = 2, _("No data")
    UNSCHEDULED = 3, _("Unscheduled")


class GTFSFeed(models.Model):
    code = models.CharField(
        primary_key=True,
        max_length=64,
        verbose_name=_("code"),
        help_text=_("Unique code of the feed.")
    )

    name = models.CharField(
        max_length=255,
        verbose_name=_("name"),
        unique=True,
        help_text=_("Full name that describes the feed."),
    )

    country = models.CharField(
        max_length=2,
        verbose_name=_("country"),
        choices=Country,
    )

    feed_url = models.URLField(
        verbose_name=_("feed URL"),
        help_text=_("URL to download the GTFS feed. Must point to a ZIP archive. "
                    "See https://gtfs.org/schedule/ for more information."),
    )

    rt_feed_url = models.URLField(
        verbose_name=_("realtime feed URL"),
        blank=True,
        default="",
        help_text=_("URL to download the GTFS-Realtime feed, in the GTFS-RT format. "
                    "See https://gtfs.org/realtime/ for more information."),
    )

    last_modified = models.DateTimeField(
        verbose_name=_("last modified date"),
        null=True,
        default=None,
    )

    etag = models.CharField(
        max_length=255,
        verbose_name=_("ETag"),
        blank=True,
        default="",
        help_text=_("If applicable, corresponds to the tag of the last downloaded file. "
                    "If it is not modified, the file is the same."),
    )

    def __str__(self):
        return f"{self.name} ({self.code})"

    class Meta:
        verbose_name = _("GTFS feed")
        verbose_name_plural = _("GTFS feeds")
        ordering = ('country', 'name',)
        indexes = (models.Index(fields=['name']),)


class Agency(models.Model):
    id = models.CharField(
        max_length=255,
        primary_key=True,
        verbose_name=_("Agency ID"),
    )

    name = models.CharField(
        max_length=255,
        verbose_name=_("Agency name"),
    )

    url = models.URLField(
        verbose_name=_("Agency URL"),
    )

    timezone = models.CharField(
        max_length=255,
        verbose_name=_("Agency timezone"),
    )

    lang = models.CharField(
        max_length=255,
        verbose_name=_("Agency language"),
        blank=True,
    )

    phone = models.CharField(
        max_length=255,
        verbose_name=_("Agency phone"),
        blank=True,
    )

    email = models.EmailField(
        verbose_name=_("Agency email"),
        blank=True,
    )

    gtfs_feed = models.ForeignKey(
        GTFSFeed,
        on_delete=models.CASCADE,
        verbose_name=_("GTFS feed"),
    )

    def __str__(self):
        return self.name

    class Meta:
        verbose_name = _("Agency")
        verbose_name_plural = _("Agencies")
        ordering = ("name",)
        indexes = (models.Index(fields=['name']), models.Index(fields=['gtfs_feed']),)


class Stop(models.Model):
    id = models.CharField(
        max_length=255,
        primary_key=True,
        verbose_name=_("Stop ID"),
    )

    code = models.CharField(
        max_length=255,
        verbose_name=_("Stop code"),
        blank=True,
    )

    name = models.CharField(
        max_length=255,
        verbose_name=_("Stop name"),
    )

    desc = models.CharField(
        max_length=255,
        verbose_name=_("Stop description"),
        blank=True,
    )

    lon = models.FloatField(
        verbose_name=_("Stop longitude"),
    )

    lat = models.FloatField(
        verbose_name=_("Stop latitude"),
    )

    zone_id = models.CharField(
        max_length=255,
        verbose_name=_("Zone ID"),
        blank=True,
    )

    url = models.URLField(
        verbose_name=_("Stop URL"),
        blank=True,
    )

    location_type = models.IntegerField(
        verbose_name=_("Location type"),
        blank=True,
        choices=LocationType,
        default=LocationType.STOP_PLATFORM,
    )

    parent_station = models.ForeignKey(
        to="Stop",
        on_delete=models.PROTECT,
        verbose_name=_("Parent station"),
        related_name="children",
        blank=True,
        null=True,
    )

    timezone = models.CharField(
        max_length=255,
        verbose_name=_("Stop timezone"),
        blank=True,
    )

    level_id = models.CharField(
        max_length=255,
        verbose_name=_("Level ID"),
        blank=True,
    )

    wheelchair_boarding = models.IntegerField(
        verbose_name=_("Wheelchair boarding"),
        blank=True,
        choices=AccessInformation,
        default=AccessInformation.NO_INFORMATION,
    )

    platform_code = models.CharField(
        max_length=255,
        verbose_name=_("Platform code"),
        blank=True,
    )

    gtfs_feed = models.ForeignKey(
        GTFSFeed,
        on_delete=models.CASCADE,
        verbose_name=_("GTFS feed"),
    )

    @property
    def stop_type(self):
        train_type = self.id.split('StopPoint:OCE')[-1].split('StopArea:OCE')[-1].split('-')[0]
        return train_type

    def __str__(self):
        return f"{self.name} ({self.id})"

    class Meta:
        verbose_name = _("Stop")
        verbose_name_plural = _("Stops")
        ordering = ("id",)
        indexes = (models.Index(fields=['name']),
                   models.Index(fields=['code']),
                   models.Index(fields=['gtfs_feed']),)


class Route(models.Model):
    id = models.CharField(
        max_length=255,
        primary_key=True,
        verbose_name=_("ID"),
    )

    agency = models.ForeignKey(
        to="Agency",
        on_delete=models.CASCADE,
        verbose_name=_("Agency"),
        related_name="routes",
        null=True,
        blank=True,
        default=None,
    )

    short_name = models.CharField(
        max_length=255,
        verbose_name=_("Route short name"),
    )

    long_name = models.CharField(
        max_length=255,
        verbose_name=_("Route long name"),
        blank=True,
    )

    desc = models.CharField(
        max_length=255,
        verbose_name=_("Route description"),
        blank=True,
    )

    type = models.IntegerField(
        verbose_name=_("Route type"),
        choices=RouteType,
    )

    url = models.URLField(
        verbose_name=_("Route URL"),
        blank=True,
    )

    color = models.CharField(
        max_length=255,
        verbose_name=_("Route color"),
        blank=True,
    )

    text_color = models.CharField(
        max_length=255,
        verbose_name=_("Route text color"),
        blank=True,
    )

    gtfs_feed = models.ForeignKey(
        GTFSFeed,
        on_delete=models.CASCADE,
        verbose_name=_("GTFS feed"),
    )

    def __str__(self):
        return self.long_name or self.short_name

    class Meta:
        verbose_name = _("Route")
        verbose_name_plural = _("Routes")
        ordering = ("id",)
        indexes = (models.Index(fields=['gtfs_feed']),)


class Trip(models.Model):
    id = models.CharField(
        max_length=255,
        primary_key=True,
        verbose_name=_("Trip ID"),
    )

    route = models.ForeignKey(
        to="Route",
        on_delete=models.CASCADE,
        verbose_name=_("Route"),
        related_name="trips",
    )

    service = models.ForeignKey(
        to="Calendar",
        on_delete=models.CASCADE,
        verbose_name=_("Service"),
        related_name="trips",
    )

    headsign = models.CharField(
        max_length=255,
        verbose_name=_("Trip headsign"),
        blank=True,
    )

    short_name = models.CharField(
        max_length=255,
        verbose_name=_("Trip short name"),
        blank=True,
    )

    direction_id = models.IntegerField(
        verbose_name=_("Direction"),
        choices=Direction,
        null=True,
    )

    block_id = models.CharField(
        max_length=255,
        verbose_name=_("Block ID"),
        blank=True,
    )

    shape_id = models.CharField(
        max_length=255,
        verbose_name=_("Shape ID"),
        blank=True,
    )

    wheelchair_accessible = models.IntegerField(
        verbose_name=_("Wheelchair accessible"),
        choices=AccessInformation,
        default=AccessInformation.NO_INFORMATION,
        null=True,
    )

    bikes_allowed = models.IntegerField(
        verbose_name=_("Bikes allowed"),
        choices=AccessInformation,
        default=AccessInformation.NO_INFORMATION,
        null=True,
    )

    gtfs_feed = models.ForeignKey(
        GTFSFeed,
        on_delete=models.CASCADE,
        verbose_name=_("GTFS feed"),
    )

    @property
    def origin(self) -> Stop | None:
        return self.stop_times.order_by('stop_sequence').first().stop if self.stop_times.exists() else None

    @property
    def destination(self) -> Stop | None:
        return self.stop_times.order_by('-stop_sequence').first().stop if self.stop_times.exists() else None

    @property
    def departure_time(self):
        if not self.stop_times.exists():
            return _("Unknown")
        dep_time = self.stop_times.order_by('stop_sequence').first().departure_time
        hours = int(dep_time.total_seconds() // 3600)
        minutes = int((dep_time.total_seconds() % 3600) // 60)
        return f"{hours:02}:{minutes:02}"

    @property
    def arrival_time(self):
        if not self.stop_times.exists():
            return _("Unknown")
        arr_time = self.stop_times.order_by('-stop_sequence').first().arrival_time
        hours = int(arr_time.total_seconds() // 3600)
        minutes = int((arr_time.total_seconds() % 3600) // 60)
        return f"{hours:02}:{minutes:02}"

    @property
    def train_type(self):
        if self.gtfs_feed.code == "FR-IDF-TN":
            return self.route.short_name
        else:
            return self.origin.stop_type

    @property
    def train_number(self):
        if self.gtfs_feed.code == "FR-IDF-TN":
            return self.short_name
        else:
            return self.headsign

    @property
    def color(self):
        if self.route.color:
            return self.route.color
        elif self.train_type == "OUIGO":
            return "E60075"
        return "FFFFFF"

    @property
    def text_color(self):
        if self.route.text_color:
            return self.route.text_color
        elif self.train_type == "OUIGO":
            return "FFFFFF"
        elif self.train_type == "TGV INOUI":
            return "9B2743"
        elif self.train_type == "INTER-CITÉS" or self.train_type == "INTER-CITÉS de nuit":
            return "404042"
        return "000000"

    @property
    def origin_destination(self):
        origin = self.origin
        origin = origin.name if origin else _("Unknown")
        destination = self.destination
        destination = destination.name if destination else _("Unknown")
        return f"{origin} {self.departure_time} → {destination} {self.arrival_time}"

    origin_destination.fget.short_description = _("Origin → Destination")

    def __str__(self):
        return self.origin_destination

    class Meta:
        verbose_name = _("Trip")
        verbose_name_plural = _("Trips")
        indexes = (models.Index(fields=['route']), models.Index(fields=['gtfs_feed']),)


class StopTime(models.Model):
    id = models.CharField(
        max_length=255,
        primary_key=True,
        verbose_name=_("ID"),
    )

    trip = models.ForeignKey(
        to="Trip",
        on_delete=models.CASCADE,
        verbose_name=_("Trip"),
        related_name="stop_times",
    )

    arrival_time = models.DurationField(
        verbose_name=_("Arrival time"),
    )

    departure_time = models.DurationField(
        verbose_name=_("Departure time"),
    )

    stop = models.ForeignKey(
        to="Stop",
        on_delete=models.CASCADE,
        verbose_name=_("Stop ID"),
        related_name="stop_times",
    )

    stop_sequence = models.IntegerField(
        verbose_name=_("Stop sequence"),
    )

    stop_headsign = models.CharField(
        max_length=255,
        verbose_name=_("Stop headsign"),
        blank=True,
    )

    pickup_type = models.IntegerField(
        verbose_name=_("Pickup type"),
        choices=PickupType,
        default=PickupType.REGULAR,
        null=True,
    )

    drop_off_type = models.IntegerField(
        verbose_name=_("Drop off type"),
        choices=PickupType,
        default=PickupType.REGULAR,
        null=True,
    )

    timepoint = models.BooleanField(
        verbose_name=_("Timepoint"),
        default=True,
        null=True,
    )

    @property
    def pretty_arrival_time(self):
        seconds = self.arrival_time.total_seconds()
        hours = int(seconds // 3600) % 24
        minutes = int((seconds % 3600) // 60)
        return f"{hours:02}:{minutes:02}"

    @property
    def pretty_departure_time(self):
        seconds = self.departure_time.total_seconds()
        hours = int(seconds // 3600) % 24
        minutes = int((seconds % 3600) // 60)
        return f"{hours:02}:{minutes:02}"

    def __str__(self):
        return f"{self.stop.name} - {self.trip_id}"

    class Meta:
        verbose_name = _("Stop time")
        verbose_name_plural = _("Stop times")
        indexes = (models.Index(fields=['stop']), models.Index(fields=['trip']),)


class Calendar(models.Model):
    id = models.CharField(
        max_length=255,
        primary_key=True,
        verbose_name=_("Service ID"),
    )

    monday = models.BooleanField(
        verbose_name=_("Monday"),
    )

    tuesday = models.BooleanField(
        verbose_name=_("Tuesday"),
    )

    wednesday = models.BooleanField(
        verbose_name=_("Wednesday"),
    )

    thursday = models.BooleanField(
        verbose_name=_("Thursday"),
    )

    friday = models.BooleanField(
        verbose_name=_("Friday"),
    )

    saturday = models.BooleanField(
        verbose_name=_("Saturday"),
    )

    sunday = models.BooleanField(
        verbose_name=_("Sunday"),
    )

    start_date = models.DateField(
        verbose_name=_("Start date"),
    )

    end_date = models.DateField(
        verbose_name=_("End date"),
    )

    gtfs_feed = models.ForeignKey(
        GTFSFeed,
        on_delete=models.CASCADE,
        verbose_name=_("GTFS feed"),
    )

    def __str__(self):
        return self.id

    class Meta:
        verbose_name = _("Calendar")
        verbose_name_plural = _("Calendars")
        ordering = ("id",)
        indexes = (models.Index(fields=['gtfs_feed']),)


class CalendarDate(models.Model):
    id = models.CharField(
        max_length=255,
        primary_key=True,
        verbose_name=_("ID"),
    )

    service = models.ForeignKey(
        to="Calendar",
        on_delete=models.CASCADE,
        verbose_name=_("Service"),
        related_name="dates",
    )

    date = models.DateField(
        verbose_name=_("Date"),
    )

    exception_type = models.IntegerField(
        verbose_name=_("Exception type"),
        choices=ExceptionType,
    )

    def __str__(self):
        return f"{self.service.id} - {self.date} - {self.exception_type}"

    class Meta:
        verbose_name = _("Calendar date")
        verbose_name_plural = _("Calendar dates")
        ordering = ("id",)
        indexes = (models.Index(fields=['service']), models.Index(fields=['date']),)


class Transfer(models.Model):
    id = models.CharField(
        max_length=255,
        primary_key=True,
        verbose_name=_("ID"),
    )

    from_stop = models.ForeignKey(
        to="Stop",
        on_delete=models.CASCADE,
        verbose_name=_("From stop"),
        related_name="transfers_from",
    )

    to_stop = models.ForeignKey(
        to="Stop",
        on_delete=models.CASCADE,
        verbose_name=_("To stop"),
        related_name="transfers_to",
    )

    transfer_type = models.IntegerField(
        verbose_name=_("Transfer type"),
        choices=TransferType,
        default=TransferType.RECOMMENDED,
    )

    min_transfer_time = models.IntegerField(
        verbose_name=_("Minimum transfer time"),
        blank=True,
    )

    class Meta:
        verbose_name = _("Transfer")
        verbose_name_plural = _("Transfers")
        ordering = ("id",)


class FeedInfo(models.Model):
    publisher_name = models.CharField(
        max_length=255,
        verbose_name=_("Feed publisher name"),
    )

    publisher_url = models.URLField(
        verbose_name=_("Feed publisher URL"),
    )

    lang = models.CharField(
        max_length=255,
        verbose_name=_("Feed language"),
    )

    start_date = models.DateField(
        verbose_name=_("Feed start date"),
    )

    end_date = models.DateField(
        verbose_name=_("Feed end date"),
    )

    version = models.CharField(
        max_length=255,
        verbose_name=_("Feed version"),
    )

    gtfs_feed = models.ForeignKey(
        GTFSFeed,
        on_delete=models.CASCADE,
        verbose_name=_("GTFS feed"),
    )

    class Meta:
        verbose_name = _("Feed info")
        verbose_name_plural = _("Feed infos")
        ordering = ("publisher_name",)
        indexes = (models.Index(fields=['gtfs_feed']),)


class TripUpdate(models.Model):
    trip = models.OneToOneField(
        to="Trip",
        on_delete=models.CASCADE,
        verbose_name=_("Trip"),
        related_name="update",
        primary_key=True,
    )

    start_date = models.DateField(
        verbose_name=_("Start date"),
    )

    start_time = models.TimeField(
        verbose_name=_("Start time"),
    )

    schedule_relationship = models.IntegerField(
        verbose_name=_("Schedule relationship"),
        choices=TripScheduleRelationship,
        default=TripScheduleRelationship.SCHEDULED,
    )

    def __str__(self):
        return str(self.trip)

    class Meta:
        verbose_name = _("Trip update")
        verbose_name_plural = _("Trip updates")
        ordering = ("start_date", "trip",)
        unique_together = ("trip", "start_date", "start_time",)
        indexes = (models.Index(fields=['trip']),)


class StopTimeUpdate(models.Model):
    trip_update = models.ForeignKey(
        to="TripUpdate",
        on_delete=models.CASCADE,
        verbose_name=_("Trip update"),
        related_name="stop_time_updates",
    )

    stop_time = models.OneToOneField(
        to="StopTime",
        on_delete=models.CASCADE,
        verbose_name=_("Stop time"),
        related_name="update",
        primary_key=True,
    )

    arrival_delay = models.DurationField(
        verbose_name=_("Arrival delay"),
    )

    arrival_time = models.DateTimeField(
        verbose_name=_("Arrival time"),
    )

    departure_delay = models.DurationField(
        verbose_name=_("Departure delay"),
    )

    departure_time = models.DateTimeField(
        verbose_name=_("Departure time"),
    )

    schedule_relationship = models.IntegerField(
        verbose_name=_("Schedule relationship"),
        choices=StopScheduleRelationship,
        default=StopScheduleRelationship.SCHEDULED,
    )

    def __str__(self):
        return str(self.trip_update)

    class Meta:
        verbose_name = _("Stop time update")
        verbose_name_plural = _("Stop time updates")
        ordering = ("trip_update", "stop_time",)
        unique_together = ("trip_update", "stop_time",)
        indexes = (models.Index(fields=['trip_update']), models.Index(fields=['stop_time']),)