/**************************************************************************
*
* Copyright (c) 2019-2023 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
*
* @author  (last)      Sean Nash
* @date    (last)      11-Aug-2023
*
* @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 "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 )
#define DIP_DRAIN_DG_FLOW_RATE_ML_MIN           60                              ///< DG drain inlet pump flow rate in mL/min.
#define DIP_DRAIN_DG_RUNTIME_MS                 ( 40 * MS_PER_SECOND )          ///< DG drain inlet pump run time in ms.

/// Post-Treatment drain reservoirs state machine.
typedef enum Drain_States
{
    DRAIN_RESERVOIR_SWITCH_STATE = 0,                                           ///< Switch reservoir state.
    DRAIN_RESERVOIR_START_DRAIN_STATE,                                          ///< Drain reservoir start drain state.
    DRAIN_RESERVOIR_DRAIN_STATE,                                                ///< Drain first reservoir state.
    DRAIN_RESERVOIR_COMPLETE_STATE,                                             ///< Drain complete state.
    NUM_OF_DRAIN_STATES                                                         ///< Number of drain states.
} DRAIN_STATE_T;

// ********** private data **********

static BOOL patientDisconnectionConfirmed;                                      ///< Flag indicates user confirms patient disconnection.
static bool patientDisconnectHandled;                                           ///< Flag indicates patient disconnect actions have been handled
static BOOL disposableRemovalConfirmed;                                         ///< Flag indicates user confirms disposable removal.
static BOOL rinseConcentrateLines;                                              ///< FLag indicates to rinse concentrate lines.
static BOOL isThisFirstDrain;                                                   ///< Flag to indicate whether this is the first time the reservoir is drained.
static BOOL dgCommandSent;                                                      ///< Flag indicates whether a DG command has been already sent.
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 DRAIN_STATE_T handleDrainReservoirSwitchState( void );
static DRAIN_STATE_T handleDrainReservoirStartDrainState( void );
static DRAIN_STATE_T handleDrainReservoirDrainState( void );

/*********************************************************************//**
 * @brief
 * The initPostTreatmentMode function initializes the Post-Treatment Mode module.
 * @details Inputs: none
 * @details Outputs: patientDisconnectionConfirmed, disposableRemovalConfirmed,
 * rinseConcentrateLines, isThisFirstDrain, currentPostTreatmentState,
 * currentDrainReservoirState, postTreatmentPublishTimerCounter, treatmentLogData,
 * dgCommandSent
 * @return none
 *************************************************************************/
void initPostTreatmentMode( void )
{
    patientDisconnectionConfirmed    = FALSE;
    patientDisconnectHandled         = FALSE;
    disposableRemovalConfirmed       = FALSE;
    rinseConcentrateLines            = FALSE;
    isThisFirstDrain                 = TRUE;
    dgCommandSent                    = FALSE;
    currentPostTreatmentState        = HD_POST_TREATMENT_DRAIN_RESERVOIRS_STATE;
    currentDrainReservoirState       = DRAIN_RESERVOIR_SWITCH_STATE;
    postTreatmentPublishTimerCounter = 0;

    // Reset treatment log data
    memset( &treatmentLogData, 0x0, sizeof( TREATMENT_LOG_DATA_PAYLOAD_T ) );
}

/*********************************************************************//**
 * @brief
 * The transitionToPostTreatmentMode function prepares for transition to
 * post-treatment mode.
 * @details Inputs: none
 * @details Outputs: none
 * @return initial state
 *************************************************************************/
U32 transitionToPostTreatmentMode( void )
{
    initPostTreatmentMode();

    // Stop any DG fill that may be in progress from an aborted treatment
    cmdStopDGFill();

    syringeDetectionRequired( FALSE );

    // Set user alarm recovery actions allowed in this mode
    setAlarmUserActionEnabled( ALARM_USER_ACTION_RESUME, FALSE );
    setAlarmUserActionEnabled( ALARM_USER_ACTION_RINSEBACK, FALSE );
    setAlarmUserActionEnabled( ALARM_USER_ACTION_END_TREATMENT, FALSE );

    doorClosedRequired( TRUE, TRUE );

    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();
    setCurrentSubState( (U32)currentDrainReservoirState );
    collectTreatmentLogData();

    exitBloodLeakNormalState();

    setVenousBubbleDetectionEnabled( FALSE );

    return currentPostTreatmentState;
}

/*********************************************************************//**
 * @brief
 * The execPostTreatmentMode function executes the Post-Treatment Mode state machine.
 * @details Inputs: none
 * @details Outputs: none
 * @return current state (sub-mode)
 *************************************************************************/
U32 execPostTreatmentMode( void )
{
    BOOL stop = isStopButtonPressed();
    HD_POST_TREATMENT_STATE_T priorSubState = currentPostTreatmentState;

    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_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 ( priorSubState != currentPostTreatmentState )
    {
        SEND_EVENT_WITH_2_U32_DATA( HD_EVENT_SUB_MODE_CHANGE, priorSubState, currentPostTreatmentState );
    }

    if ( ++postTreatmentPublishTimerCounter >= getU32OverrideValue( &postTreatmentModePublishInterval ) )
    {
        POST_TREATMENT_STATE_DATA_T data;

        data.currentPostTreatmentState = (U32)currentPostTreatmentState;
        data.currentPostTxDrainState   = (U32)currentDrainReservoirState;

        broadcastData( MSG_ID_HD_POST_TREATMENT_STATE_DATA, COMM_BUFFER_OUT_CAN_HD_BROADCAST, (U08*)&data, sizeof( POST_TREATMENT_STATE_DATA_T ) );
        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 requestPostTxNext requests to go to next screen after treatment record
 * review screen in post-treatment mode
 * @details Inputs: currentPostTreatmentState
 * @details Outputs: none
 * @return none
 *************************************************************************/
void requestPostTxNext( void )
{
    BOOL accepted = FALSE;
    REQUEST_REJECT_REASON_CODE_T rejReason = REQUEST_REJECT_REASON_NONE;

    if ( currentPostTreatmentState >= HD_POST_TREATMENT_DISPOSABLE_REMOVAL_STATE )
    {
        accepted = TRUE;
    }
    else
    {
        rejReason = REQUEST_REJECT_REASON_DRAIN_NOT_COMPLETE;
    }

    // Respond to user request to move to next screen
    sendPostTxNextCmdResponse( accepted, (U32)rejReason );
}

/*********************************************************************//**
 * @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_RESERVOIR_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
 * @details Outputs: given alarm action executed
 * @param action ID of alarm action to execute
 * @return none
 *************************************************************************/
void signalAlarmActionToPostTreatmentMode( ALARM_ACTION_T action )
{
    // Do nothing for now
}

/*********************************************************************//**
 * @brief
 * The checkPostTreatmentPatientDisconnect function checks UI patient
 * disconnect and handles valve release, pump homing, door requirement.
 * @details Inputs: patientDisconnectionConfirmed, patientDisconnectHandled
 * @details Outputs: patientDisconnectHandled, home pumps, open valves,
 *                   retract syringe, door closed requirement.
 * @return none
 *************************************************************************/
static void checkPostTreatmentPatientDisconnect( void )
{
    VALVE_T valve;

    if ( ( TRUE == patientDisconnectionConfirmed ) && ( FALSE == patientDisconnectHandled ) )
    {
        // home valves and pumps while front door is still closed.
        homeBloodPump();
        homeDialInPump();
        homeDialOutPump();
        for ( valve = VDI; valve < NUM_OF_VALVES; valve++ )
        {
            setValvePosition( valve, VALVE_POSITION_A_INSERT_EJECT );
        }
        if ( ( getTreatmentParameterF32( TREATMENT_PARAM_HEPARIN_BOLUS_VOLUME ) > 0.0 ) ||
             ( getTreatmentParameterF32( TREATMENT_PARAM_HEPARIN_DISPENSE_RATE ) > 0.0 ) )
        {
            retractSyringePump();
        }
        doorClosedRequired( FALSE, FALSE );
        // NOTE: patientDisconnectionConfirmed is NOT reset because
        // Patient Disconnect State needs it to move to next state
        patientDisconnectHandled = TRUE;
    }
}

/*********************************************************************//**
 * @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_RESERVOIR_COMPLETE_STATE == currentDrainReservoirState )
    {
        // Done with draining the reservoirs
        state = HD_POST_TREATMENT_PATIENT_DISCONNECTION_STATE;
        setCurrentSubState( NO_SUB_STATE );
    }
    checkPostTreatmentPatientDisconnect();

    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 )
{
    HD_POST_TREATMENT_STATE_T state = HD_POST_TREATMENT_PATIENT_DISCONNECTION_STATE;
    OPN_CLS_STATE_T frontDoor       = getSwitchStatus( FRONT_DOOR );
    OPN_CLS_STATE_T pumpTrack       = getSwitchStatus( PUMP_TRACK_SWITCH );

    checkPostTreatmentPatientDisconnect();

#ifndef _RELEASE_
    if ( SW_CONFIG_ENABLE_VALUE == getSoftwareConfigStatus( SW_CONFIG_DISABLE_UI_INTERACTION ) )
    {
        patientDisconnectionConfirmed = TRUE;
    }
#endif

    if ( TRUE == patientDisconnectionConfirmed )
    {
        patientDisconnectionConfirmed = FALSE;
        patientDisconnectHandled      = FALSE;
        state                         = HD_POST_TREATMENT_DISPOSABLE_REMOVAL_STATE;
    }

    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;

#ifndef _RELEASE_
    if ( SW_CONFIG_ENABLE_VALUE == getSoftwareConfigStatus( SW_CONFIG_DISABLE_UI_INTERACTION ) )
    {
        disposableRemovalConfirmed = TRUE;
    }
#endif

    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 )
{
    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_HD_CARTRIDGE_REMOVAL_FAILURE );
    }
    else
    {
        clearAlarmCondition( ALARM_ID_HD_CARTRIDGE_REMOVAL_FAILURE );
    }

    if ( TRUE == isSyringeDetected() )
    {
#ifndef _RELEASE_
        if ( getSoftwareConfigStatus( SW_CONFIG_DISABLE_SYRINGE_PUMP_ALARMS ) != SW_CONFIG_ENABLE_VALUE )
#endif
        {
            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 )
    {
        BOOL isSyringePumpDetected = isSyringeDetected();

        if ( FALSE == isSyringePumpDetected )
        {
            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 )
{
    DRAIN_STATE_T priorSubState = currentDrainReservoirState;
    DG_OP_MODE_T         dgMode = getDGOpMode();

    if ( DG_MODE_STAN == dgMode )
    {   // In case DG was reset/power cycled, re-start DG so we can drain the reservoirs
        cmdStartDG();
        currentDrainReservoirState = DRAIN_RESERVOIR_SWITCH_STATE;
    }
    else if ( DG_MODE_FAUL == dgMode )
    {   // In case DG faulted, skip DG drain states and move on
        currentDrainReservoirState = DRAIN_RESERVOIR_COMPLETE_STATE;
    }

    // Execute post-tx drain state machine
    switch ( currentDrainReservoirState )
    {
        case DRAIN_RESERVOIR_SWITCH_STATE:
            currentDrainReservoirState = handleDrainReservoirSwitchState();
            break;

        case DRAIN_RESERVOIR_START_DRAIN_STATE:
            currentDrainReservoirState = handleDrainReservoirStartDrainState();
            break;

        case DRAIN_RESERVOIR_DRAIN_STATE:
            currentDrainReservoirState = handleDrainReservoirDrainState();
            break;

        case DRAIN_RESERVOIR_COMPLETE_STATE:
            // Nothing to be done in this state (terminal state).
            break;

        default:
            currentDrainReservoirState = DRAIN_RESERVOIR_SWITCH_STATE;
            SET_ALARM_WITH_2_U32_DATA( ALARM_ID_HD_SOFTWARE_FAULT, SW_FAULT_ID_POST_TREATMENT_RESERVOIR_MGMT_INVALID_STATE, currentDrainReservoirState );
            break;
    }

    if ( priorSubState != currentDrainReservoirState )
    {
        setCurrentSubState( (U32)currentDrainReservoirState );
        SEND_EVENT_WITH_2_U32_DATA( HD_EVENT_SUB_STATE_CHANGE, priorSubState, currentDrainReservoirState );
    }
}

/*********************************************************************//**
 * @brief
 * The handleDrainReservoirSwitchState function waits for DG to complete
 * reservoir switch.
 * @details Inputs: isThisFirstDrain, dgCommandSent
 * @details Outputs: rinseConcentrateLines, dgCommandSent
 * @return next state
 *************************************************************************/
static DRAIN_STATE_T handleDrainReservoirSwitchState( void )
{
    DRAIN_STATE_T state  = DRAIN_RESERVOIR_SWITCH_STATE;
    DG_OP_MODE_T  dgMode = getDGOpMode();

    if ( ( DG_MODE_GENE == dgMode ) && ( DG_GEN_IDLE_MODE_STATE_FLUSH_WATER == getDGSubMode() ) )
    {
        DG_CMD_RESPONSE_T dgCmdResp;

        if ( FALSE == dgCommandSent )
        {
            DG_SWITCH_RSRVRS_CMD_T rsrvrCmd;

            // The first drain starts with setting the reservoir 2 as active and then the first reservoir is drained
            rsrvrCmd.reservoirID            = (U32)( TRUE == isThisFirstDrain ? DG_RESERVOIR_2 : DG_RESERVOIR_1 );
            rsrvrCmd.useLastTrimmerHeaterDC = FALSE;
            dgCommandSent                   = TRUE;

            cmdSetDGActiveReservoir( &rsrvrCmd );
        }
        else if ( TRUE == getDGCommandResponse( DG_CMD_SWITCH_RESERVOIR, &dgCmdResp ) )
        {
            if ( DG_CMD_REQUEST_REJECT_REASON_NONE == dgCmdResp.rejectCode )
            {
                if ( TRUE == hasDGCompletedReservoirSwitch() )
                {
                    // Check if the reservoir switch has been completed
                    // If this is the first drain, rinse the concentrate lines too
                    rinseConcentrateLines = ( TRUE == isThisFirstDrain ? TRUE : FALSE );
                    dgCommandSent         = FALSE;
                    state                 = DRAIN_RESERVOIR_START_DRAIN_STATE;
                }
            }
            else
            {
                dgCommandSent = FALSE;
            }
        }
    }

    return state;
}

/*********************************************************************//**
 * @brief
 * The handleDrainReservoirStartDrainState function sends command to DG to
 * send the drain command to DG.
 * @details Inputs: rinseConcentrateLines, dgCommandSent
 * @details Outputs: dgCommandSent
 * @return next state
 *************************************************************************/
static DRAIN_STATE_T handleDrainReservoirStartDrainState( void )
{
    DRAIN_STATE_T state = DRAIN_RESERVOIR_START_DRAIN_STATE;

    if ( ( DG_MODE_GENE == getDGOpMode() ) && ( DG_GEN_IDLE_MODE_STATE_FLUSH_WATER == getDGSubMode() ) )
    {
        DG_CMD_RESPONSE_T dgCmdResp;

        // If DG is still in the Idle mode in the flush water state and the drain command has not been sent, send the drain command
        if ( FALSE == dgCommandSent )
        {
            cmdStartDGDrain( DRAIN_RESERVOIR_TO_VOLUME_ML, TRUE, rinseConcentrateLines, TRUE );
            dgCommandSent = TRUE;
        }
        else if ( ( TRUE == getDGCommandResponse( DG_CMD_START_DRAIN, &dgCmdResp ) ) && ( dgCmdResp.rejectCode != DG_CMD_REQUEST_REJECT_REASON_NONE ) )
        {
            dgCommandSent = FALSE;
        }
    }
    else if ( DG_MODE_DRAI == getDGOpMode() )
    {
        // If the drain has started, transition to drain state
        dgCommandSent = FALSE;
        state         = DRAIN_RESERVOIR_DRAIN_STATE;
    }

    return state;
}

/*********************************************************************//**
 * @brief
 * The handleDrainReservoirDrainState function waits for DG to finish draining
 * a reservoir.
 * @details Inputs: isThisFirstDrain
 * @details Outputs: isThisFirstDrain
 * @return next state
 *************************************************************************/
static DRAIN_STATE_T handleDrainReservoirDrainState( void )
{
    DRAIN_STATE_T state = DRAIN_RESERVOIR_DRAIN_STATE;

    if ( ( DG_MODE_GENE == getDGOpMode() ) && ( DG_GEN_IDLE_MODE_STATE_FLUSH_WATER == getDGSubMode() ) )
    {
        // Decide which state should the state machine transition
        state            = ( TRUE == isThisFirstDrain ? DRAIN_RESERVOIR_SWITCH_STATE : DRAIN_RESERVOIR_COMPLETE_STATE );
        // At least the first reservoir has been drained so far so set the variable to false.
        isThisFirstDrain = FALSE;
    }

    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;
}

/**@}*/