parse-trains

#!/usr/bin/env python

import json
import re
import sys

from split.parse import unpack, loop, fare_file


try:
    timetable_dir = sys.argv[2]
    id = re.match('timetable-data-(\d+)', timetable_dir).group(1)
    timetable_file = '%s/ttisf%s.mca' % (timetable_dir, id)
except IndexError:
    print("Specify the timetable ID as the second parameter", file=sys.stderr)
    sys.exit(1)


lookups = {
    "SR": ( ('code', 2), ('id', 6), ('dir', 1), ('quota', 1), ('sleeping_ind', 1) ),

    'BS': ( ('type', 1), ('uid', 6), ('date_from', 6), ('date_to', 6), ('days', 7), ('bh', 1), ('status', 1),
        ('category', 2), ('identity', 4), ('headcode', 4), ('course_indicator', 1), ('service_code', 8),
        ('portion_id', 1), ('power_type', 3), ('timing_load', 4), ('speed', 3), ('opchars', 6),
        ('train_class', 1), ('sleepers', 1), ('reservations', 1), ('connection_indicator', 1),
        ('catering', 4), ('branding', 4), ('spare', 1), ('stp', 1),
    ),
    'BX': ( ('dummy', 4), ('uic', 5), ('atoc', 2), ('ats', 1) ),
    'LO': ( ('location', 8), ('dep_sched', 5), ('dep_public', 4), ('platform', 3), ('line', 3), ('eng_allow', 2), ('pathing_allow', 2), ('activity', 12), ('perf_allow', 2) ),
    'LI': ( ('location', 8), ('arr_sched', 5), ('dep_sched', 5), ('pass_sched', 5), ('arr_public', 4), ('dep_public', 4), ('platform', 3), ('line', 3), ('path', 3), ('activity', 12), ('eng_allow', 2), ('pathing_allow', 2), ('perf_allow', 2)  ),
    'LT': ( ('location', 8), ('arr_sched', 5), ('arr_public', 4), ('platform', 4), ('path', 3), ('activity', 12) ),
}

TRAINS = set()
for row in loop(fare_file('RST')):
    record_type, cf_mkr = row[0:2], row[2]

    if cf_mkr == 'F': continue
    if record_type != 'SR': continue
    row = unpack(row[3:], *lookups[record_type])
    TRAINS.add(row['id'])

stations = {}
started = None
data = {}
for line in open(timetable_file):
    if line.startswith('TI'):
        tiploc = line[2:9].strip()
        crs = line[53:56].strip()
        if not crs: continue
        stations[tiploc] = crs
        continue

    if line.startswith('BSN'):
        if line[3:9] in TRAINS:
            started = line[3:9]
        else:
            started = None

    if started:
        type = line[0:2]
        if type == 'CR': continue
        row = unpack(line[2:], *lookups[type])
        if type == 'BS':
            if row['days'][1] == '0': continue
            del row['reservations'], row['power_type'], row['train_class'], row['speed'], row['headcode'], row['category'], row['course_indicator'], row['connection_indicator'], row['branding'], row['type'], row['status'], row['bh'], row['catering'], row['spare'], row['portion_id'], row['opchars'], row['sleepers'], row['timing_load'], row['identity'], row['stp'], row['service_code']
            assert row['uid'] == started
            data[row['uid']] = row
        elif type == 'BX':
            continue
        elif type in ('LO', 'LI', 'LT'):
            if row['location'] not in stations: continue
            if row.get('arr_sched') == '': continue
            if started not in data: continue
            data[started].setdefault('stops', {})[stations[row['location']]] = [ row.get('arr_public'), row.get('dep_public') ]
        else:
            raise Exception

json.dump(data, open('data/trains.json', 'w'), separators=(',', ': '), indent=2, sort_keys=True)