Index: firmware/App/Modes/ModePostTreat.c
===================================================================
diff -u -r9bb9f9bb8b96f057b2ad8dd42c5ac4f0861c1d90 -r65d0f3fed34c836640f502d80a8aeb4e10ddf5ac
--- firmware/App/Modes/ModePostTreat.c	(.../ModePostTreat.c)	(revision 9bb9f9bb8b96f057b2ad8dd42c5ac4f0861c1d90)
+++ firmware/App/Modes/ModePostTreat.c	(.../ModePostTreat.c)	(revision 65d0f3fed34c836640f502d80a8aeb4e10ddf5ac)
@@ -7,8 +7,8 @@
 *
 * @file    ModePostTreat.c
 *
-* @author  (last)      Quang Nguyen
-* @date    (last)      24-Aug-2021
+* @author  (last)      Dara Navaei
+* @date    (last)      12-Nov-2021
 *
 * @author  (original)  Dara Navaei
 * @date    (original)  05-Nov-2019
@@ -63,7 +63,8 @@
 static BOOL disposableRemovalConfirmed = FALSE;                     ///< Flag indicates user confirms disposable removal.
 
 static BOOL isDrainStarted = FALSE;                                 ///< Flag indicates a drain operation has been started.
-static BOOL isDrainCompleted = FALSE;
+static BOOL isDrainCompleted = FALSE;                               ///< Flag indicates drain has completed.
+static BOOL rinseConcentrateLines = FALSE;                          ///< FLag indicates to rinse concentrate lines.
 
 static HD_POST_TREATMENT_STATE_T currentPostTreatmentState;         ///< Current state of post-treatment mode state machine.
 static DRAIN_STATE_T currentDrainReservoirState;                    ///< Current drain reservoir state.
@@ -75,8 +76,9 @@
 
 // ********** private function prototypes **********
 
-static HD_POST_TREATMENT_STATE_T handlePostTreatmentPatientDisconnectionState( void );
 static HD_POST_TREATMENT_STATE_T handlePostTreatmentDrainReservoirsState( void );
+static HD_POST_TREATMENT_STATE_T handlePostTreatmentPatientDisconnectionState( void );
+static HD_POST_TREATMENT_STATE_T handlePostTreatmentDisposableRemovalState( void );
 static HD_POST_TREATMENT_STATE_T handlePostTreatmentVerifyState( void );
 
 static void execDrainReservoirs( void );
@@ -99,7 +101,8 @@
     disposableRemovalConfirmed = FALSE;
     isDrainStarted = FALSE;
     isDrainCompleted = FALSE;
-    currentPostTreatmentState = HD_POST_TREATMENT_PATIENT_DISCONNECTION_STATE;
+    rinseConcentrateLines = FALSE;
+    currentPostTreatmentState = HD_POST_TREATMENT_DRAIN_RESERVOIRS_STATE;
     currentDrainReservoirState = DRAIN_FIRST_RESERVOIR_START_STATE;
     postTreatmentPublishTimerCounter = 0;
 
@@ -113,9 +116,9 @@
  * post-treatment mode.
  * @details Inputs: none
  * @details Outputs: none
- * @return none
+ * @return initial state
  *************************************************************************/
-void transitionToPostTreatmentMode( void )
+U32 transitionToPostTreatmentMode( void )
 {
     initPostTreatmentMode();
 
@@ -141,6 +144,8 @@
     signalDialInPumpHardStop();
 
     collectTreatmentLogData();
+
+    return currentPostTreatmentState;
 }
 
 /*********************************************************************//**
@@ -156,12 +161,16 @@
 
     switch ( currentPostTreatmentState )
     {
+        case HD_POST_TREATMENT_DRAIN_RESERVOIRS_STATE:
+            currentPostTreatmentState = handlePostTreatmentDrainReservoirsState();
+            break;
+
         case HD_POST_TREATMENT_PATIENT_DISCONNECTION_STATE:
             currentPostTreatmentState = handlePostTreatmentPatientDisconnectionState();
             break;
 
-        case HD_POST_TREATMENT_DRAIN_RESERVOIRS_STATE:
-            currentPostTreatmentState = handlePostTreatmentDrainReservoirsState();
+        case HD_POST_TREATMENT_DISPOSABLE_REMOVAL_STATE:
+            currentPostTreatmentState = handlePostTreatmentDisposableRemovalState();
             break;
 
         case HD_POST_TREATMENT_VERIFY_STATE:
@@ -175,7 +184,9 @@
 
     if ( ++postTreatmentPublishTimerCounter >= getU32OverrideValue( &postTreatmentModePublishInterval ) )
     {
-        broadcastPostTreatmentState( currentPostTreatmentState );
+        U32 state = (U32)currentPostTreatmentState;
+
+        broadcastData( MSG_ID_HD_POST_TREATMENT_STATE, COMM_BUFFER_OUT_CAN_HD_BROADCAST, (U08*)&state, sizeof( U32 ) );
         postTreatmentPublishTimerCounter = 0;
     }
 
@@ -317,6 +328,28 @@
 
 /*********************************************************************//**
  * @brief
+ * The handlePostTreatmentDrainReservoirsState function executes drain reservoirs
+ * operation state machine and transition to next state once completed.
+ * @details Inputs: currentDrainReservoirState
+ * @details Outputs: executed drain reservoirs operation state machine.
+ * @return current state (sub-mode)
+ *************************************************************************/
+static HD_POST_TREATMENT_STATE_T handlePostTreatmentDrainReservoirsState( void )
+{
+    HD_POST_TREATMENT_STATE_T state = HD_POST_TREATMENT_DRAIN_RESERVOIRS_STATE;
+
+    execDrainReservoirs();
+
+    if ( DRAIN_COMPLETE_STATE == currentDrainReservoirState )
+    {
+        state = HD_POST_TREATMENT_PATIENT_DISCONNECTION_STATE;
+    }
+
+    return state;
+}
+
+/*********************************************************************//**
+ * @brief
  * The handlePostTreatmentPatientConnectionState function waits for user
  * confirms patient disconnection.
  * @details Inputs: none
@@ -338,11 +371,11 @@
         if ( TRUE == patientDisconnectionConfirmed )
         {
             patientDisconnectionConfirmed = FALSE;
-            state = HD_POST_TREATMENT_DRAIN_RESERVOIRS_STATE;
+            state = HD_POST_TREATMENT_DISPOSABLE_REMOVAL_STATE;
 
-            for ( valve = VDI; valve < NUM_OF_VALVES; ++valve )
+            for ( valve = VDI; valve < NUM_OF_VALVES; valve++ )
             {
-                homeValve( valve );
+                setValvePosition( valve, VALVE_POSITION_A_INSERT_EJECT );
             }
 
             homeBloodPump();
@@ -365,28 +398,24 @@
 
 /*********************************************************************//**
  * @brief
- * The handlePostTreatmentDrainReservoirsState function executes drain reservoirs
- * operation state machine and transition to next state once completed.
- * @details Inputs: currentDrainReservoirState
- * @details Outputs: executed drain reservoirs operation state machine.
+ * The handlePostTreatmentDisposableRemovalState function waits for user
+ * to confirm disposable removal.
+ * @details Inputs: disposableRemovalConfirmed
+ * @details Outputs: processed disposable removal confirmation
  * @return current state (sub-mode)
  *************************************************************************/
-static HD_POST_TREATMENT_STATE_T handlePostTreatmentDrainReservoirsState( void )
+static HD_POST_TREATMENT_STATE_T handlePostTreatmentDisposableRemovalState( void )
 {
-    HD_POST_TREATMENT_STATE_T state = HD_POST_TREATMENT_DRAIN_RESERVOIRS_STATE;
+    HD_POST_TREATMENT_STATE_T state = HD_POST_TREATMENT_DISPOSABLE_REMOVAL_STATE;
 
-    execDrainReservoirs();
-
-    if ( DRAIN_COMPLETE_STATE == currentDrainReservoirState )
-    {
 #ifdef SKIP_UI_INTERACTION
-        disposableRemovalConfirmed = TRUE;
+    disposableRemovalConfirmed = TRUE;
 #endif
-        if ( TRUE == disposableRemovalConfirmed )
-        {
-            disposableRemovalConfirmed = FALSE;
-            state = HD_POST_TREATMENT_VERIFY_STATE;
-        }
+
+    if ( TRUE == disposableRemovalConfirmed )
+    {
+        disposableRemovalConfirmed = FALSE;
+        state = HD_POST_TREATMENT_VERIFY_STATE;
     }
 
     return state;
@@ -419,7 +448,7 @@
         {
             activateAlarmNoData( ALARM_ID_HD_SYRINGE_DETECTED );
         }
-
+        // Wait for cartridge and syringe to be removed and for DG to finish any drains/fills.
         if ( ( TRUE == isCartridgeRemoved ) && ( FALSE == isSyringeDetected() ) )
         {
             cmdStopDG();
@@ -491,7 +520,7 @@
         }
         else
         {
-            cmdStartDGDrain( DRAIN_RESERVOIR_TO_VOLUME_ML, TRUE );
+            cmdStartDGDrain( DRAIN_RESERVOIR_TO_VOLUME_ML, TRUE, rinseConcentrateLines, TRUE );
         }
     }
 
@@ -512,8 +541,14 @@
 
     if ( ( DG_MODE_GENE == getDGOpMode() ) && ( DG_GEN_IDLE_MODE_STATE_FLUSH_WATER == getDGSubMode() ) )
     {
-        state = DRAIN_FIRST_RESERVOIR_STATE;
-        cmdStartDGDrain( DRAIN_RESERVOIR_TO_VOLUME_ML, TRUE );
+        cmdSetDGActiveReservoir( DG_RESERVOIR_2 );
+
+        if ( TRUE == hasDGCompletedReservoirSwitch() )
+        {
+            state = DRAIN_FIRST_RESERVOIR_STATE;
+            rinseConcentrateLines = TRUE;
+            cmdStartDGDrain( DRAIN_RESERVOIR_TO_VOLUME_ML, TRUE, rinseConcentrateLines, TRUE );
+        }
     }
 
     return state;
@@ -571,7 +606,8 @@
     if ( TRUE == hasDGCompletedReservoirSwitch() )
     {
         state = DRAIN_SECOND_RESERVOIR_STATE;
-        cmdStartDGDrain( DRAIN_RESERVOIR_TO_VOLUME_ML, TRUE );
+        rinseConcentrateLines = FALSE;
+        cmdStartDGDrain( DRAIN_RESERVOIR_TO_VOLUME_ML, TRUE, rinseConcentrateLines, TRUE );
     }
 
     return state;
@@ -594,7 +630,7 @@
         isDrainStarted = TRUE;
     }
 
-    // Drain has started and DG goes to re-circ mode flush lines state means drain completed
+    // Drain has started and DG goes to generation idle mode flush lines state means drain completed
     if ( ( TRUE == isDrainStarted ) && ( DG_GEN_IDLE_MODE_STATE_FLUSH_LINES == getDGSubMode() ) )
     {
         isDrainCompleted = TRUE;
@@ -606,6 +642,7 @@
         {
             isDrainStarted = FALSE;
             isDrainCompleted = FALSE;
+            rinseConcentrateLines = FALSE;
             state = DRAIN_COMPLETE_STATE;
         }
     }