fix: names corrections for multiple units

src/utils/trainUtils.ts

@@ -1,3 +1,11 @@
+const unitNameCorrections: Record<string, string[]> = {
+  '2EN57': ['EN57', 'EN57'],
+  '201E': ['ET22'],
+  '4E': ['EU07'],
+  M62: ['ST44'],
+  CTLR4C: ['ST44']
+};
+
 export const getRegionNameById = (id: string) => {
   switch (id) {
     case 'eu':
@@ -20,10 +28,18 @@ export const getRegionNameById = (id: string) => {
   }
 };
 
-export const unitNameCorrections: Record<string, string> = {
-  '2EN57': 'EN57',
-  '201E': 'ET22',
-  '4E': 'EU07',
-  M62: 'ST44',
-  CTLR4C: 'ST44',
-};
+export function getHeadUnits(stockString: string) {
+  const stockList = stockString.split(';').slice(0, 3);
+
+  return stockList.reduce((acc, unitType, i) => {
+    if (i != 0 && !/-\d+$/.test(unitType)) return acc;
+
+    const unitName = unitType.slice(0, unitType.indexOf('-'));
+
+    const correctedNames = unitNameCorrections[unitName] ?? [unitName];
+
+    acc.push(...correctedNames);
+
+    return acc;
+  }, [] as string[]);
+}