buses.py

#!/usr/bin/env python3
import enum, json
from sys import stderr
from datetime import date, time, datetime, timedelta
from copy import copy
from misc import *
from geometry import *
Suunta = enum.Enum('Suunta', [('Taaksepäin', 0), ('Eteenpäin', 1)])

def muunna_ajovuoro_tunniste(tunniste):
	return tunniste

class Ajovuoro:
	def __init__(self, tunniste, linja, palvelu, kyltti, suunta):
		self.tunniste, self.linja, self.palvelu, self.kyltti, self.suunta = tunniste, linja, \
			palvelu, kyltti, suunta
		self.reitti = []
		self.nimi = muunna_ajovuoro_tunniste(tunniste)
	def __repr__(self):
		return 'ajot[%r]' % self.nimi
	def pysäkkiReitillä(self, pysäkki):
		for pysähdys in self.reitti:
			if pysähdys.pysäkki is pysäkki:
				return pysähdys
		else:
			return None
	def ajetaan_päivänä(self, päivä):
		try:
			return self.palvelu in palvelut_per_päivä[päivä]
		except KeyError:
			return False
	def suppea_reitti(self, pysäkistä = None):
		if pysäkistä and pysäkistä in self.reitti:
			reitti = copy(self.reitti)
			reitti = reitti[reitti.index(pysäkistä):]
		else:
			reitti = self.reitti
		käytetyt_alueet = set()
		tulos = []
		for pysähdys in reitti:
			pysäkki = pysähdys.pysäkki
			if pysäkki.alue and pysäkki.alue not in käytetyt_alueet:
				käytetyt_alueet.add(pysäkki.alue)
				tulos.append(pysäkki.alue)
		return tulos

class Linja:
	def __init__(self, tietue):
		self.tunniste = tietue['route_id']
		self.viite = tietue['route_short_name']
		self.selite = tietue['route_long_name']
	def __repr__(self):
		return 'linjat[%r]' % self.viite

class Palvelu:
	def __init__(self, tunniste):
		self.tunniste = tunniste
		self.päivät = set()
	def __repr__(self):
		return 'palvelut[%r]' % self.tunniste

class Pysäkki:
	def __init__(self, tunniste, nimi, sijainti):
		self.tunniste, self.nimi, self.sijainti = tunniste, nimi, sijainti
		self.cluster = None
		self.pairs = set() # samannimiset lähellä olevat pysäkit
	def __repr__(self):
		return 'pysäkit[%r]' % self.tunniste
	def schedule(self, max_amount = 50):
		'''
			Hakee tämän pysäkin seuraavat `määrä` lähtöä. Päätepysäkille saapuvia busseja ei
			lasketa. Palauttaa pysähdykset listana jossa alkiot ovat muotoa (aika, pysähdys),
			jossa:
			- `aika` on saapumishetki muotoa datetime ja
			- `pysähdys` on vastaava Pysähdys olio.
			
			Mikäli pysäkille ei ole määrätty riittävästi pysähdyksiä kalenterissa, tuloslista
			jää alimittaiseksi, mahdollisesti jopa tyhjäksi.
		'''
		result = []
		date = tänään()
		# Niin kauan kuin aikatauluja ei ole vielä tarpeeksi,
		while len(result) < max_amount:
			try:
				# hae nykyisen päivän aikataulut ja lisää ne,
				result += self.schedule_for_day(date)
			except ValueError:
				# paitsi jos mentiin kalenterin ulkopuolelle, jolloin lopetetaan,
				break
			# ja siirry seuraavaan päivään.
			date += timedelta(1)
		# Typistä lopputulos haluttuun tulosmäärään.
		return result[:max_amount]
	def schedule_for_day(self, date):
		'''
			Hakee pysäkin aikataulut tiettynä päivänä.
		'''
		# Jos päädyttiin aikataulukalenterin ulkopuolelle, niin tuotetaan virhe. Jos vain
		# palautettaisiin tyhjä tulos, niin algoritmi jatkaisi etsintää loputtomiin.
		if date > viimeinen_käyttöpäivä:
			raise ValueError('tried to retrieve schedule for date %s which is outside schedule data' % date)
		result = []
		# Jokaiselle ajovuorolle,
		for trip in ajot.values():
			# jos tämä ajovuoro ajetaan tänä päivänä
			if trip.ajetaan_päivänä(date):
				# ja jos tämä ajo pysähtyy tällä pysäkillä, ei kuitenkaan saapuen
				# päätepysäkille,
				stop = trip.pysäkkiReitillä(self)
				if stop and stop is not trip.reitti[-1]:
					# ja jos tämä pysähdys on tulevaisuudessa,
					stop_time = datetime.combine(date, time()) + stop.saapumisaika
					if stop_time >= nyt():
						# lisää pysähdys listaan.
						result.append({
							'time': stop_time,
							'trip': trip,
							'stop': stop,
						})
		# Lajittele lopputulos saapumisajan mukaan.
		result.sort(key = lambda schedule_entry: schedule_entry['time'])
		return result

class Pysähdys:
	def __init__(self, saapumisaika, lähtöaika, pysäkki, ajo):
		self.saapumisaika, self.lähtöaika, self.pysäkki, self.ajo = saapumisaika, lähtöaika, \
			pysäkki, ajo
	def __repr__(self):
		return 'Pysähdys(%r, %r, %r, %r)' % (self.saapumisaika, self.lähtöaika, self.pysäkki, self.ajo)

linjat = {}
linjat_per_tunniste = {}
ajot = {}
ajot_per_numero = {}
palvelut = {}
pysäkit = {}
all_clusters = set()

print('Ladataan linjat... ', file = stderr, end = '', flush = True)
with open('gtfs/routes.txt') as tiedosto:
	for rivi in lue_csv(tiedosto):
		linja = Linja(rivi)
		linja.tunniste = linja.tunniste
		linjat[linja.viite] = linja
		linjat_per_tunniste[linja.tunniste] = linja
print('%d linjaa' % len(linjat), file = stderr)

print('Ladataan ajot... ', file = stderr, end = '', flush = True)
with open('gtfs/trips.txt') as tiedosto:
	for rivi in lue_csv(tiedosto, muunnokset = {'direction_id': lambda k: Suunta(int(k))}):
		if rivi['service_id'] not in palvelut:
			palvelut[rivi['service_id']] = Palvelu(rivi['service_id'])
		linja = linjat_per_tunniste[rivi['route_id']]
		ajo = Ajovuoro(tunniste = rivi['trip_id'],
			linja = linja,
			palvelu = palvelut[rivi['service_id']],
			kyltti = rivi['trip_headsign'],
			suunta = rivi['direction_id'])
		assert ajo.nimi not in ajot
		ajot[ajo.nimi] = ajo
print('%d ajoa' % len(ajot), file = stderr)

def lue_päiväys(teksti):
	return date(int(teksti[:4]), int(teksti[4:6]), int(teksti[6:]))

def lue_aika(teksti):
	tunti, minuutti, sekunti = map(int, teksti.split(':'))
	return timedelta(hours = tunti, minutes = minuutti, seconds = sekunti)

print('Ladataan päiväykset... ', file = stderr, flush = True)

viimeinen_käyttöpäivä = date.today()
palvelut_per_päivä = {}

with open('gtfs/calendar_dates.txt') as tiedosto:
	for rivi in lue_csv(tiedosto):
		palvelu = palvelut[rivi['service_id']]
		päivä = lue_päiväys(rivi['date'])
		palvelu.päivät.add(päivä)
		if päivä not in palvelut_per_päivä:
			palvelut_per_päivä[päivä] = set()
		palvelut_per_päivä[päivä].add(palvelu)
		viimeinen_käyttöpäivä = max(päivä, viimeinen_käyttöpäivä)

def palvelut_käytössä(päivä):
	for palvelu in palvelut.values():
		if päivä in palvelu.päivät:
			yield palvelu

print('Ladataan pysäkit... ', file = stderr, end = '', flush = True)
with open('gtfs/stops.txt') as file:
	for rivi in lue_csv(file):
		sijainti = Sijainti(float(rivi['stop_lat']), float(rivi['stop_lon']))
		pysäkki = Pysäkki(rivi['stop_id'], rivi['stop_name'], sijainti)
		pysäkit[pysäkki.tunniste] = pysäkki
with open('regions-per-stop.json') as file:
	for pysäkkitunniste, alue in json.load(file).items():
		pysäkit[pysäkkitunniste].alue = alue
print('%d pysäkkiä' % len(pysäkit), file = stderr)

class BusStopCluster:
	def __init__(self):
		self.stops = set()
		self._center = None
		self.name = None
	def add_stop(self, stop):
		assert not stop.cluster
		stop.cluster = self
		self.stops.add(stop)
		self._center = None
	@property
	def center(self):
		if not self._center:
			if self.stops:
				from statistics import median
				pointtype = type(next(iter(self.stops)).sijainti)
				self._center = pointtype(
					median(stop.sijainti.x for stop in self.stops),
					median(stop.sijainti.y for stop in self.stops),
				)
			else:
				raise ValueError('an empty cluster has no center point')
		return self._center
	def merge(self, other):
		for bus_stop in other.stops:
			bus_stop.cluster = self
		self.stops |= other.stops
		other.stops = set()
		other._center = None
	def schedule(self, max_amount = 50):
		result = []
		for stop in self.stops:
			result += stop.schedule(max_amount)
		result.sort(key = lambda schedule_entry: schedule_entry['time'])
		return result[:max_amount]

from collections import defaultdict
bus_stops_by_name = defaultdict(set)
for bus_stop in pysäkit.values():
	bus_stops_by_name[bus_stop.nimi].add(bus_stop)
bus_stops_by_name = dict(bus_stops_by_name)

# ryhmittele pysäkit nimen mukaan
all_clusters = []
def cluster_bus_stops():
	sorted_bus_stops = sorted(pysäkit.values(), key = lambda bus_stop: bus_stop.nimi)
	for bus_stop in sorted_bus_stops:
		if not bus_stop.cluster:
			stops_to_cluster = {bus_stop}
			# etsi pysäkin samannimiset vastaparit
			for pair_candidate in bus_stops_by_name[bus_stop.nimi]:
				distance = pair_candidate.sijainti.etäisyys(bus_stop.sijainti)
				if pair_candidate is not bus_stop and distance <= 0.3:
					stops_to_cluster.add(pair_candidate)
			for stop_to_cluster in stops_to_cluster:
				if stop_to_cluster.cluster:
					cluster = stop_to_cluster.cluster
					break
			else:
				cluster = BusStopCluster()
				all_clusters.append(cluster)
			for stop_to_cluster in stops_to_cluster:
				if not stop_to_cluster.cluster:
					cluster.add_stop(stop_to_cluster)
	# Merkitse muistiin pysäkkien vastaparit käyttäen hyväksi tämänhetkistä ryhmittelytietoa
	for bus_stop in pysäkit.values():
		if bus_stop.cluster:
			bus_stop.pairs = bus_stop.cluster.stops - {bus_stop}
	# Ryhmitä ne pysäkit, joilla ei ollut omaa vastaparia, muiden pysäkkien kanssa
	for bus_stop in sorted_bus_stops:
		if len(bus_stop.cluster.stops) == 1:
			possibilities = set()
			for cluster in all_clusters:
				if cluster is not bus_stop.cluster:
					distance = cluster.center.etäisyys(bus_stop.sijainti)
					if distance <= 0.3:
						possibilities.add((distance, cluster))
			if possibilities:
				best = min(possibilities)[1]
				all_clusters.remove(bus_stop.cluster)
				best.merge(bus_stop.cluster)

def shared_elements_in_n_sets(sets):
	from itertools import combinations
	result = set()
	for pair in combinations(sets, 2):
		result |= pair[0] & pair[1]
	return result

def name_clusters():
	from collections import defaultdict
	from pprint import pprint
	clusters_per_name = defaultdict(set)
	for cluster in all_clusters:
		name_representing_stop = min((len(pysäkki.tunniste), pysäkki.tunniste, pysäkki) for pysäkki in cluster.stops)[2]
		name = name_representing_stop.nimi.lower().replace(' ', '-')
		clusters_per_name[name].add(cluster)
	for name, clusters in clusters_per_name.items():
		if len(clusters) == 1:
			# Ryhmä on ainoa jolla on varaus tälle nimelle. Sen kuin vaan.
			next(iter(clusters)).name = name
		else:
			# Olisiko kaikki klusterit eri alueilla?
			common_regions = shared_elements_in_n_sets({stop.alue for stop in cluster.stops} for cluster in clusters)
			# Esitys: ryhmä -> ne alueet jotka ovat tälle ryhmälle ainutlaatuisia
			proposal = {
				cluster: {stop.alue for stop in cluster.stops} - common_regions - {None}
				for cluster in clusters
			}
			# Jos enintään yksi klusteri tässä esityksessä on kokonaan ilman omaa aluetta, jolla se voisi eritellä,
			# niin nimetään klusterit näiden alueiden mukaan.
			# Se klusteri jolla ei ole omaa aluetta (jos on) jätetään ilman aluepäätettä.
			if sum([1 for unique_areas in proposal.values() if not unique_areas]) <= 1:
				for cluster, unique_areas in proposal.items():
					individual_cluster_name = name
					if unique_areas:
						individual_cluster_name += '-' + min(unique_areas).lower().replace(' ', '-')
					cluster.name = individual_cluster_name
			else:
				# Typerä reunatapaus. Indeksoidaan numeroin...
				for n, (_, cluster) in enumerate(sorted(
					min((stop.tunniste.lower(), cluster) for stop in cluster.stops)
					for cluster in clusters
				), 1):
					individual_cluster_name = name + '-' + str(n)
					cluster.name = individual_cluster_name

print('Ryhmitellään pysäkit...')
cluster_bus_stops()
name_clusters()

clusters_by_name = {}
for cluster in all_clusters:
	assert cluster.name not in clusters_by_name
	clusters_by_name[cluster.name] = cluster

print('Ladataan aikataulut... ', end = '', flush = True, file = stderr)
with open('gtfs/stop_times.txt') as file:
	rivimäärä = sum(line.count('\n') for line in file)
	laskettu = 0
	file.seek(0)
	for rivi in lue_csv(file):
		ajo = ajot[muunna_ajovuoro_tunniste(rivi['trip_id'])]
		saapumisaika = lue_aika(rivi['arrival_time'])
		lähtöaika = lue_aika(rivi['departure_time'])
		pysäkki = pysäkit[rivi['stop_id']]
		ajo.reitti.append(Pysähdys(saapumisaika, lähtöaika, pysäkki, ajo))
		laskettu += 1
		if laskettu % 1000 == 0:
			print('\rLadataan aikataulut... %.1f%%' % (laskettu * 100 / rivimäärä), end = ' ', file = stderr)
print('\rLadataan aikataulut... ladattu', file = stderr)