Ordonner des horaires de train

class Train:
    """
    Representation of a train with multiple stops
    """

    def __init__(self, type, id, days=frozenset(), **stops):
        if len(stops) < 2:
            raise ValueError('At least 2 stops are needed')
        self.type = type
        self.id = id
        self.days = days
        self.stops = stops

    def __iter__(self):
        "Iterate over all (stop, stop_time) couples"
        return iter(self.stops.items())

    def iter_parts(self):
        "Iterable over all segments of the trip (couples of stop items)"
        it = iter(self)
        src_item = next(it)
        for dst_item in it:
            yield src_item, dst_item
            src_item = dst_item

    @property
    def departure(self):
        return next(iter(self.stops.keys()))

    @property
    def departure_time(self):
        return next(iter(self.stops.values()))

    @property
    def arrival(self):
        return next(reversed(self.stops.keys()))

    @property
    def arrival_time(self):
        return next(reversed(self.stops.values()))

import csv
import datetime
import sys

from .train import Train


time_fmt = '%H:%M'
day_names = ('L', 'Ma', 'Me', 'J', 'V', 'S', 'D')
days_mapping = {name: i for i, name in enumerate(day_names)}


def load(f=sys.stdin):
    "Load trains from a CSV file"
    reader = csv.reader(f)
    return map(load_train, reader)


def load_train(row):
    """
    Load train from a row

    load_train(['T', '#1', 'SD', 'viridian', '12:30', 'pewter', '13:00', 'cerulean', '14:00'])
    => Train('T', '#1', {5, 6}, viridian=time(12, 30), pewter=time(13, 0), cerulean=time(14, 0))
    """
    train_type, train_id, days, *stops = row
    days = {i for d, i in days_mapping.items() if d in days}
    stops = zip(stops[::2], stops[1::2])
    stops = {
        stop: datetime.datetime.strptime(time, time_fmt).time()
        for stop, time in stops
    }
    return Train(train_type, train_id, days, **stops)

import datetime


def diff_time(t1, t2):
    "Compute difference between two time objects"
    t1 = datetime.datetime.combine(datetime.date.min, t1)
    t2 = datetime.datetime.combine(datetime.date.min, t2)
    # t2 in the next day
    if t2 < t1:
        t2 += datetime.timedelta(days=1)
    return (t2 - t1).total_seconds()


def compute_distances(trains):
    """
    Compute minimal distances to neighboring stops (previous & next)
    for all stops
    Return two distances dicts (minimal distances to stop & from stop)

    compute_distances([
        Train(viridian=time(10), pewter=time(20)),
        Train(viridian=time(10), pewter=time(30)),
        Train(pewter=time(50), cerulean=time(80)),
    ])
    => (
        {
            'viridian': (0, 0, None),
            'pewter': (10, 1, 'viridian'),
            'cerulean': (30, 1, 'pewter'),
        },
        {
            'cerulean': (0, 0, None),
            'pewter': (30, 1, 'cerulean'),
            'viridian': (10, 1, 'pewter'),
        },
    )
    """
    distances_to = {}
    distances_from = {}
    default = (0, 0, None)

    for train in trains:
        for (src, time_src), (dst, time_dst) in train.iter_parts():
            dist = diff_time(time_src, time_dst)
            distances_to.setdefault(src, default)
            distances_from.setdefault(dst, default)

            old_dist, _, _ = distances_to.get(dst, default)
            if not old_dist or dist < old_dist:
                distances_to[dst] = dist, 1, src

            old_dist, _, _ = distances_from.get(src, default)
            if not old_dist or dist < old_dist:
                distances_from[src] = dist, 1, dst

    return distances_to, distances_from


def normalize_distances(distances):
    """
    Normalize a distances dict
    Re-compute the distances to make them absolute from departure.arrival

    normalize_distances({
        'viridian': (0, 0, None),
        'pewter': (10, 1, 'viridian'),
        'cerulean': (30, 1, 'pewter'),
    })
    => {
        'viridian': (0, 0, None),
        'pewter': (10, 1, None),
        'cerulean': (40, 2, None),
    }
    """
    def set_absolute_dist(stop):
        # Walk recursively through stops to update min_dist & n_stops
        dist, n, src = distances[stop]

        if src is not None:
            d, i, src = set_absolute_dist(src)
            dist += d
            n += i
            distances[stop] = dist, n, src

        return dist, n, src

    for stop in list(distances):
        set_absolute_dist(stop)

    return distances


def get_sorted_stops(trains):
    """
    Get a sorted list of all stops by computing distances

    get_sorted_stops([
        Train(viridian=time(10), pewter=time(20)),
        Train(viridian=time(10), pewter=time(30)),
        Train(pewter=time(50), cerulean=time(80)),
    ])
    => ['viridian', 'pewter', 'cerulean']
    """
    distances_to, distances_from = compute_distances(trains)
    normalize_distances(distances_to)
    normalize_distances(distances_from)

    def sort_key(key):
        # Stops are sorted by closest from departure & furthest from arrival
        dist1, n1, _ = distances_to[key]
        dist2, n2, _, = distances_from[key]
        return dist1 - dist2, n1 - n2

    return sorted(distances_to, key=sort_key)

from contextlib import contextmanager


class _Group:
    "Identify a mergeable group"
    def __init__(self, key):
        self.key = key
        self.set = {self}

    def merge(self, rhs, merge_key_func=None):
        # Merge = union of both groups sets
        self.set.update(rhs.set)
        if merge_key_func:
            self.key = merge_key_func(self.key, rhs.key)
        for item in rhs.set:
            item.key = self.key
            item.set = self.set

    def clear(self):
        self.set.clear()


class HashInstance:
    "Wrap any Python object to make it hashable"
    def __init__(self, obj):
        self.obj = obj

    def __eq__(self, rhs):
        if not isinstance(rhs, __class__):
            return NotImplemented
        return self.obj is rhs.obj

    def __hash__(self):
        return object.__hash__(self.obj)


class Grouper:
    "Abstraction to manipulate groups over common objects"
    def __init__(self):
        self._groups = set()
        self._groups_by_object = {}

    def add(self, obj, key=None):
        "Create a new group for an object"
        group = _Group(key)
        # Register the object in the group
        group.obj = obj
        self._groups.add(group)
        obj = HashInstance(obj)
        self._groups_by_object[obj] = group

    def _obj_group(self, obj):
        "Get the group associated to an object"
        objh = HashInstance(obj)
        return self._groups_by_object[objh]

    def get_key(self, obj):
        return self._obj_group(obj).key

    def equal(self, lhs, rhs):
        "Compare groups of two objects"
        lgroup = self._obj_group(lhs)
        rgroup = self._obj_group(rhs)
        return lgroup.set is rgroup.set

    def merge(self, lhs, rhs, merge_key=None):
        "Merge groups of two objects"
        lgroup = self._obj_group(lhs)
        rgroup = self._obj_group(rhs)
        lgroup.merge(rgroup, merge_key_func=merge_key)

    def groups(self):
        "Iterate over objects by group"
        for group in self._groups:
            yield group.key, [g.obj for g in group.set]

    def clear(self):
        "Clear all groups"
        for group in self._groups:
            del group.obj
            group.clear()
        self._groups.clear()
        self._groups_by_object.clear()


@contextmanager
def grouper():
    "Create & clear a Grouper object"
    g = Grouper()
    try:
        yield g
    finally:
        g.clear()

from .utils import grouper


def get_next_train(train, trains):
    "Get direct/earliest connection after a train"
    trains = [t for t in trains if t.departure_time > train.arrival_time]
    if trains:
        return min(trains, key=lambda t: t.departure_time)


def get_prev_train(train, trains):
    "Get direct/latest connection before a train"
    trains = [t for t in trains if t.arrival_time < train.departure_time]
    if trains:
        return max(trains, key=lambda t: t.arrival_time)


def group_trains(trains, sorted_stops):
    """
    Group trains by connections
    All trains that are connected belong to a same group
    Return sorted groups as lists of trains

    group_trains([
        Train('T', '#1', viridian=time(12, 30), pewter=time(13, 0)),
        Train('T', '#2', viridian=time(16, 0), pewter=time(16, 30)),
        Train('T', '#3', pewter=time(13, 30), cerulean=time(14, 30)),
        Train('T', '#4', pewter=time(17, 0), cerulean=time(18, 0)),
    ], ['viridian', 'pewter', 'cerulean'])
    => [
        [Train('T', '#1', ...), Train('T', '#3', ...)],
        [Train('T', '#2', ...), Train('T', '#4', ...)],
    ]
    """
    trains_from_stop = {stop: [] for stop in sorted_stops}
    trains_to_stop = {stop: [] for stop in sorted_stops}

    with grouper() as g:
        for train in trains:
            # Groups are associated to a departure time
            # when merged, groups get associated to earlier departure times
            g.add(train, train.departure_time)
            trains_from_stop[train.departure].append(train)
            trains_to_stop[train.arrival].append(train)

        for stop in sorted_stops:
            # Connect each train with next one
            for train in trains_from_stop[stop]:
                next_train = get_next_train(train, trains_from_stop[train.arrival])
                if next_train:
                    g.merge(train, next_train, merge_key=min)

            # Connect each train with previous one
            for train in trains_to_stop[stop]:
                prev_train = get_prev_train(train, trains_to_stop[train.departure])
                if prev_train:
                    g.merge(train, prev_train, merge_key=min)

        # Sort trains in groups by time of arrival & departure
        grouped_trains = {
            time: sorted(trains, key=lambda t: (t.departure_time, t.arrival_time))
            for time, trains in g.groups()
        }
        # Sort groups by departure time of the first train
        return [trains for _, trains in sorted(grouped_trains.items())]

def get_day_groups(trains):
    """
    Get a list of all distinctive sets of days

    get_day_groups([
        Train(days={1, 2, 3}),
        Train(days={0, 1}),
        Train(days={4}),
    ])
    => [{0}, {1}, {2, 3}, {4}]
    """
    day_sets = {frozenset(train.days) for train in trains}
    all_days = frozenset.union(*day_sets)
    groups = {all_days}

    for day_set in day_sets:
        couples = ((group & day_set, group - day_set) for group in groups)
        groups = {group for couple in couples for group in couple if group}

    return sorted(groups, key=tuple)

day_names = ('L', 'Ma', 'Me', 'J', 'V', 'S', 'D')
day_full_names = (
    'Lundi', 'Mardi', 'Mercredi', 'Jeudi', 'Vendredi', 'Samedi', 'Dimanche',
)


def dump_days(days, full=False):
    "Get printable format for days set"
    if full:
        return ', '.join(day_full_names[i] for i in sorted(days))
    else:
        return ''.join(day_names[i] for i in sorted(days))

from .parse import dump_days


def iter_table(sorted_stops, grouped_trains, filter_days=set(), *, print_days=True):
    header = ['Train']
    if print_days:
        header.append('Jours')
    header.extend(sorted_stops)
    yield header

    for group in grouped_trains:
        yield
        for train in group:
            if not (train.days >= filter_days):
                continue
            row = [f'{train.type} {train.id}']
            if print_days:
                row.append(dump_days(train.days))
            stops = (train.stops.get(stop) for stop in sorted_stops)
            row.extend(time.strftime(time_fmt) if time else '' for time in stops)
            yield row