From d31040f5d8023a81a1011fe2b55e63992724d917 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Adri=C3=A1n=20Panella?= <ianchi74@outlook.com>
Date: Sun, 4 Jul 2021 18:09:09 -0500
Subject: [PATCH] hlw8012: fix constants for BL0937 (#1973)

Co-authored-by: Jesse Hills <3060199+jesserockz@users.noreply.github.com>
---
 esphome/components/hlw8012/hlw8012.cpp | 34 +++++++++++++++++---------
 esphome/components/hlw8012/hlw8012.h   | 12 +++++++++
 esphome/components/hlw8012/sensor.py   | 11 +++++++++
 tests/test1.yaml                       |  1 +
 4 files changed, 46 insertions(+), 12 deletions(-)

diff --git a/esphome/components/hlw8012/hlw8012.cpp b/esphome/components/hlw8012/hlw8012.cpp
index 356dbd0bf4..79c25a45b0 100644
--- a/esphome/components/hlw8012/hlw8012.cpp
+++ b/esphome/components/hlw8012/hlw8012.cpp
@@ -6,15 +6,32 @@ namespace hlw8012 {
 
 static const char *const TAG = "hlw8012";
 
+// valid for HLW8012 and CSE7759
 static const uint32_t HLW8012_CLOCK_FREQUENCY = 3579000;
-static const float HLW8012_REFERENCE_VOLTAGE = 2.43f;
 
 void HLW8012Component::setup() {
+  float reference_voltage = 0;
   ESP_LOGCONFIG(TAG, "Setting up HLW8012...");
   this->sel_pin_->setup();
   this->sel_pin_->digital_write(this->current_mode_);
   this->cf_store_.pulse_counter_setup(this->cf_pin_);
   this->cf1_store_.pulse_counter_setup(this->cf1_pin_);
+
+  // Initialize multipliers
+  if (this->sensor_model_ == HLW8012_SENSOR_MODEL_BL0937) {
+    reference_voltage = 1.218f;
+    this->power_multiplier_ =
+        reference_voltage * reference_voltage * this->voltage_divider_ / this->current_resistor_ / 1721506.0f;
+    this->current_multiplier_ = reference_voltage / this->current_resistor_ / 94638.0f;
+    this->voltage_multiplier_ = reference_voltage * this->voltage_divider_ / 15397.0f;
+  } else {
+    // HLW8012 and CSE7759 have same reference specs
+    reference_voltage = 2.43f;
+    this->power_multiplier_ = reference_voltage * reference_voltage * this->voltage_divider_ / this->current_resistor_ *
+                              64.0f / 24.0f / HLW8012_CLOCK_FREQUENCY;
+    this->current_multiplier_ = reference_voltage / this->current_resistor_ * 512.0f / 24.0f / HLW8012_CLOCK_FREQUENCY;
+    this->voltage_multiplier_ = reference_voltage * this->voltage_divider_ * 256.0f / HLW8012_CLOCK_FREQUENCY;
+  }
 }
 void HLW8012Component::dump_config() {
   ESP_LOGCONFIG(TAG, "HLW8012:");
@@ -49,25 +66,18 @@ void HLW8012Component::update() {
     return;
   }
 
-  const float v_ref_squared = HLW8012_REFERENCE_VOLTAGE * HLW8012_REFERENCE_VOLTAGE;
-  const float power_multiplier_micros =
-      64000000.0f * v_ref_squared * this->voltage_divider_ / this->current_resistor_ / 24.0f / HLW8012_CLOCK_FREQUENCY;
-  float power = cf_hz * power_multiplier_micros / 1000000.0f;
+  float power = cf_hz * this->power_multiplier_;
 
   if (this->change_mode_at_ != 0) {
     // Only read cf1 after one cycle. Apparently it's quite unstable after being changed.
     if (this->current_mode_) {
-      const float current_multiplier_micros =
-          512000000.0f * HLW8012_REFERENCE_VOLTAGE / this->current_resistor_ / 24.0f / HLW8012_CLOCK_FREQUENCY;
-      float current = cf1_hz * current_multiplier_micros / 1000000.0f;
+      float current = cf1_hz * this->current_multiplier_;
       ESP_LOGD(TAG, "Got power=%.1fW, current=%.1fA", power, current);
       if (this->current_sensor_ != nullptr) {
         this->current_sensor_->publish_state(current);
       }
     } else {
-      const float voltage_multiplier_micros =
-          256000000.0f * HLW8012_REFERENCE_VOLTAGE * this->voltage_divider_ / HLW8012_CLOCK_FREQUENCY;
-      float voltage = cf1_hz * voltage_multiplier_micros / 1000000.0f;
+      float voltage = cf1_hz * this->voltage_multiplier_;
       ESP_LOGD(TAG, "Got power=%.1fW, voltage=%.1fV", power, voltage);
       if (this->voltage_sensor_ != nullptr) {
         this->voltage_sensor_->publish_state(voltage);
@@ -81,7 +91,7 @@ void HLW8012Component::update() {
 
   if (this->energy_sensor_ != nullptr) {
     cf_total_pulses_ += raw_cf;
-    float energy = cf_total_pulses_ * power_multiplier_micros / 3600 / 1000000.0f;
+    float energy = cf_total_pulses_ * this->power_multiplier_ / 3600;
     this->energy_sensor_->publish_state(energy);
   }
 
diff --git a/esphome/components/hlw8012/hlw8012.h b/esphome/components/hlw8012/hlw8012.h
index af1f2e9a8c..52fb03c020 100644
--- a/esphome/components/hlw8012/hlw8012.h
+++ b/esphome/components/hlw8012/hlw8012.h
@@ -10,6 +10,12 @@ namespace hlw8012 {
 
 enum HLW8012InitialMode { HLW8012_INITIAL_MODE_CURRENT = 0, HLW8012_INITIAL_MODE_VOLTAGE };
 
+enum HLW8012SensorModels {
+  HLW8012_SENSOR_MODEL_HLW8012 = 0,
+  HLW8012_SENSOR_MODEL_CSE7759,
+  HLW8012_SENSOR_MODEL_BL0937
+};
+
 class HLW8012Component : public PollingComponent {
  public:
   void setup() override;
@@ -20,6 +26,7 @@ class HLW8012Component : public PollingComponent {
   void set_initial_mode(HLW8012InitialMode initial_mode) {
     current_mode_ = initial_mode == HLW8012_INITIAL_MODE_CURRENT;
   }
+  void set_sensor_model(HLW8012SensorModels sensor_model) { sensor_model_ = sensor_model; }
   void set_change_mode_every(uint32_t change_mode_every) { change_mode_every_ = change_mode_every; }
   void set_current_resistor(float current_resistor) { current_resistor_ = current_resistor; }
   void set_voltage_divider(float voltage_divider) { voltage_divider_ = voltage_divider; }
@@ -38,6 +45,7 @@ class HLW8012Component : public PollingComponent {
   uint32_t change_mode_every_{8};
   float current_resistor_{0.001};
   float voltage_divider_{2351};
+  HLW8012SensorModels sensor_model_{HLW8012_SENSOR_MODEL_HLW8012};
   uint64_t cf_total_pulses_{0};
   GPIOPin *sel_pin_;
   GPIOPin *cf_pin_;
@@ -48,6 +56,10 @@ class HLW8012Component : public PollingComponent {
   sensor::Sensor *current_sensor_{nullptr};
   sensor::Sensor *power_sensor_{nullptr};
   sensor::Sensor *energy_sensor_{nullptr};
+
+  float voltage_multiplier_{0.0f};
+  float current_multiplier_{0.0f};
+  float power_multiplier_{0.0f};
 };
 
 }  // namespace hlw8012
diff --git a/esphome/components/hlw8012/sensor.py b/esphome/components/hlw8012/sensor.py
index 6454a9fcc9..e24e995eba 100644
--- a/esphome/components/hlw8012/sensor.py
+++ b/esphome/components/hlw8012/sensor.py
@@ -11,6 +11,7 @@ from esphome.const import (
     CONF_POWER,
     CONF_ENERGY,
     CONF_SEL_PIN,
+    CONF_MODEL,
     CONF_VOLTAGE,
     CONF_VOLTAGE_DIVIDER,
     DEVICE_CLASS_CURRENT,
@@ -31,11 +32,19 @@ AUTO_LOAD = ["pulse_counter"]
 hlw8012_ns = cg.esphome_ns.namespace("hlw8012")
 HLW8012Component = hlw8012_ns.class_("HLW8012Component", cg.PollingComponent)
 HLW8012InitialMode = hlw8012_ns.enum("HLW8012InitialMode")
+HLW8012SensorModels = hlw8012_ns.enum("HLW8012SensorModels")
+
 INITIAL_MODES = {
     CONF_CURRENT: HLW8012InitialMode.HLW8012_INITIAL_MODE_CURRENT,
     CONF_VOLTAGE: HLW8012InitialMode.HLW8012_INITIAL_MODE_VOLTAGE,
 }
 
+MODELS = {
+    "HLW8012": HLW8012SensorModels.HLW8012_SENSOR_MODEL_HLW8012,
+    "CSE7759": HLW8012SensorModels.HLW8012_SENSOR_MODEL_CSE7759,
+    "BL0937": HLW8012SensorModels.HLW8012_SENSOR_MODEL_BL0937,
+}
+
 CONF_CF1_PIN = "cf1_pin"
 CONF_CF_PIN = "cf_pin"
 CONFIG_SCHEMA = cv.Schema(
@@ -62,6 +71,7 @@ CONFIG_SCHEMA = cv.Schema(
         ),
         cv.Optional(CONF_CURRENT_RESISTOR, default=0.001): cv.resistance,
         cv.Optional(CONF_VOLTAGE_DIVIDER, default=2351): cv.positive_float,
+        cv.Optional(CONF_MODEL, default="HLW8012"): cv.enum(MODELS, upper=True),
         cv.Optional(CONF_CHANGE_MODE_EVERY, default=8): cv.All(
             cv.uint32_t, cv.Range(min=1)
         ),
@@ -99,3 +109,4 @@ async def to_code(config):
     cg.add(var.set_voltage_divider(config[CONF_VOLTAGE_DIVIDER]))
     cg.add(var.set_change_mode_every(config[CONF_CHANGE_MODE_EVERY]))
     cg.add(var.set_initial_mode(INITIAL_MODES[config[CONF_INITIAL_MODE]]))
+    cg.add(var.set_sensor_model(config[CONF_MODEL]))
diff --git a/tests/test1.yaml b/tests/test1.yaml
index 08e1d63534..78dc40cf8e 100644
--- a/tests/test1.yaml
+++ b/tests/test1.yaml
@@ -514,6 +514,7 @@ sensor:
     voltage_divider: 2351
     change_mode_every: 16
     initial_mode: VOLTAGE
+    model: hlw8012
   - platform: total_daily_energy
     power_id: hlw8012_power
     name: 'HLW8012 Total Daily Energy'