From 784934024ef9a0b7db67fcf831a02688102e021b Mon Sep 17 00:00:00 2001
From: Harald Barth <haba@kth.se>
Date: Thu, 23 Nov 2023 10:47:43 +0100
Subject: [PATCH] Bugfix: Do not turn off all tracks on change ; give better
 power messages

---
 CommandDistributor.cpp | 53 +++++++++++++++++++++++++++++++++++-------
 TrackManager.cpp       | 16 +++++++------
 TrackManager.h         |  5 +++-
 3 files changed, 57 insertions(+), 17 deletions(-)

diff --git a/CommandDistributor.cpp b/CommandDistributor.cpp
index 1f11e28..a30fcd3 100644
--- a/CommandDistributor.cpp
+++ b/CommandDistributor.cpp
@@ -254,21 +254,56 @@ void  CommandDistributor::broadcastPower() {
     if (TrackManager::getPower(t, pstr))
       broadcastReply(COMMAND_TYPE, F("<p%s>\n"),pstr);
 
+  byte trackcount=0;
+  byte oncount=0;
+  byte offcount=0;
+  for(byte t=0; t<TrackManager::MAX_TRACKS; t++) {
+    if (TrackManager::isActive(t)) {
+      trackcount++;
+      // do not call getPower(t) unless isActive(t)!
+      if (TrackManager::getPower(t) == POWERMODE::ON)
+	oncount++;
+      else
+	offcount++;
+    }
+  }
+  //DIAG(F("t=%d on=%d off=%d"), trackcount, oncount, offcount);
+
+  char state='2';
+  if (oncount==0 || offcount == trackcount)
+    state = '0';
+  else if (oncount == trackcount) {
+    state = '1';
+  }
+
+  // additional info about MAIN, PROG and JOIN
   bool main=TrackManager::getMainPower()==POWERMODE::ON;
   bool prog=TrackManager::getProgPower()==POWERMODE::ON;
   bool join=TrackManager::isJoined();
+  //DIAG(F("m=%d p=%d j=%d"), main, prog, join);
   const FSH * reason=F("");
-  char state='1';
-  if (main && prog && join) reason=F(" JOIN");
-  else if (main && prog);
-  else if (main) reason=F(" MAIN");
-  else if (prog) reason=F(" PROG");
-  else state='0';
-  broadcastReply(COMMAND_TYPE, F("<p%c%S>\n"),state,reason);
+  if (join) {
+    reason = F("JOIN");
+    broadcastReply(COMMAND_TYPE, F("<p1 %S>\n"),reason);
+  } else {
+    if (main) {
+      //reason = F("MAIN");
+      broadcastReply(COMMAND_TYPE, F("<p1 MAIN>\n"));
+    }
+    if (prog) {
+      //reason = F("PROG");
+      broadcastReply(COMMAND_TYPE, F("<p1 PROG>\n"));
+    }
+  }
+
+  if (state != '2')
+    broadcastReply(COMMAND_TYPE, F("<p%c>\n"),state);
 #ifdef CD_HANDLE_RING
-  broadcastReply(WITHROTTLE_TYPE, F("PPA%c\n"), main?'1':'0');
+  // send '1' if all main are on, otherwise global state (which in that case is '0' or '2')
+  broadcastReply(WITHROTTLE_TYPE, F("PPA%c\n"), main?'1': state);
 #endif
-  LCD(2,F("Power %S%S"),state=='1'?F("On"):F("Off"),reason);  
+
+  LCD(2,F("Power %S %S"),state=='1'?F("On"): ( state=='0'? F("Off") : F("SC") ),reason);
 }
 
 void CommandDistributor::broadcastRaw(clientType type, char * msg) {
diff --git a/TrackManager.cpp b/TrackManager.cpp
index 846c163..ac44c76 100644
--- a/TrackManager.cpp
+++ b/TrackManager.cpp
@@ -54,7 +54,6 @@ const int16_t HASH_KEYWORD_INV = 11857;
 MotorDriver * TrackManager::track[MAX_TRACKS];
 int16_t TrackManager::trackDCAddr[MAX_TRACKS];
 
-POWERMODE TrackManager::mainPowerGuess=POWERMODE::OFF;
 byte TrackManager::lastTrack=0;
 bool TrackManager::progTrackSyncMain=false; 
 bool TrackManager::progTrackBoosted=false; 
@@ -210,6 +209,9 @@ void TrackManager::setDCSignal(int16_t cab, byte speedbyte) {
 bool TrackManager::setTrackMode(byte trackToSet, TRACK_MODE mode, int16_t dcAddr) {
     if (trackToSet>lastTrack || track[trackToSet]==NULL) return false;
 
+    // Remember track mode we came from for later
+    TRACK_MODE oldmode = track[trackToSet]->getMode();
+
     //DIAG(F("Track=%c Mode=%d"),trackToSet+'A', mode);
     // DC tracks require a motorDriver that can set brake!
     if (mode & TRACK_MODE_DC) {
@@ -262,7 +264,7 @@ bool TrackManager::setTrackMode(byte trackToSet, TRACK_MODE mode, int16_t dcAddr
 	  track[t]->setPower(POWERMODE::OFF);
 	  track[t]->setMode(TRACK_MODE_NONE);
 	  track[t]->makeProgTrack(false);     // revoke prog track special handling
-    streamTrackState(NULL,t);
+	  streamTrackState(NULL,t);
 	}
       track[trackToSet]->makeProgTrack(true); // set for prog track special handling
     } else {
@@ -270,7 +272,6 @@ bool TrackManager::setTrackMode(byte trackToSet, TRACK_MODE mode, int16_t dcAddr
     }
     track[trackToSet]->setMode(mode);
     trackDCAddr[trackToSet]=dcAddr;
-    streamTrackState(NULL,trackToSet);
 
     // When a track is switched, we must clear any side effects of its previous 
     // state, otherwise trains run away or just dont move.
@@ -337,10 +338,11 @@ bool TrackManager::setTrackMode(byte trackToSet, TRACK_MODE mode, int16_t dcAddr
         applyDCSpeed(trackToSet);
     }
 
-    // Normal running tracks are set to the global power state 
-    track[trackToSet]->setPower(
-      (mode & (TRACK_MODE_MAIN | TRACK_MODE_DC | TRACK_MODE_EXT | TRACK_MODE_BOOST)) ?
-        mainPowerGuess : POWERMODE::OFF);
+    // Turn off power if we changed the mode of this track
+    if (mode != oldmode)
+      track[trackToSet]->setPower(POWERMODE::OFF);
+    streamTrackState(NULL,trackToSet);
+
     //DIAG(F("TrackMode=%d"),mode);
     return true; 
 }
diff --git a/TrackManager.h b/TrackManager.h
index f2a357d..12adb4e 100644
--- a/TrackManager.h
+++ b/TrackManager.h
@@ -77,9 +77,13 @@ class TrackManager {
     static void loop();
     static POWERMODE getMainPower();
     static POWERMODE getProgPower();
+    static inline POWERMODE getPower(byte t) { return track[t]->getPower(); }
     static bool getPower(byte t, char s[]);
     static void setJoin(bool join);
     static bool isJoined() { return progTrackSyncMain;}
+    static inline bool isActive (byte tr) {
+      if (tr > lastTrack) return false;
+      return track[tr]->getMode() & (TRACK_MODE_MAIN|TRACK_MODE_PROG|TRACK_MODE_DC|TRACK_MODE_BOOST|TRACK_MODE_EXT);}
     static void setJoinRelayPin(byte joinRelayPin);
     static void sampleCurrent();
     static void reportGauges(Print* stream);
@@ -108,7 +112,6 @@ class TrackManager {
     static void addTrack(byte t, MotorDriver* driver);
     static byte lastTrack;
     static byte nextCycleTrack;
-    static POWERMODE mainPowerGuess;
     static void applyDCSpeed(byte t);
 
     static int16_t trackDCAddr[MAX_TRACKS];  // dc address if TRACK_MODE_DC