Добавлен дополнительный класс, изменена логика прорисовки колес

DrawningLocomotive.java

@@ -3,7 +3,6 @@ import java.util.Random;
 
 public class DrawningLocomotive {
     public EntityLocomotive Locomotive;
-    public ExtraWheelsDraw extraWheelsDraw;
     public IDrawningExtra drawningExtra;
     /// Левая координата отрисовки локомотива
     protected float _startPosX;
@@ -22,13 +21,12 @@ public class DrawningLocomotive {
     public DrawningLocomotive(int speed, float weight, Color bodyColor)
     {
         int randExtra = random.nextInt(2);
-        extraWheelsDraw = new ExtraWheelsDraw(randExtra, bodyColor);
         switch (random.nextInt(3)){
             case 0:
-                drawningExtra = new ExtraStarWheelDraw(randExtra);
+                drawningExtra = new ExtraStarWheelDraw(randExtra, bodyColor);
                 break;
             case 1:
-                drawningExtra = new ExtraRoundWheelDraw(randExtra);
+                drawningExtra = new ExtraRoundWheelDraw(randExtra, bodyColor);
                 break;
             case 2:
                 drawningExtra = new ExtraWheelsDraw(randExtra, bodyColor);
@@ -37,6 +35,11 @@ public class DrawningLocomotive {
         Locomotive = new EntityLocomotive(speed, weight, bodyColor);
     }
 
+    public DrawningLocomotive(EntityLocomotive locomotive, IDrawningExtra extra) {
+        drawningExtra = extra;
+        Locomotive = locomotive;
+    }
+
     // Новый конструктор
     protected DrawningLocomotive (int speed, float weight, Color bodyColor, int locomotiveWidth, int locomotiveHeight)
     {
@@ -122,8 +125,6 @@ public class DrawningLocomotive {
         //дверь
         g.setColor(Color.BLACK);
         g.fillRect( (int)_startPosX + 50, (int)_startPosY + 10, 10, 20);
-        //колеса
-        extraWheelsDraw.DrawExtra((int)_startPosX, (int)_startPosY, g);
         //extra
         drawningExtra.DrawExtra((int)_startPosX, (int)_startPosY, g);
         //движок

DrawningWarmlyLocomotive.java

@@ -5,6 +5,12 @@ public class DrawningWarmlyLocomotive extends DrawningLocomotive{
         super(speed, weight, bodyColor,  140,  70);
         Locomotive = new EntityWarmlyLocomotive(speed, weight, bodyColor, extraColor, pipe, storage);
     }
+
+    public DrawningWarmlyLocomotive(EntityLocomotive locomotive, IDrawningExtra extra) {
+        super(locomotive, extra);
+        Locomotive = locomotive;
+    }
+
     @Override
     public void DrawTransport(Graphics2D g)
     {

EntityWithExtraCreator.java

@@ -0,0 +1,47 @@
+import java.util.Random;
+
+public class EntityWithExtraCreator <T extends EntityLocomotive, U extends IDrawningExtra> {
+    private T[] entityArr;
+    private U[] extraArr;
+
+    int entitiesCount = 0;
+    int extraCount = 0;
+
+    public EntityWithExtraCreator(int countEntities, int countExtra) {
+        entityArr = (T[]) new Object[countEntities];
+        extraArr = (U[]) new Object[countExtra];
+    }
+
+    public int Insert(T entityLocomotive) {
+        if(entitiesCount < entityArr.length) {
+            entityArr[entitiesCount] = entityLocomotive;
+            entitiesCount++;
+            return entitiesCount - 1;
+        }
+        return -1;
+    }
+
+    public int Insert (U extra) {
+        if(extraCount < extraArr.length) {
+            extraArr[extraCount] = extra;
+            extraCount++;
+            return extraCount - 1;
+        }
+        return -1;
+    }
+
+    public DrawningLocomotive getEntityWithExtra() {
+        Random random = new Random();
+        int getEntityRandomIndex = random.nextInt(entityArr.length);
+        int getExtraRandomIndex = random.nextInt(extraArr.length);
+
+        T locomotive = entityArr[getEntityRandomIndex];
+        U extra = extraArr[getExtraRandomIndex];
+
+        if (locomotive instanceof EntityWarmlyLocomotive) {
+            return  new DrawningWarmlyLocomotive(locomotive, extra);
+        }
+        return new DrawningLocomotive(locomotive, extra);
+    }
+
+}

ExtraRoundWheelDraw.java

@@ -2,6 +2,7 @@ import java.awt.*;
 
 public class ExtraRoundWheelDraw  implements  IDrawningExtra{
     private WheelsCount wheelsCount = WheelsCount.Two;
+    private ExtraWheelsDraw extraWheelsDraw;
     public void setExtraNum(int num) {
         switch (num) {
             case 0: {
@@ -17,11 +18,13 @@ public class ExtraRoundWheelDraw  implements  IDrawningExtra{
         }
     }
 
-    public ExtraRoundWheelDraw (int num) {
+    public ExtraRoundWheelDraw (int num, Color bodyColor) {
         setExtraNum(num);
+        extraWheelsDraw = new ExtraWheelsDraw(num, bodyColor);
     }
 
     public void DrawExtra(int startPosX, int startPosY, Graphics2D g) {
+        extraWheelsDraw.DrawExtra(startPosX, startPosY, g);
         g.setColor(Color.BLACK);
         g.fillOval(startPosX + 5, startPosY + 35, 10, 10);
         g.fillOval(startPosX + 95, startPosY + 35, 10, 10);

ExtraStarWheelDraw.java

@@ -2,6 +2,7 @@ import java.awt.*;
 
 public class ExtraStarWheelDraw implements IDrawningExtra{
     private WheelsCount wheelsCount = WheelsCount.Two;
+    private ExtraWheelsDraw extraWheelsDraw;
     public void setExtraNum(int num) {
         switch (num) {
             case 0: {
@@ -17,11 +18,13 @@ public class ExtraStarWheelDraw implements IDrawningExtra{
         }
     }
 
-    public ExtraStarWheelDraw (int num) {
+    public ExtraStarWheelDraw (int num, Color bodyColor) {
         setExtraNum(num);
+        extraWheelsDraw = new ExtraWheelsDraw(num, bodyColor);
     }
 
     public void DrawExtra(int startPosX, int startPosY, Graphics2D g) {
+        extraWheelsDraw.DrawExtra(startPosX, startPosY, g);
         DrawStarOnWheel(startPosX, startPosY + 30, g);
         DrawStarOnWheel(startPosX + 90, startPosY + 30, g);
         switch (wheelsCount) {

FormLocomotive.java

@@ -38,8 +38,8 @@ public class FormLocomotive extends JComponent{
         modifiedButton.addActionListener(e -> {
             Random rnd = new Random();
 
-            Color colorFirst = JColorChooser.showDialog(null, "Цвет", new Color(rnd.nextInt(256), rnd.nextInt(256),rnd.nextInt(256)));
+            Color colorFirst = JColorChooser.showDialog(null, "Color", new Color(rnd.nextInt(256), rnd.nextInt(256),rnd.nextInt(256)));
-            Color colorSecond = JColorChooser.showDialog(null, "Цвет", new Color(rnd.nextInt(256), rnd.nextInt(256),rnd.nextInt(256)));
+            Color colorSecond = JColorChooser.showDialog(null, "Color", new Color(rnd.nextInt(256), rnd.nextInt(256),rnd.nextInt(256)));
 
             _locomotive = new DrawningWarmlyLocomotive(rnd.nextInt(200) + 100, rnd.nextInt(1000) + 1000,
                     colorFirst,