2 files changed, 96 insertions, 6 deletions

diff --git a/lib/Travel/Status/DE/DBWagenreihung.pm b/lib/Travel/Status/DE/DBWagenreihung.pm
index 962211a..0bc7ae6 100644
--- a/lib/Travel/Status/DE/DBWagenreihung.pm
+++ b/lib/Travel/Status/DE/DBWagenreihung.pm
@@ -158,6 +158,10 @@ sub train_no {
 sub train_subtype {
 	my ($self) = @_;
 
+	if ( exists $self->{train_subtype} ) {
+		return $self->{train_subtype};
+	}
+
 	my @wagons = $self->wagons;
 
 	my %ml = (
@@ -210,7 +214,8 @@ sub train_subtype {
 		return undef;
 	}
 
-	return $likelihood[0];
+	$self->{train_subtype} = $likelihood[0];
+	return $self->{train_subtype};
 }
 
 sub wagons {
@@ -241,6 +246,17 @@ sub wagons {
 	@{ $self->{wagons} } = sort {
 		$a->{position}->{start_percent} <=> $b->{position}->{start_percent}
 	} @{ $self->{wagons} };
+
+	# ->train_subtype calls ->wagons, so this call must not be made before
+	# $self->{wagons} has beet set.
+	my $tt = $self->train_subtype;
+
+	if ($tt) {
+		for my $wagon ( @{ $self->{wagons} } ) {
+			$wagon->set_traintype($tt);
+		}
+	}
+
 	return @{ $self->{wagons} // [] };
 }
 
diff --git a/lib/Travel/Status/DE/DBWagenreihung/Wagon.pm b/lib/Travel/Status/DE/DBWagenreihung/Wagon.pm
index 5c2957e..308a6b0 100644
--- a/lib/Travel/Status/DE/DBWagenreihung/Wagon.pm
+++ b/lib/Travel/Status/DE/DBWagenreihung/Wagon.pm
@@ -9,11 +9,55 @@ use parent 'Class::Accessor';
 use Carp qw(cluck);
 
 our $VERSION = '0.00';
-
 Travel::Status::DE::DBWagenreihung::Wagon->mk_ro_accessors(
-	qw(attributes class_type has_ac has_accessibility has_bistro has_compartments
-	  has_multipurpose is_dosto is_interregio is_locomotive is_powercar number
-	  model section type)
+	qw(attributes class_type has_ac has_accessibility has_bahn_comfort
+	  has_bike_storage has_bistro has_compartments has_family_area
+	  has_phone_area has_quiet_area is_dosto is_interregio is_locomotive
+	  is_powercar number model multipurpose section type)
+);
+
+our %type_attributes = (
+	'ICE 1' => [
+		undef, ['has_quiet_area'], undef, ['has_quiet_area'],    # 1 2 3 4
+		['has_family_area'], undef, ['has_bahn_comfort'],        # 5 6 7
+		undef, undef, undef, ['has_bahn_comfort'],               # 8 9 (10) 11
+		['has_quiet_area'], undef, undef                         # 12 (13) 14
+	],
+	'ICE 2' => [
+		undef, ['has_quiet_area'], ['has_bahn_comfort'],
+		['has_family_area'],                                     # 1 2 3 4
+		undef, ['has_bahn_comfort'],
+		[ 'has_quiet_area', 'has_phone_area' ]                   # 5 6 7
+	],
+	'ICE 3' => [
+		['has_quiet_area'], undef, undef, undef,                 # 1 2 3 (4)
+		['has_family_area'], undef, ['has_bahn_comfort'],        # 5 6 7
+		[ 'has_quiet_area', 'has_phone_area', 'has_bahn_comfort' ], undef  # 8 9
+	],
+	'ICE 3 V' => [
+		['has_quiet_area'], undef, undef, ['has_family_area'],    # 1 2 3 4
+		['has_bahn_comfort'], ['has_bahn_comfort'], undef, undef, # 5 6 (7) 8
+		[ 'has_quiet_area', 'has_phone_area' ]                    # 9
+	],
+	'ICE 4' => [
+		['has_bike_storage'], undef, ['has_quiet_area'], undef,
+		undef,                                                    # 1 2 3 4 5
+		undef, ['has_bahn_comfort'], undef, ['has_family_area'],  # 6 7 (8) 9
+		undef, ['has_bahn_comfort'], undef, undef,
+		['has_quiet_area']    # 10 11 12 (13) 14
+	],
+	'ICE T 411' => [
+		['has_quiet_area'], ['has_quiet_area'], undef,
+		['has_family_area'],    # 1 2 3 4
+		undef, undef, ['has_bahn_comfort'],
+		[ 'has_quiet_area', 'has_bahn_comfort' ]    # (5) 6 7 8
+	],
+	'ICE T 415' => [
+		['has_quiet_area'], ['has_quiet_area'], ['has_bahn_comfort'],
+		undef,                                      # 1 2 3 (4)
+		undef, undef, ['has_family_area'],
+		[ 'has_quiet_area', 'has_bahn_comfort' ]    # (5) (6) 7 8
+	],
 );
 
 sub new {
@@ -91,7 +135,7 @@ sub parse_type {
 	}
 
 	if ( $type =~ m{d} ) {
-		$self->{has_multipurpose} = 1;
+		$self->{multipurpose} = 1;
 		push( @desc, 'Mehrzweck' );
 	}
 
@@ -132,6 +176,36 @@ sub parse_type {
 	$self->{attributes} = \@desc;
 }
 
+sub set_traintype {
+	my ( $self, $tt ) = @_;
+
+	if ( not $self->{number} or not exists( $type_attributes{$tt} ) ) {
+		return;
+	}
+
+	if ( $self->{number} !~ m{^\d+$} ) {
+		return;
+	}
+
+	my $index = $self->{number} - 1;
+
+	if ( $index >= 30 ) {
+		$index -= 30;
+	}
+	elsif ( $index >= 20 ) {
+		$index -= 20;
+	}
+
+	if ( not $type_attributes{$tt}[$index] ) {
+		return;
+	}
+
+	for my $attr ( @{ $type_attributes{$tt}[$index] } ) {
+		$self->{$attr} = 1;
+		say "$index -> $attr";
+	}
+}
+
 sub is_first_class {
 	my ($self) = @_;