org.opentripplanner.ext.flex.FlexIndex Maven / Gradle / Ivy

Go to download
package org.opentripplanner.ext.flex;

import com.google.common.collect.ArrayListMultimap;
import com.google.common.collect.HashMultimap;
import com.google.common.collect.Multimap;
import org.opentripplanner.common.geometry.HashGridSpatialIndex;
import org.opentripplanner.ext.flex.trip.FlexTrip;
import org.opentripplanner.model.FeedScopedId;
import org.opentripplanner.model.FlexStopLocation;
import org.opentripplanner.model.FlexLocationGroup;
import org.opentripplanner.model.Route;
import org.opentripplanner.model.SimpleTransfer;
import org.opentripplanner.model.StopLocation;
import org.opentripplanner.model.Trip;
import org.opentripplanner.routing.graph.Graph;

import java.util.HashMap;
import java.util.Map;
import java.util.stream.Stream;

public class FlexIndex {
  public Multimap transfersToStop = ArrayListMultimap.create();

  public Multimap flexTripsByStop = HashMultimap.create();

  public Multimap locationGroupsByStop = ArrayListMultimap.create();

  public HashGridSpatialIndex locationIndex = new HashGridSpatialIndex();

  public Map routeById = new HashMap<>();

  public Map tripById = new HashMap<>();

  public FlexIndex(Graph graph) {
    for (SimpleTransfer transfer : graph.transfersByStop.values()) {
      transfersToStop.put(transfer.to, transfer);
    }
    for (FlexTrip flexTrip : graph.flexTripsById.values()) {
      routeById.put(flexTrip.getTrip().getRoute().getId(), flexTrip.getTrip().getRoute());
      tripById.put(flexTrip.getTrip().getId(), flexTrip.getTrip());
      for (StopLocation stop : flexTrip.getStops()) {
        if (stop instanceof FlexLocationGroup) {
          for (StopLocation stopElement : ((FlexLocationGroup) stop).getLocations()) {
            flexTripsByStop.put(stopElement, flexTrip);
          }
        } else {
          flexTripsByStop.put(stop, flexTrip);
        }
      }
    }
    for (FlexLocationGroup flexLocationGroup : graph.locationGroupsById.values()) {
      for (StopLocation stop : flexLocationGroup.getLocations()) {
        locationGroupsByStop.put(stop, flexLocationGroup);
      }
    }
    for (FlexStopLocation flexStopLocation : graph.locationsById.values()) {
      locationIndex.insert(flexStopLocation.getGeometry().getEnvelopeInternal(), flexStopLocation);
    }
  }

  Stream getFlexTripsByStop(StopLocation stopLocation) {
    return flexTripsByStop.get(stopLocation).stream();
  }
}