Index: firmware/App/Modes/ModePostTreat.c
===================================================================
diff -u -r8c00197ce69e80f1967aa3f2eb36beb3573f36f3 -r40e8c3a97babf6910cc6e9602d691f8e56691c17
--- firmware/App/Modes/ModePostTreat.c	(.../ModePostTreat.c)	(revision 8c00197ce69e80f1967aa3f2eb36beb3573f36f3)
+++ firmware/App/Modes/ModePostTreat.c	(.../ModePostTreat.c)	(revision 40e8c3a97babf6910cc6e9602d691f8e56691c17)
@@ -1,37 +1,93 @@
 /**************************************************************************
 *
-* Copyright (c) 2019-2020 Diality Inc. - All Rights Reserved.
+* Copyright (c) 2019-2022 Diality Inc. - All Rights Reserved.
 *
 * 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 ModePostTreat.c
+* @file    ModePostTreat.c
 *
-* @author (last) Sean Nash
-* @date (last) 08-Oct-2020
+* @author  (last)      Dara Navaei
+* @date    (last)      31-Mar-2022
 *
-* @author (original) Dara Navaei
-* @date (original) 05-Nov-2019
+* @author  (original)  Dara Navaei
+* @date    (original)  05-Nov-2019
 *
 ***************************************************************************/
 
+#include <string.h>  // For memset()
+
 #include "AlarmLamp.h"
+#include "Buttons.h"
 #include "BloodFlow.h"
 #include "DialInFlow.h"
 #include "DialOutFlow.h"
-#include "Buttons.h"
+#include "DGInterface.h"
+#include "FPGA.h"
 #include "OperationModes.h"
+#include "PresOccl.h"
 #include "ModePostTreat.h"
+#include "ModeTreatment.h"
+#include "ModeTreatmentParams.h"
+#include "RTC.h"
+#include "SampleWater.h"
+#include "SystemCommMessages.h"
+#include "TaskGeneral.h"
+#include "Timers.h"
+#include "Valves.h"
 
 /**
  *  @addtogroup HDPostTreatmentMode
  *  @{
  */
 
+// ********** private definitions **********
+
+/// Interval (ms/task time) at which the post-treatment state data is published on the CAN bus.
+#define POST_TREATMENT_DATA_PUB_INTERVAL        ( 250 / TASK_GENERAL_INTERVAL )
+
+/// Post-Treatment drain reservoirs state machine.
+typedef enum Drain_States
+{
+    DRAIN_FIRST_RESERVOIR_START_STATE = 0,      ///< Drain first reservoir start state.
+    DRAIN_FIRST_RESERVOIR_STATE,                ///< Drain first reservoir state.
+    DRAIN_SWITCH_RESERVOIR_STATE,               ///< Switch reservoir state.
+    DRAIN_SECOND_RESERVOIR_STATE,               ///< Drain second reservoir state.
+    DRAIN_COMPLETE_STATE,                       ///< Drain complete state.
+    NUM_OF_DRAIN_STATES                         ///< Number of drain states.
+} DRAIN_STATE_T;
+
 // ********** private data **********
 
+static BOOL patientDisconnectionConfirmed = FALSE;                  ///< Flag indicates user confirms patient disconnection.
+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;                               ///< 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.
+static U32 postTreatmentPublishTimerCounter;                        ///< Timer counter used to schedule post-treatment data broadcast.
+/// Interval (in task intervals) at which to publish post-treatment mode data to CAN bus.
+static OVERRIDE_U32_T postTreatmentModePublishInterval = { POST_TREATMENT_DATA_PUB_INTERVAL, POST_TREATMENT_DATA_PUB_INTERVAL, POST_TREATMENT_DATA_PUB_INTERVAL, 0 };
+
+static TREATMENT_LOG_DATA_PAYLOAD_T treatmentLogData;               ///< Treatment parameters record for logging.
+
 // ********** private function prototypes **********
 
+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 );
+static BOOL isReservoirDrainStarted( void );
+static DRAIN_STATE_T handleDrainFirstReservoirStartState( void );
+static DRAIN_STATE_T handleDrainFirstReservoirState( void );
+static DRAIN_STATE_T handleDrainSwitchReservoirState( void );
+static DRAIN_STATE_T handleDrainSecondReservoirState( void );
+
 /*********************************************************************//**
  * @brief
  * The initPostTreatmentMode function initializes the Post-Treatment Mode module.
@@ -41,6 +97,17 @@
  *************************************************************************/
 void initPostTreatmentMode( void )
 {
+    patientDisconnectionConfirmed = FALSE;
+    disposableRemovalConfirmed = FALSE;
+    isDrainStarted = FALSE;
+    isDrainCompleted = FALSE;
+    rinseConcentrateLines = FALSE;
+    currentPostTreatmentState = HD_POST_TREATMENT_DRAIN_RESERVOIRS_STATE;
+    currentDrainReservoirState = DRAIN_FIRST_RESERVOIR_START_STATE;
+    postTreatmentPublishTimerCounter = 0;
+
+    // Reset treatment log data
+    memset( &treatmentLogData, 0x0, sizeof( TREATMENT_LOG_DATA_PAYLOAD_T ) );
 }
 
 /*********************************************************************//**
@@ -49,16 +116,36 @@
  * post-treatment mode.
  * @details Inputs: none
  * @details Outputs: none
- * @return none
+ * @return initial state
  *************************************************************************/
-void transitionToPostTreatmentMode( void )
+U32 transitionToPostTreatmentMode( void )
 {
-    // TODO - stop any DG fill that may be in progress from an aborted treatment
+    initPostTreatmentMode();
 
+    // Stop any DG fill that may be in progress from an aborted treatment
+    cmdStopDGFill();
+
     // Set user alarm recovery actions allowed in this mode
-    setAlarmUserActionEnabled( ALARM_USER_ACTION_RESUME, TRUE );
+    setAlarmUserActionEnabled( ALARM_USER_ACTION_RESUME, FALSE );
     setAlarmUserActionEnabled( ALARM_USER_ACTION_RINSEBACK, FALSE );
     setAlarmUserActionEnabled( ALARM_USER_ACTION_END_TREATMENT, FALSE );
+
+    cmdStopDGTrimmerHeater();
+
+    setValvePosition( VDI, VALVE_POSITION_C_CLOSE );
+    setValvePosition( VDO, VALVE_POSITION_C_CLOSE );
+    setValvePosition( VBA, VALVE_POSITION_C_CLOSE );
+    setValvePosition( VBV, VALVE_POSITION_C_CLOSE );
+    setValveAirTrap( STATE_CLOSED );
+
+    stopSyringePump();
+    signalBloodPumpHardStop();
+    signalDialOutPumpHardStop();
+    signalDialInPumpHardStop();
+
+    collectTreatmentLogData();
+
+    return currentPostTreatmentState;
 }
 
 /*********************************************************************//**
@@ -72,13 +159,162 @@
 {
     BOOL stop = isStopButtonPressed();
 
-    requestNewOperationMode( MODE_STAN );  // TODO - implement post treatment mode
+    switch ( currentPostTreatmentState )
+    {
+        case HD_POST_TREATMENT_DRAIN_RESERVOIRS_STATE:
+            currentPostTreatmentState = handlePostTreatmentDrainReservoirsState();
+            break;
 
-    return 0; // TODO - return current state
+        case HD_POST_TREATMENT_PATIENT_DISCONNECTION_STATE:
+            currentPostTreatmentState = handlePostTreatmentPatientDisconnectionState();
+            break;
+
+        case HD_POST_TREATMENT_DISPOSABLE_REMOVAL_STATE:
+            currentPostTreatmentState = handlePostTreatmentDisposableRemovalState();
+            break;
+
+        case HD_POST_TREATMENT_VERIFY_STATE:
+            currentPostTreatmentState = handlePostTreatmentVerifyState();
+            break;
+
+        default:
+            SET_ALARM_WITH_2_U32_DATA( ALARM_ID_HD_SOFTWARE_FAULT, SW_FAULT_ID_MODE_POST_TREATMENT_INVALID_STATE, currentPostTreatmentState );
+            break;
+    }
+
+    if ( ++postTreatmentPublishTimerCounter >= getU32OverrideValue( &postTreatmentModePublishInterval ) )
+    {
+        U32 state = (U32)currentPostTreatmentState;
+
+        broadcastData( MSG_ID_HD_POST_TREATMENT_STATE, COMM_BUFFER_OUT_CAN_HD_BROADCAST, (U08*)&state, sizeof( U32 ) );
+        postTreatmentPublishTimerCounter = 0;
+    }
+
+    return currentPostTreatmentState;
 }
 
 /*********************************************************************//**
  * @brief
+ * The collectTreatmentLogData function collects treatment data.
+ * @details Inputs: none
+ * @details Outputs: collected treatment data
+ * @return none
+ *************************************************************************/
+void collectTreatmentLogData( void )
+{
+    // Reset treatment log data
+    memset( &treatmentLogData, 0x0, sizeof( TREATMENT_LOG_DATA_PAYLOAD_T ) );
+
+    treatmentLogData.bloodFlowRate_mL_min = getTreatmentParameterU32( TREATMENT_PARAM_BLOOD_FLOW );
+    treatmentLogData.dialysateFlowRate_mL_min = getTreatmentParameterU32( TREATMENT_PARAM_DIALYSATE_FLOW );
+    treatmentLogData.treatmentDuration_sec = SEC_PER_MIN * getTreatmentParameterU32( TREATMENT_PARAM_TREATMENT_DURATION );
+    treatmentLogData.actualTreatmentDur_sec = getActualTreatmentTimeSecs();
+
+    treatmentLogData.acidConcentrate = getTreatmentParameterU32( TREATMENT_PARAM_ACID_CONCENTRATE );
+    treatmentLogData.bicarbConcentrate = getTreatmentParameterU32( TREATMENT_PARAM_BICARB_CONCENTRATE );
+
+    treatmentLogData.dialysateTemperature_degC = getTreatmentParameterF32( TREATMENT_PARAM_DIALYSATE_TEMPERATURE );
+    treatmentLogData.dialyzerType = getTreatmentParameterU32( TREATMENT_PARAM_DIALYZER_TYPE );
+    treatmentLogData.treatmentStartDateAndTime = getTreatmentStartTimeStamp();
+    treatmentLogData.treatmentEndDateAndTime = getTreatmentEndTimeStamp();
+
+    // There is treatment start time but treatment not completed means treatment ends early
+    if ( ( treatmentLogData.treatmentStartDateAndTime != 0 ) && ( FALSE == isTreatmentCompleted() ) )
+    {
+        treatmentLogData.treatmentEndDateAndTime = getRTCTimestamp();
+    }
+
+    treatmentLogData.avgBloodFlow_mL_min = getTreatmentAvgBloodFlowRate();
+    treatmentLogData.avgDialysateFlow_mL_min = getTreatmentAvgDialysateFlowRate();
+    treatmentLogData.dialysateVolumeUsed_L = treatmentLogData.avgDialysateFlow_mL_min * treatmentLogData.actualTreatmentDur_sec / ( ML_PER_LITER * SEC_PER_MIN );
+    treatmentLogData.avgDialysateTemperature_degC = getTreatmentAvgDialysateTemp();
+
+    // original UF Volume/Rate
+    treatmentLogData.originUFVolume_L = getUltrafiltrationVolumeOriginal();
+    treatmentLogData.originUFRate_mL_min = getUltrafiltrationRateOriginal();
+    // Target UF Volume/Rate
+    treatmentLogData.targetUFVolume_L = getTreatmentParameterF32( TREATMENT_PARAM_UF_VOLUME );
+    treatmentLogData.targetUFRate_mL_min = ( ( treatmentLogData.targetUFVolume_L * ML_PER_LITER * SEC_PER_MIN ) / treatmentLogData.treatmentDuration_sec );
+    // Actual UF Volume/Rate
+    if ( treatmentLogData.actualTreatmentDur_sec > 0 )
+    {
+        treatmentLogData.actualUFVolume_L = getUltrafiltrationVolumeCollected() / ML_PER_LITER;
+        treatmentLogData.actualUFRate_mL_min = getUltrafiltrationVolumeCollected() * SEC_PER_MIN / treatmentLogData.actualTreatmentDur_sec;
+    }
+    else
+    {
+        treatmentLogData.actualUFVolume_L = 0.0;
+        treatmentLogData.actualUFRate_mL_min = 0.0;
+    }
+
+    treatmentLogData.salineBolusVolume_mL = getTotalSalineBolusVolumeDelivered();
+
+    treatmentLogData.heparinBolusVolume_mL = getTreatmentParameterF32( TREATMENT_PARAM_HEPARIN_BOLUS_VOLUME );
+    treatmentLogData.heparinDispenseRate_mL_hr = getTreatmentParameterF32( TREATMENT_PARAM_HEPARIN_DISPENSE_RATE );
+    treatmentLogData.heparinPreStop_min = getTreatmentParameterU32( TREATMENT_PARAM_HEPARIN_PRE_STOP_TIME );
+    treatmentLogData.heparinDeliveredVolume_mL = getSyringePumpVolumeDelivered();
+    treatmentLogData.heparinType = getTreatmentParameterU32( TREATMENT_PARAM_HEPARIN_TYPE );
+
+    treatmentLogData.avgArterialPressure_mmHg = getTreatmentAvgArterialPressure();
+    treatmentLogData.avgVenousPressure_mmHg = getTreatmentAvgVenousPressure();
+
+    treatmentLogData.waterSampleTestResult = (U32)getSampleWaterResult();
+}
+
+/*********************************************************************//**
+ * @brief
+ * The sendTreatmentLogDataToUI sends treatment log data to UI.
+ * @details Inputs: treatmentLogData
+ * @details Outputs: Sent treatment log data to UI
+ * @return none
+ *************************************************************************/
+void sendTreatmentLogDataToUI( void )
+{
+    BOOL accepted = TRUE;
+    REQUEST_REJECT_REASON_CODE_T rejReason = REQUEST_REJECT_REASON_NONE;
+
+    // Send treatment log data to UI
+    sendTreatmentLogData( accepted, rejReason, &treatmentLogData );
+}
+
+/*********************************************************************//**
+ * @brief
+ * The signalUserConfirmPatientDisconnection signals post-treatment mode
+ * user has confirmed patient disconnection.
+ * @details Inputs: none
+ * @details Outputs: patientDisconnectedConfirm
+ * @return none
+ *************************************************************************/
+void signalUserConfirmPatientDisconnection( void )
+{
+    patientDisconnectionConfirmed = TRUE;
+}
+
+/*********************************************************************//**
+ * @brief
+ * The signalUserConfirmDisposableRemoval signals post-treatment mode
+ * user has confirmed disposable removal.
+ * @details Inputs: none
+ * @details Outputs: disposableRemovalConfirmed
+ * @return none
+ *************************************************************************/
+void signalUserConfirmDisposableRemoval( void )
+{
+    BOOL accepted = FALSE;
+    REQUEST_REJECT_REASON_CODE_T rejReason = REQUEST_REJECT_REASON_DRAIN_NOT_COMPLETE;
+
+    if ( DRAIN_COMPLETE_STATE == currentDrainReservoirState )
+    {
+        disposableRemovalConfirmed = TRUE;
+        accepted = TRUE;
+        rejReason = REQUEST_REJECT_REASON_NONE;
+    }
+
+    sendDisposableRemovalConfirmResponse( accepted, rejReason );
+}
+
+/*********************************************************************//**
+ * @brief
  * The signalAlarmActionToPostTreatmentMode function executes the given alarm action
  * as appropriate while in PostTreatment Mode.
  * @details Inputs: none
@@ -91,4 +327,382 @@
 
 }
 
+/*********************************************************************//**
+ * @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
+ * @details Outputs: processed patient disconnection confirmation
+ * @return current state (sub-mode)
+ *************************************************************************/
+static HD_POST_TREATMENT_STATE_T handlePostTreatmentPatientDisconnectionState( void )
+{
+    F32 const bolusVol = getTreatmentParameterF32( TREATMENT_PARAM_HEPARIN_BOLUS_VOLUME );
+    F32 const hepRate = getTreatmentParameterF32( TREATMENT_PARAM_HEPARIN_DISPENSE_RATE );
+    HD_POST_TREATMENT_STATE_T state = HD_POST_TREATMENT_PATIENT_DISCONNECTION_STATE;
+    VALVE_T valve;
+
+    if (( STATE_CLOSED == getSwitchStatus( FRONT_DOOR ) ) && (STATE_CLOSED == getSwitchStatus( PUMP_TRACK_SWITCH )))
+    {
+        if ( SW_CONFIG_ENABLE_VALUE == getSoftwareConfigStatus( SW_CONFIG_DISABLE_UI_INTERACTION ) )
+        {
+            patientDisconnectionConfirmed = TRUE;
+        }
+
+        if ( TRUE == patientDisconnectionConfirmed )
+        {
+            patientDisconnectionConfirmed = FALSE;
+            state = HD_POST_TREATMENT_DISPOSABLE_REMOVAL_STATE;
+
+            for ( valve = VDI; valve < NUM_OF_VALVES; valve++ )
+            {
+                setValvePosition( valve, VALVE_POSITION_A_INSERT_EJECT );
+            }
+
+            homeBloodPump();
+            homeDialInPump();
+            homeDialOutPump();
+
+            if ( ( bolusVol > 0.0 ) || ( hepRate > 0.0 ) )
+            {
+                retractSyringePump();
+            }
+        }
+    }
+    else
+    {
+        activateAlarmNoData( ALARM_ID_CARTRIDGE_DOOR_OPENED );
+    }
+
+    return state;
+}
+
+/*********************************************************************//**
+ * @brief
+ * 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 handlePostTreatmentDisposableRemovalState( void )
+{
+    HD_POST_TREATMENT_STATE_T state = HD_POST_TREATMENT_DISPOSABLE_REMOVAL_STATE;
+
+    if ( SW_CONFIG_ENABLE_VALUE == getSoftwareConfigStatus( SW_CONFIG_DISABLE_UI_INTERACTION ) )
+    {
+        disposableRemovalConfirmed = TRUE;
+    }
+
+    if ( TRUE == disposableRemovalConfirmed )
+    {
+        disposableRemovalConfirmed = FALSE;
+        state = HD_POST_TREATMENT_VERIFY_STATE;
+    }
+
+    return state;
+}
+
+/*********************************************************************//**
+ * @brief
+ * The handlePostTreatmentVerifyState function verifies cartridge removed,
+ * syringe removed before transition to standby mode.
+ * @details Inputs: doors' status, syringe status, occlusion sensors value
+ * @details Outputs: requested transition to standby mode
+ * @return current state (sub-mode)
+ *************************************************************************/
+static HD_POST_TREATMENT_STATE_T handlePostTreatmentVerifyState( void )
+{
+    // Assuming the cartridge door is removed already.
+    // This is used when the disable cartridge removal step software configuration is enabled
+    BOOL isCartridgeRemoved = TRUE;
+
+#ifndef _RELEASE_
+    if ( getSoftwareConfigStatus( SW_CONFIG_DISABLE_CARTRIDGE_REMOVAL_STEP ) != SW_CONFIG_ENABLE_VALUE )
+#endif
+    {
+        isCartridgeRemoved = isCartridgeUnloaded();
+    }
+
+    if ( FALSE == isCartridgeRemoved )
+    {
+        activateAlarmNoData( ALARM_ID_CARTRIDGE_REMOVAL_FAILURE );
+    }
+
+    if ( TRUE == isSyringeDetected() )
+    {
+        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();
+        requestNewOperationMode( MODE_STAN );
+    }
+
+    return HD_POST_TREATMENT_VERIFY_STATE;
+}
+
+/*********************************************************************//**
+ * @brief
+ * The execDrainReservoirs function executes the drain reservoirs operation.
+ * @details Inputs: currentDrainReservoirState
+ * @details Outputs: currentDrainReservoirState
+ * @return current state
+ *************************************************************************/
+static void execDrainReservoirs( void )
+{
+    switch ( currentDrainReservoirState )
+    {
+        case DRAIN_FIRST_RESERVOIR_START_STATE:
+            currentDrainReservoirState = handleDrainFirstReservoirStartState();
+            break;
+
+        case DRAIN_FIRST_RESERVOIR_STATE:
+            currentDrainReservoirState = handleDrainFirstReservoirState();
+            break;
+
+        case DRAIN_SWITCH_RESERVOIR_STATE:
+            currentDrainReservoirState = handleDrainSwitchReservoirState();
+            break;
+
+        case DRAIN_SECOND_RESERVOIR_STATE:
+            currentDrainReservoirState = handleDrainSecondReservoirState();
+            break;
+
+        case DRAIN_COMPLETE_STATE:
+            break;
+
+        default:
+            currentDrainReservoirState = DRAIN_FIRST_RESERVOIR_START_STATE;
+            SET_ALARM_WITH_2_U32_DATA( ALARM_ID_HD_SOFTWARE_FAULT, SW_FAULT_ID_POST_TREATMENT_RESERVOIR_MGMT_INVALID_STATE, currentDrainReservoirState );
+            break;
+    }
+}
+
+/*********************************************************************//**
+ * @brief
+ * The isReservoirDrainStarted function checks if reservoir drain has started.
+ * @details Inputs: none
+ * @details Outputs: processed drain command response
+ * @return TRUE if reservoir drain has started, otherwise FALSE
+ *************************************************************************/
+static BOOL isReservoirDrainStarted( void )
+{
+    DG_CMD_RESPONSE_T dgCmdResp;
+    BOOL result = FALSE;
+
+    if ( TRUE == getDGCommandResponse( DG_CMD_START_DRAIN, &dgCmdResp ) )
+    {
+        if ( DG_CMD_REQUEST_REJECT_REASON_NONE == dgCmdResp.rejectCode )
+        {
+            result = TRUE;
+        }
+        else
+        {
+            cmdStartDGDrain( DRAIN_RESERVOIR_TO_VOLUME_ML, TRUE, rinseConcentrateLines, TRUE );
+        }
+    }
+
+    return result;
+}
+
+/*********************************************************************//**
+ * @brief
+ * The handleDrainReservoirOneState function sends command to DG to
+ * drain first reservoir.
+ * @details Inputs: none
+ * @details Outputs: drained reservoir one
+ * @return next state
+ *************************************************************************/
+static DRAIN_STATE_T handleDrainFirstReservoirStartState( void )
+{
+    DRAIN_STATE_T state = DRAIN_FIRST_RESERVOIR_START_STATE;
+
+    if ( ( DG_MODE_GENE == getDGOpMode() ) && ( DG_GEN_IDLE_MODE_STATE_FLUSH_WATER == getDGSubMode() ) )
+    {
+        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;
+}
+
+/*********************************************************************//**
+ * @brief
+ * The handleDrainFirstReservoirState function waits for DG to finish draining
+ * first reservoir and sends command to switch reservoir.
+ * @details Inputs: none
+ * @details Outputs: drained first reservoir
+ * @return next state
+ *************************************************************************/
+static DRAIN_STATE_T handleDrainFirstReservoirState( void )
+{
+    DRAIN_STATE_T state = DRAIN_FIRST_RESERVOIR_STATE;
+
+    if ( TRUE == isReservoirDrainStarted() )
+    {
+        isDrainStarted = TRUE;
+    }
+
+    // Drain has started and DG goes to re-circ mode flush lines state means drain completed
+    if ( ( TRUE == isDrainStarted ) && ( DG_GEN_IDLE_MODE_STATE_FLUSH_WATER == getDGSubMode() ) )
+    {
+        isDrainCompleted = TRUE;
+    }
+
+    if ( TRUE == isDrainCompleted )
+    {
+        if ( ( DG_MODE_GENE == getDGOpMode() ) && ( DG_GEN_IDLE_MODE_STATE_FLUSH_WATER == getDGSubMode() ) )
+        {
+            isDrainStarted = FALSE;
+            isDrainCompleted = FALSE;
+            state = DRAIN_SWITCH_RESERVOIR_STATE;
+            cmdSetDGActiveReservoir( getDGInactiveReservoir() );
+        }
+    }
+
+    return state;
+}
+
+/*********************************************************************//**
+ * @brief
+ * The handleDrainSwitchReservoirState function waits for DG to complete
+ * reservoir switch and sends drain command.
+ * @details Inputs: none
+ * @details Outputs: proceed to next state after DG switched reservoir
+ * @return next state
+ *************************************************************************/
+static DRAIN_STATE_T handleDrainSwitchReservoirState( void )
+{
+    DRAIN_STATE_T state = DRAIN_SWITCH_RESERVOIR_STATE;
+
+    if ( TRUE == hasDGCompletedReservoirSwitch() )
+    {
+        state = DRAIN_SECOND_RESERVOIR_STATE;
+        rinseConcentrateLines = FALSE;
+        cmdStartDGDrain( DRAIN_RESERVOIR_TO_VOLUME_ML, TRUE, rinseConcentrateLines, TRUE );
+    }
+
+    return state;
+}
+
+/*********************************************************************//**
+ * @brief
+ * The handleDrainReservoirTwoState function waits for DG to finish draining
+ * second reservoir.
+ * @details Inputs: none
+ * @details Outputs: drained second reservoir
+ * @return next state
+ *************************************************************************/
+static DRAIN_STATE_T handleDrainSecondReservoirState( void )
+{
+    DRAIN_STATE_T state = DRAIN_SECOND_RESERVOIR_STATE;
+
+    if ( TRUE == isReservoirDrainStarted() )
+    {
+        isDrainStarted = TRUE;
+    }
+
+    // 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_WATER == getDGSubMode() ) )
+    {
+        isDrainCompleted = TRUE;
+    }
+
+    if ( TRUE == isDrainCompleted )
+    {
+        if ( ( DG_MODE_GENE == getDGOpMode() ) && ( DG_GEN_IDLE_MODE_STATE_FLUSH_WATER == getDGSubMode() ) )
+        {
+            isDrainStarted = FALSE;
+            isDrainCompleted = FALSE;
+            rinseConcentrateLines = FALSE;
+            state = DRAIN_COMPLETE_STATE;
+        }
+    }
+
+    return state;
+}
+
+
+/*************************************************************************
+ *                       TEST SUPPORT FUNCTIONS
+ *************************************************************************/
+
+
+/*********************************************************************//**
+ * @brief
+ * The testSetPostTreatmentModePublishIntervalOverride function sets the override of the
+ * post-treatment mode data publication interval.
+ * @details Inputs: none
+ * @details Outputs: postTreatmentModePublishInterval
+ * @param ms milliseconds between post-treatment mode broadcasts
+ * @return TRUE if override set successful, FALSE if not
+ *************************************************************************/
+BOOL testSetPostTreatmentModePublishIntervalOverride( U32 ms )
+{
+    BOOL result = FALSE;
+
+    if ( TRUE == isTestingActivated() )
+    {
+        U32 intvl = ms / TASK_GENERAL_INTERVAL;
+
+        result = TRUE;
+        postTreatmentModePublishInterval.ovData = intvl;
+        postTreatmentModePublishInterval.override = OVERRIDE_KEY;
+    }
+
+    return result;
+}
+
+/*********************************************************************//**
+ * @brief
+ * The testResetPostTreatmentModePublishIntervalOverride function resets the override of the
+ * post-treatment mode data publication interval.
+ * @details Inputs: none
+ * @details Outputs: postTreatmentModePublishInterval
+ * @return TRUE if override reset successful, FALSE if not
+ *************************************************************************/
+BOOL testResetPostTreatmentModePublishIntervalOverride( void )
+{
+    BOOL result = FALSE;
+
+    if ( TRUE == isTestingActivated() )
+    {
+        result = TRUE;
+        postTreatmentModePublishInterval.override = OVERRIDE_RESET;
+        postTreatmentModePublishInterval.ovData = postTreatmentModePublishInterval.ovInitData;
+    }
+
+    return result;
+}
+
 /**@}*/