From ece342f037ef0298e6e144aeb556d36d20419dcc Mon Sep 17 00:00:00 2001
From: Asbelos <asbelos@btinternet.com>
Date: Thu, 31 Mar 2022 22:19:13 +0100
Subject: [PATCH] DC power fix

---
 MotorDriver.cpp  |  4 ---
 TrackManager.cpp | 69 +++++++++++++++++++++++++++++++++---------------
 TrackManager.h   |  1 +
 3 files changed, 48 insertions(+), 26 deletions(-)

diff --git a/MotorDriver.cpp b/MotorDriver.cpp
index ecad826..517e2dd 100644
--- a/MotorDriver.cpp
+++ b/MotorDriver.cpp
@@ -97,10 +97,6 @@ bool MotorDriver::isPWMCapable() {
 void MotorDriver::setPower(POWERMODE mode) {
   bool on=mode==POWERMODE::ON;
   if (on) {
-    // toggle brake before turning power on - resets overcurrent error
-    // on the Pololu board if brake is wired to ^D2.
-    setBrake(true);
-    setBrake(false);
     IODevice::write(powerPin,HIGH);
   }
   else IODevice::write(powerPin,LOW);
diff --git a/TrackManager.cpp b/TrackManager.cpp
index 13a3280..5618644 100644
--- a/TrackManager.cpp
+++ b/TrackManager.cpp
@@ -134,17 +134,8 @@ bool TrackManager::setTrackMode(byte trackToSet, TRACK_MODE mode, int16_t dcAddr
         // DC tracks need to be given speed of the throttle for that cab address
         // otherwise will not match other tracks on same cab.
         // This also needs to allow for inverted DCX
-
-        int16_t speed1=DCC::getThrottleSpeed(dcAddr);
-        byte speedByte;
-        if (speed1<0) speedByte=0;
-        else {
-            speedByte=speed1;
-            bool direction=DCC::getThrottleDirection(dcAddr);
-            if (mode==TRACK_MODE_DCX) direction=!direction;  
-            if (direction) speedByte|=0x80;
-        }
-        track[trackToSet]->setDCSignal(speedByte); 
+        applyDCSpeed(trackToSet);
+ 
     }
     else {
         // DCC tracks need to have the brake set off or they will not work.
@@ -168,6 +159,20 @@ bool TrackManager::setTrackMode(byte trackToSet, TRACK_MODE mode, int16_t dcAddr
     return true; 
 }
 
+void TrackManager::applyDCSpeed(byte t) {
+    
+        int16_t speed1=DCC::getThrottleSpeed(trackDCAddr[t]);
+        byte speedByte;
+        if (speed1<0) speedByte=0;
+        else {
+            speedByte=speed1;
+            bool direction=DCC::getThrottleDirection(trackDCAddr[t]);
+            if (trackMode[t]==TRACK_MODE_DCX) direction=!direction;  
+            if (direction) speedByte|=0x80;
+        }
+        track[t]->setDCSignal(speedByte);
+}
+
 bool TrackManager::parseJ(Print *stream, int16_t params, int16_t p[])
 {
     
@@ -240,18 +245,38 @@ MotorDriver * TrackManager::getProgDriver() {
     FOR_EACH_TRACK(t)
         if (trackMode[t]==TRACK_MODE_PROG) return track[t];
     return NULL;
-}        
+} 
+
 void TrackManager::setPower2(bool setProg,POWERMODE mode) {
-    if (setProg) {
-         APPLY_BY_MODE(TRACK_MODE_PROG,setPower(mode))
-    }
-    else {
-        mainPowerGuess=mode; 
-        FOR_EACH_TRACK(t)
-        if (track[t] 
-             && trackMode[t]!=TRACK_MODE_OFF
-             && trackMode[t]!=TRACK_MODE_PROG 
-             ) track[t]->setPower(mode);
+    if (!setProg) mainPowerGuess=mode; 
+    FOR_EACH_TRACK(t) {
+        MotorDriver * driver=track[t]; 
+        if (!driver) continue; 
+        switch (trackMode[t]) {
+            case TRACK_MODE_MAIN:
+                if (setProg) break; 
+                // toggle brake before turning power on - resets overcurrent error
+                // on the Pololu board if brake is wired to ^D2.   
+                driver->setBrake(true);
+                driver->setBrake(false); // DCC runs with brake off
+                driver->setPower(mode);  
+                break; 
+            case TRACK_MODE_DC:
+            case TRACK_MODE_DCX:
+                if (setProg) break; 
+                driver->setBrake(true); // DC starts with brake on
+                applyDCSpeed(t);        // speed match DCC throttles
+                driver->setPower(mode);
+                break;  
+            case TRACK_MODE_PROG:
+                if (!setProg) break; 
+                driver->setBrake(true);
+                driver->setBrake(false);
+                driver->setPower(mode);
+                break;  
+            case TRACK_MODE_OFF:
+                break;
+        }
     }
 }
   
diff --git a/TrackManager.h b/TrackManager.h
index dde4b90..8b12f86 100644
--- a/TrackManager.h
+++ b/TrackManager.h
@@ -78,6 +78,7 @@ class TrackManager {
     static byte lastTrack;
     static byte nextCycleTrack;
     static POWERMODE mainPowerGuess;
+    static void applyDCSpeed(byte t);
 
     static MotorDriver* track[MAX_TRACKS];
     static TRACK_MODE trackMode[MAX_TRACKS];