From dd6e3bba5dbbad5cb3bd7e45d70fd4af4c072522 Mon Sep 17 00:00:00 2001
From: steffenaxer <26229392+steffenaxer@users.noreply.github.com>
Date: Tue, 23 Jan 2024 07:50:25 +0100
Subject: [PATCH] Clear PassengerGroupIdentifier

---
 .../extension/companions/DrtCompanionRideGenerator.java    | 7 ++++++-
 .../drt/extension/companions/DrtCompanionUtils.java        | 5 +++++
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/contribs/drt-extensions/src/main/java/org/matsim/contrib/drt/extension/companions/DrtCompanionRideGenerator.java b/contribs/drt-extensions/src/main/java/org/matsim/contrib/drt/extension/companions/DrtCompanionRideGenerator.java
index 8f1839b67d6..d6019872aa3 100644
--- a/contribs/drt-extensions/src/main/java/org/matsim/contrib/drt/extension/companions/DrtCompanionRideGenerator.java
+++ b/contribs/drt-extensions/src/main/java/org/matsim/contrib/drt/extension/companions/DrtCompanionRideGenerator.java
@@ -57,8 +57,8 @@ final class DrtCompanionRideGenerator implements BeforeMobsimListener, AfterMobs
 	private final Scenario scenario;
 	private final String drtMode;
 	private final int maxCapacity;
-
 	private final Set<Id<Person>> companionAgentIds = new HashSet<>();
+	private final Set<Leg> drtLegs = new HashSet<>();
 	private WeightedRandomSelection<Integer> sampler;
 
 	private final Map<Id<PassengerGroupIdentifier.PassengerGroup>, List<GroupLeg>> passengerGroups = new HashMap<>();
@@ -113,6 +113,7 @@ private void addCompanionAgents() {
 
 				for (Leg leg : trip.getLegsOnly()) {
 					if (leg.getMode().equals(this.drtMode)) {
+						this.drtLegs.add(leg);
 						// Initial person travels now in a group
 						Id<PassengerGroupIdentifier.PassengerGroup> currentGroupIdentifier = getGroupIdentifier();
 						DrtCompanionUtils.setPassengerGroupIdentifier(leg, currentGroupIdentifier);
@@ -221,6 +222,10 @@ private void removeCompanionAgents() {
 		passengerGroups.clear();
 		personIdentifierSuffix.set(0);
 		LOG.info("Removed # {} drt companion agents", counter);
+
+		// Remove attribute from legs of initial drt riders
+		this.drtLegs.stream().forEach(DrtCompanionUtils::removePassengerGroupIdentifier);
+		this.drtLegs.clear();
 	}
 
 	@Override
diff --git a/contribs/drt-extensions/src/main/java/org/matsim/contrib/drt/extension/companions/DrtCompanionUtils.java b/contribs/drt-extensions/src/main/java/org/matsim/contrib/drt/extension/companions/DrtCompanionUtils.java
index cc2f7caed31..8646f201d3a 100644
--- a/contribs/drt-extensions/src/main/java/org/matsim/contrib/drt/extension/companions/DrtCompanionUtils.java
+++ b/contribs/drt-extensions/src/main/java/org/matsim/contrib/drt/extension/companions/DrtCompanionUtils.java
@@ -109,4 +109,9 @@ public static void setPassengerGroupIdentifier(Leg leg, Id<PassengerGroupIdentif
 		leg.getAttributes().putAttribute(GROUP_IDENTIFIER_ATTRIBUTE, passengerGroupIdentifierId);
 	}
 
+	public static void removePassengerGroupIdentifier(Leg leg)
+	{
+		leg.getAttributes().removeAttribute(GROUP_IDENTIFIER_ATTRIBUTE);
+	}
+
 }