Index: firmware/App/Modes/TreatmentEnd.c
===================================================================
diff -u -r844f98879b7425c207b58562e623ab960adbc357 -r8bd1ae47aa13a843aa8abd6321ddc050deacb4a6
--- firmware/App/Modes/TreatmentEnd.c	(.../TreatmentEnd.c)	(revision 844f98879b7425c207b58562e623ab960adbc357)
+++ firmware/App/Modes/TreatmentEnd.c	(.../TreatmentEnd.c)	(revision 8bd1ae47aa13a843aa8abd6321ddc050deacb4a6)
@@ -5,13 +5,13 @@
 * THIS CODE MAY NOT BE COPIED OR REPRODUCED IN ANY FORM, IN PART OR IN
 * WHOLE, WITHOUT THE EXPLICIT PERMISSION OF THE COPYRIGHT OWNER.
 *
-* @file TreatmentEnd.c
+* @file    TreatmentEnd.c
 *
-* @author (last) Sean Nash
-* @date (last) 04-Feb-2021
+* @author  (last)      Sean Nash
+* @date    (last)      19-Aug-2021
 *
-* @author (original) Sean
-* @date (original) 04-Feb-2021
+* @author  (original)  Sean Nash
+* @date    (original)  05-Feb-2021
 *
 ***************************************************************************/
 
@@ -125,12 +125,13 @@
     signalDialOutPumpHardStop();
     stopSyringePump();
     setBloodPumpTargetFlowRate( TX_END_BP_FLOW_RATE_ML_MIN, MOTOR_DIR_FORWARD, PUMP_CONTROL_MODE_CLOSED_LOOP );
+    cmdStopDGTrimmerHeater();
 
     // Continue air trap control
     startAirTrapControl();
 
     // Set user alarm recovery actions allowed in this sub-mode
-    setAlarmUserActionEnabled( ALARM_USER_ACTION_RESUME, FALSE );
+    setAlarmUserActionEnabled( ALARM_USER_ACTION_RESUME, TRUE );  // Resume in this sub-mode indicates resume wait for user rinseback request (not resume treatment)
     setAlarmUserActionEnabled( ALARM_USER_ACTION_RINSEBACK, TRUE );
     setAlarmUserActionEnabled( ALARM_USER_ACTION_END_TREATMENT, TRUE );
 }
@@ -146,6 +147,8 @@
  *************************************************************************/
 static void setupForTxEndWait4RinsebackState( void )
 {
+    // No resume alarm action allowed
+    setAlarmUserActionEnabled( ALARM_USER_ACTION_RESUME, FALSE );
     // Open VBA and VBV valves to allow flow from saline bag and to patient venous line
     setValvePosition( VBA, VALVE_POSITION_B_OPEN );
     setValvePosition( VBV, VALVE_POSITION_B_OPEN );
@@ -167,6 +170,8 @@
  *************************************************************************/
 static void setupForTxEndPausedState( void )
 {
+    // Resume alarm action will take us back to TxEndWait4RinsebackState
+    setAlarmUserActionEnabled( ALARM_USER_ACTION_RESUME, TRUE );
     // Stop blood pump
     signalBloodPumpHardStop();
     // Close arterial and venous lines
@@ -219,6 +224,19 @@
 {
     TREATMENT_END_STATE_T result = TREATMENT_END_WAIT_FOR_RINSEBACK_STATE;
 
+    // Set reservoir 2 as active so we can drain reservoir 1 ASAP    // TODO - I think this should be handled by treatment reservoir mgmt state machine
+    if ( ( DG_MODE_GENE == getDGOpMode() ) && ( DG_GEN_IDLE_MODE_STATE_FLUSH_WATER == getDGSubMode() ) )
+    {
+        if ( getDGInactiveReservoir() != DG_RESERVOIR_1 )
+        {
+            cmdSetDGActiveReservoir( DG_RESERVOIR_2 );
+        }
+        else
+        {
+            cmdStartDGDrain( DRAIN_RESERVOIR_TO_VOLUME_ML, TRUE, FALSE, TRUE );
+        }
+    }
+
     // Is alarm stop or user reconnect requested?
     if ( TRUE == doesAlarmStatusIndicateStop() )
     {
@@ -279,12 +297,12 @@
         // Ensure we do not sit in stopped state for too long
         if ( bloodSittingTimerCtr > WARN_TIME_BLOOD_SITTING )
         {
-            activateAlarmNoData( ALARM_ID_BLOOD_SITTING_WARNING );
+            activateAlarmNoData( ALARM_ID_END_OF_TREATMENT_ALARM );
         }
         if ( bloodSittingTimerCtr > MAX_TIME_BLOOD_SITTING )
         {
             // Raise the alarm
-            activateAlarmNoData( ALARM_ID_TREATMENT_STOPPED_NO_RINSEBACK );
+            activateAlarmNoData( ALARM_ID_END_OF_TREATMENT_HIGH );
         }
     }