szczanieckiej/traffic/convert_stops.go

// SPDX-FileCopyrightText: Adam Evyčędo
//
// SPDX-License-Identifier: AGPL-3.0-or-later

package traffic

import (
	"database/sql"
	"encoding/csv"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"sort"
	"strings"

	"git.sr.ht/~sircmpwn/go-bare"
)

func dropInputStopsIndex(c feedConverter) feedConverter {
	c.stopsInputIndex = map[string]int64{}
	return c
}

func readTripsThroughStopsIndex(c feedConverter) (feedConverter, error) {
	index := map[string]int64{}

	path := c.TmpFeedPath
	forEachRow(filepath.Join(path, "tripsthroughstop.csv"), func(offset int64, fields map[string]int, record []string) error {
		stopID := record[fields["stop_id"]]
		if _, ok := index[stopID]; !ok {
			index[stopID] = offset
		}
		return nil
	})

	c.stopsInputIndex = index
	return c, nil
}

func dropInputTripsIndex(c feedConverter) feedConverter {
	c.tripsInputIndex = map[string]int64{}
	return c
}

func dropInputRoutesIndex(c feedConverter) feedConverter {
	c.routesInputIndex = map[string]int64{}
	return c
}

func convertStops(c feedConverter) (feedConverter, error) { // O(n:stops) ; (translations, file:tripsThroughStop, file:tripChangeOpts, tripOffsets -- stopsOffsetsByCode:CodeIndex, stopsOffsetsByName:map[name][]offsets >> stops)
	path := c.TmpFeedPath
	var outputOffset uint = 0
	stopsOffsetsByName := map[string][]uint{}
	stopsOffsetsByCode := CodeIndex{}
	stops := map[string]string{}
	maxStopTripsLength := 0

	result, err := os.Create(filepath.Join(path, "stops.bare"))
	if err != nil {
		return c, fmt.Errorf("while creating file: %w", err)
	}
	defer result.Close()

	tripsThroughStopFile, err := os.Open(filepath.Join(path, "tripsthroughstop.csv"))
	if err != nil {
		return c, fmt.Errorf("while opening tripsThroughStop file: %w", err)
	}
	defer tripsThroughStopFile.Close()

	tripsThroughStop := csv.NewReader(tripsThroughStopFile)
	tripsThroughStopHeader, err := tripsThroughStop.Read()
	if err != nil {
		return c, fmt.Errorf("while reading tripsThroughStop header: %w", err)
	}

	tripsThroughStopFields := map[string]int{}
	for i, headerField := range tripsThroughStopHeader {
		tripsThroughStopFields[headerField] = i
	}

	cacheDir, err := os.UserCacheDir()
	if err != nil {
		return c, fmt.Errorf("while getting cache dir: %w", err)
	}
	db, err := sql.Open("sqlite3", filepath.Join(cacheDir, "turntable.db"))
	if err != nil {
		return c, fmt.Errorf("while opening db: %w", err)
	}

	err = forEachRow(filepath.Join(path, "stops.txt"), func(offset int64, fields map[string]int, record []string) error {
		if f, ok := fields["location_type"]; ok && record[f] != "" && record[f] != "0" {
			// NOTE for now ignore everything that’s not a stop/platform
			// TODO use Portals (location_type == 2) to show on map if platform has a parent (location_type == 1) that has a Portal
			// TODO use location_type in {3,4} for routing inside stations (with pathways, transfers, and levels)
			return nil
		}

		stop := Stop{}

		stopID := record[fields["stop_id"]]

		stopTrips := map[string]StopOrder{}
		if position, ok := c.stopsInputIndex[stopID]; ok {
			tripsThroughStopFile.Seek(position, 0)
			for {
				tripsThroughStopRecord, err := readCsvLine(tripsThroughStopFile, -1, len(tripsThroughStopHeader))
				if err != nil {
					if err == io.EOF {
						break
					} else {
						return fmt.Errorf("while reading tripsThroughStop record: %w", err)
					}
				}
				recordStopID := tripsThroughStopRecord[tripsThroughStopFields["stop_id"]]
				if stopID != recordStopID {
					break
				}
				tripID := tripsThroughStopRecord[tripsThroughStopFields["trip_id"]]
				var sequence int
				fmt.Sscanf(tripsThroughStopRecord[tripsThroughStopFields["sequence"]], "%d", &sequence)
				stopTrips[tripID] = StopOrder{
					Sequence:   sequence,
					TripOffset: c.tripsOffsets[tripID],
				}
			}
			stopTripsLength := len(stopTrips)
			if maxStopTripsLength < stopTripsLength {
				maxStopTripsLength = stopTripsLength
			}
		}

		stop.Id = stopID

		templates := []string{"stop_code", "stop_id", "stop_name", "platform_code"}
		stop.Code = c.Feed.Flags().StopIdFormat
		for _, template := range templates {
			stop.Code = strings.Replace(stop.Code, "{{"+template+"}}", record[fields[template]], -1)
		}
		stop.Name = c.Feed.Flags().StopName
		for _, template := range templates {
			// TODO if '{{template}}' is empty
			stop.Name = strings.Replace(stop.Name, "{{"+template+"}}", record[fields[template]], -1)
		}
		if field, ok := fields["zone_id"]; ok {
			stop.Zone = record[field]
		}
		stop.NodeName = record[fields["stop_name"]]

		stops[record[fields["stop_id"]]] = stop.Code

		if field, ok := fields["stop_timezone"]; ok {
			stop.Timezone = record[field]
		}

		if c.feedInfo.Language == "mul" {
			key := record[fields["stop_name"]]
			if _, ok := c.translations[stop.NodeName][c.defaultLanguage]; !ok {
				stop.TranslatedNames = []Translation{{Language: c.defaultLanguage, Value: stop.Name}}
				stop.TranslatedNodeNames = []Translation{{Language: c.defaultLanguage, Value: stop.NodeName}}
			} else {
				stop.TranslatedNames = []Translation{{Language: c.defaultLanguage, Value: strings.ReplaceAll(stop.Name, key, c.translations[key][c.defaultLanguage])}}
				stop.TranslatedNodeNames = []Translation{{Language: c.defaultLanguage, Value: c.translations[key][c.defaultLanguage]}}
			}
			for language, value := range c.translations[key] {
				if language == c.defaultLanguage {
					continue
				}
				stop.TranslatedNames = append(stop.TranslatedNames, Translation{Language: c.defaultLanguage, Value: strings.ReplaceAll(stop.Name, key, value)})
				stop.TranslatedNodeNames = append(stop.TranslatedNodeNames, Translation{Language: c.defaultLanguage, Value: c.translations[key][value]})
			}
		}

		var lat, lon float64
		fmt.Sscanf(record[fields["stop_lat"]], "%f", &lat)
		fmt.Sscanf(record[fields["stop_lon"]], "%f", &lon)
		stop.Position = Position{lat, lon}

		stop.ChangeOptions = []ChangeOption{}
		stop.Order = stopTrips

		rows, err := db.Query("select line_name, headsign from change_options where stop_id = ?", stopID)
		if err != nil {
			return fmt.Errorf("while querrying change options: %w", err)
		}
		for rows.Next() {
			var (
				lineName string
				headsign string
			)
			rows.Scan(&lineName, &headsign)
			stop.ChangeOptions = append(stop.ChangeOptions, ChangeOption{
				LineName:            lineName,
				Headsign:            headsign,
				TranslatedHeadsigns: []Translation{},
				// TODO add translations
			})
		}

		sort.Slice(stop.ChangeOptions, func(i, j int) bool {
			var num1, num2 int
			_, err1 := fmt.Sscanf(stop.ChangeOptions[i].LineName, "%d", &num1)
			_, err2 := fmt.Sscanf(stop.ChangeOptions[j].LineName, "%d", &num2)
			if err1 != nil && err2 != nil {
				return stop.ChangeOptions[i].LineName < stop.ChangeOptions[j].LineName
			} else if err1 != nil {
				return false
			} else if err2 != nil {
				return true
			} else {
				return num1 < num2
			}
		})

		bytes, err := bare.Marshal(&stop)
		if err != nil {
			return fmt.Errorf("while marshalling: %w", err)
		}
		b, err := result.Write(bytes)
		if err != nil {
			return fmt.Errorf("while writing: %w", err)
		}
		if len(stop.TranslatedNames) == 0 {
			stopsOffsetsByName[stop.Name] = append(stopsOffsetsByName[stop.Name], outputOffset)
		}
		for _, v := range stop.TranslatedNames {
			stopsOffsetsByName[v.Value] = append(stopsOffsetsByName[v.Value], outputOffset)
		}
		stopsOffsetsByCode[stop.Code] = outputOffset
		outputOffset += uint(b)

		return nil
	})

	if maxStopTripsLength > 12288 {
		log.Printf("maximum length of StopOrder is %d, more than 12288, which may need to be tweaked", maxStopTripsLength)
	}

	c.StopsCodeIndex = stopsOffsetsByCode
	c.StopsNameIndex = stopsOffsetsByName
	c.Stops = stops
	os.Remove(filepath.Join(cacheDir, "turntable.db"))
	return c, err
}