diff --git a/include/mqtt.h b/include/mqtt.h
index e58de6e6..cd2e7473 100644
--- a/include/mqtt.h
+++ b/include/mqtt.h
@@ -7,6 +7,7 @@
 
 #define EEPROM_CHK 1
 #define EEPROM_STATE 0
+#define EEPROM_SG 2
 
 #define MQTT_attr "espaltherma/ATTR"
 #define MQTT_lwt "espaltherma/LWT"
@@ -53,22 +54,66 @@ void sendValues()
 #endif
 }
 
-void saveEEPROM(uint8_t state){
-    EEPROM.write(EEPROM_STATE,state);
+void digitalWriteSgPins(uint8_t state)
+{
+  if (state == 0)
+  {
+    // Set SG 0 mode => SG1 = INACTIVE, SG2 = INACTIVE
+    digitalWrite(PIN_SG1, SG_RELAY_INACTIVE_STATE);
+    digitalWrite(PIN_SG2, SG_RELAY_INACTIVE_STATE);
+  }
+  else if (state == 1)
+  {
+    // Set SG 1 mode => SG1 = INACTIVE, SG2 = ACTIVE
+    digitalWrite(PIN_SG1, SG_RELAY_INACTIVE_STATE);
+    digitalWrite(PIN_SG2, SG_RELAY_ACTIVE_STATE);
+  }
+  else if (state == 2)
+  {
+    // Set SG 2 mode => SG1 = ACTIVE, SG2 = INACTIVE
+    digitalWrite(PIN_SG1, SG_RELAY_ACTIVE_STATE);
+    digitalWrite(PIN_SG2, SG_RELAY_INACTIVE_STATE);
+  }
+  else if (state == 3)
+  {
+    // Set SG 3 mode => SG1 = ACTIVE, SG2 = ACTIVE
+    digitalWrite(PIN_SG1, SG_RELAY_ACTIVE_STATE);
+    digitalWrite(PIN_SG2, SG_RELAY_ACTIVE_STATE);
+  }
+}
+
+void saveSgState(uint8_t state)
+{
+    EEPROM.write(EEPROM_SG, state);
     EEPROM.commit();
 }
 
-void readEEPROM(){
-  if ('R' == EEPROM.read(EEPROM_CHK)){
-    digitalWrite(PIN_THERM,EEPROM.read(EEPROM_STATE));
-    mqttSerial.printf("Restoring previous state: %s",(EEPROM.read(EEPROM_STATE) == PIN_THERM_ACTIVE_STATE)? "On":"Off" );
+void saveThermState(uint8_t state)
+{
+    EEPROM.write(EEPROM_STATE, state);
+    EEPROM.commit();
+}
+
+void restoreEEPROM()
+{
+  if ('R' == EEPROM.read(EEPROM_CHK)) {
+    mqttSerial.printf("Restoring previous state: %s", (EEPROM.read(EEPROM_STATE) == PIN_THERM_ACTIVE_STATE) ? "On" : "Off");
+    digitalWrite(PIN_THERM, EEPROM.read(EEPROM_STATE));
+    #ifdef PIN_SG1
+    digitalWriteSgPins(EEPROM.read(EEPROM_SG));
+    #endif
   }
-  else{
-    mqttSerial.printf("EEPROM not initialized (%d). Initializing...",EEPROM.read(EEPROM_CHK));
-    EEPROM.write(EEPROM_CHK,'R');
-    EEPROM.write(EEPROM_STATE,!PIN_THERM_ACTIVE_STATE);
+  else
+  {
+    mqttSerial.printf("EEPROM not initialized (%d). Initializing...", EEPROM.read(EEPROM_CHK));
+    EEPROM.write(EEPROM_CHK, 'R');
     EEPROM.commit();
-    digitalWrite(PIN_THERM,!PIN_THERM_ACTIVE_STATE);
+    saveThermState(!PIN_THERM_ACTIVE_STATE);
+    digitalWrite(PIN_THERM, !PIN_THERM_ACTIVE_STATE);
+    #ifdef PIN_SG1
+    saveSgState(0);
+    digitalWriteSgPins(0);
+    #endif
   }
 }
 
@@ -83,28 +128,38 @@ void reconnectMqtt()
     if (client.connect("ESPAltherma-dev", MQTT_USERNAME, MQTT_PASSWORD, MQTT_lwt, 0, true, "Offline"))
     {
       Serial.println("connected!");
+      #ifdef MQTT_HA_DISCOVERY
       client.publish("homeassistant/sensor/espAltherma/config", "{\"name\":\"AlthermaSensors\",\"stat_t\":\"~/LWT\",\"avty_t\":\"~/LWT\",\"pl_avail\":\"Online\",\"pl_not_avail\":\"Offline\",\"uniq_id\":\"espaltherma\",\"device\":{\"identifiers\":[\"ESPAltherma\"]}, \"~\":\"espaltherma\",\"json_attr_t\":\"~/ATTR\"}", true);
-      client.publish(MQTT_lwt, "Online", true);
       client.publish("homeassistant/switch/espAltherma/config", "{\"name\":\"Altherma\",\"cmd_t\":\"~/POWER\",\"stat_t\":\"~/STATE\",\"pl_off\":\"OFF\",\"pl_on\":\"ON\",\"~\":\"espaltherma\"}", true);
+      #endif
+      client.publish(MQTT_lwt, "Online", true);
 
       // Subscribe
       client.subscribe("espaltherma/POWER");
 #ifdef PIN_SG1
       // Smart Grid
+      #ifdef MQTT_HA_DISCOVERY
       client.publish("homeassistant/select/espAltherma/sg/config", "{\"availability\":[{\"topic\":\"espaltherma/LWT\",\"payload_available\":\"Online\",\"payload_not_available\":\"Offline\"}],\"availability_mode\":\"all\",\"unique_id\":\"espaltherma_sg\",\"device\":{\"identifiers\":[\"ESPAltherma\"],\"manufacturer\":\"ESPAltherma\",\"model\":\"M5StickC PLUS ESP32-PICO\",\"name\":\"ESPAltherma\"},\"icon\":\"mdi:solar-power\",\"name\":\"EspAltherma Smart Grid\",\"command_topic\":\"espaltherma/sg/set\",\"command_template\":\"{% if value == 'Free Running' %} 0 {% elif value == 'Forced Off' %} 1 {% elif value == 'Recommended On' %} 2 {% elif value == 'Forced On' %} 3 {% else %} 0 {% endif %}\",\"options\":[\"Free Running\",\"Forced Off\",\"Recommended On\",\"Forced On\"],\"state_topic\":\"espaltherma/sg/state\",\"value_template\":\"{% set mapper = { '0':'Free Running', '1':'Forced Off', '2':'Recommended On', '3':'Forced On' } %} {% set word = mapper[value] %} {{ word }}\"}", true);
+      #endif
       client.subscribe("espaltherma/sg/set");
-      client.publish("espaltherma/sg/state", "0");
+      char state[1];
+      sprintf(state, "%d", EEPROM.read(EEPROM_SG));
+      client.publish("espaltherma/sg/state", state);
 #endif
 
 #ifdef SAFETY_RELAY_PIN
       // Safety relay
+      #ifdef MQTT_HA_DISCOVERY
       client.publish("homeassistant/switch/espAltherma/safety/config", "{\"name\":\"Altherma Safety\",\"cmd_t\":\"~/SAFETY\",\"stat_t\":\"~/SAFETY_STATE\",\"pl_off\":\"0\",\"pl_on\":\"1\",\"~\":\"espaltherma\"}", true);
+      #endif
       client.subscribe("espaltherma/SAFETY");
 #endif
 
 #ifndef PIN_SG1
+      #ifdef MQTT_HA_DISCOVERY
       // Publish empty retained message so discovered entities are removed from HA
       client.publish("homeassistant/select/espAltherma/sg/config", "", true);
+      #endif
 #endif
     }
     else
@@ -133,18 +188,18 @@ void callbackTherm(byte *payload, unsigned int length)
   if (payload[1] == 'F')
   { //turn off
     digitalWrite(PIN_THERM, !PIN_THERM_ACTIVE_STATE);
-    saveEEPROM(!PIN_THERM_ACTIVE_STATE);
+    saveThermState(!PIN_THERM_ACTIVE_STATE);
     client.publish("espaltherma/STATE", "OFF", true);
     mqttSerial.println("Turned OFF");
   }
   else if (payload[1] == 'N')
   { //turn on
     digitalWrite(PIN_THERM, PIN_THERM_ACTIVE_STATE);
-    saveEEPROM(PIN_THERM_ACTIVE_STATE);
+    saveThermState(PIN_THERM_ACTIVE_STATE);
     client.publish("espaltherma/STATE", "ON", true);
     mqttSerial.println("Turned ON");
   }
-  else if (payload[0] == 'R')//R(eset/eboot)
+  else if (payload[0] == 'R') //R(eset/eboot)
   {
     mqttSerial.println("Rebooting");
     delay(100);
@@ -165,32 +220,32 @@ void callbackSg(byte *payload, unsigned int length)
   if (payload[0] == '0')
   {
     // Set SG 0 mode => SG1 = INACTIVE, SG2 = INACTIVE
-    digitalWrite(PIN_SG1, SG_RELAY_INACTIVE_STATE);
-    digitalWrite(PIN_SG2, SG_RELAY_INACTIVE_STATE);
+    digitalWriteSgPins(0);
+    saveSgState(0);
     client.publish("espaltherma/sg/state", "0");
     Serial.println("Set SG mode to 0 - Normal operation");
   }
   else if (payload[0] == '1')
   {
     // Set SG 1 mode => SG1 = INACTIVE, SG2 = ACTIVE
-    digitalWrite(PIN_SG1, SG_RELAY_INACTIVE_STATE);
-    digitalWrite(PIN_SG2, SG_RELAY_ACTIVE_STATE);
+    digitalWriteSgPins(1);
+    saveSgState(1);
     client.publish("espaltherma/sg/state", "1");
     Serial.println("Set SG mode to 1 - Forced OFF");
   }
   else if (payload[0] == '2')
   {
     // Set SG 2 mode => SG1 = ACTIVE, SG2 = INACTIVE
-    digitalWrite(PIN_SG1, SG_RELAY_ACTIVE_STATE);
-    digitalWrite(PIN_SG2, SG_RELAY_INACTIVE_STATE);
+    digitalWriteSgPins(2);
+    saveSgState(2);
     client.publish("espaltherma/sg/state", "2");
     Serial.println("Set SG mode to 2 - Recommended ON");
   }
   else if (payload[0] == '3')
   {
     // Set SG 3 mode => SG1 = ACTIVE, SG2 = ACTIVE
-    digitalWrite(PIN_SG1, SG_RELAY_ACTIVE_STATE);
-    digitalWrite(PIN_SG2, SG_RELAY_ACTIVE_STATE);
+    digitalWriteSgPins(3);
+    saveSgState(3);
     client.publish("espaltherma/sg/state", "3");
     Serial.println("Set SG mode to 3 - Forced ON");
   }
diff --git a/src/main.cpp b/src/main.cpp
index 9a4612ec..5dc88eb6 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -328,7 +328,7 @@ void setup()
 #endif
 
   EEPROM.begin(10);
-  readEEPROM();//Restore previous state
+  restoreEEPROM();//Restore previous state
   mqttSerial.print("Setting up wifi...");
   setup_wifi();
   ArduinoOTA.setHostname("ESPAltherma");
diff --git a/src/setup.h b/src/setup.h
index 5ae80839..b5e83465 100644
--- a/src/setup.h
+++ b/src/setup.h
@@ -15,6 +15,7 @@
 #define MQTT_PASSWORD ""//leave empty if not set (bad!)
 #define MQTT_PORT 1883
 //#define MQTT_ENCRYPTED // uncomment if MQTT connection is encrypted via TLS
+#define MQTT_HA_DISCOVERY // comment to disable HomeAssistant discovery messages
 
 #define FREQUENCY 30000 //query values every 30 sec