From ebaa84617feb8b07e54c6bac012c285e4c077d01 Mon Sep 17 00:00:00 2001
From: Chris Nussbaum <chrisnussbaum@gmail.com>
Date: Tue, 17 Aug 2021 14:16:02 -0500
Subject: [PATCH] Total daily energy methods (#2163)

Co-authored-by: Chris Nussbaum <chris.nussbaum@protolabs.com>
---
 .../components/total_daily_energy/sensor.py   | 11 ++++++++++
 .../total_daily_energy/total_daily_energy.cpp | 21 ++++++++++++++++++-
 .../total_daily_energy/total_daily_energy.h   |  9 ++++++++
 3 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/esphome/components/total_daily_energy/sensor.py b/esphome/components/total_daily_energy/sensor.py
index ec38daaf57..f6e92149d6 100644
--- a/esphome/components/total_daily_energy/sensor.py
+++ b/esphome/components/total_daily_energy/sensor.py
@@ -7,6 +7,7 @@ from esphome.const import (
     DEVICE_CLASS_ENERGY,
     LAST_RESET_TYPE_AUTO,
     STATE_CLASS_MEASUREMENT,
+    CONF_METHOD,
 )
 
 DEPENDENCIES = ["time"]
@@ -14,6 +15,12 @@ DEPENDENCIES = ["time"]
 CONF_POWER_ID = "power_id"
 CONF_MIN_SAVE_INTERVAL = "min_save_interval"
 total_daily_energy_ns = cg.esphome_ns.namespace("total_daily_energy")
+TotalDailyEnergyMethod = total_daily_energy_ns.enum("TotalDailyEnergyMethod")
+TOTAL_DAILY_ENERGY_METHODS = {
+    "trapezoid": TotalDailyEnergyMethod.TOTAL_DAILY_ENERGY_METHOD_TRAPEZOID,
+    "left": TotalDailyEnergyMethod.TOTAL_DAILY_ENERGY_METHOD_LEFT,
+    "right": TotalDailyEnergyMethod.TOTAL_DAILY_ENERGY_METHOD_RIGHT,
+}
 TotalDailyEnergy = total_daily_energy_ns.class_(
     "TotalDailyEnergy", sensor.Sensor, cg.Component
 )
@@ -33,6 +40,9 @@ CONFIG_SCHEMA = (
             cv.Optional(
                 CONF_MIN_SAVE_INTERVAL, default="0s"
             ): cv.positive_time_period_milliseconds,
+            cv.Optional(CONF_METHOD, default="right"): cv.enum(
+                TOTAL_DAILY_ENERGY_METHODS, lower=True
+            ),
         }
     )
     .extend(cv.COMPONENT_SCHEMA)
@@ -50,3 +60,4 @@ async def to_code(config):
     time_ = await cg.get_variable(config[CONF_TIME_ID])
     cg.add(var.set_time(time_))
     cg.add(var.set_min_save_interval(config[CONF_MIN_SAVE_INTERVAL]))
+    cg.add(var.set_method(config[CONF_METHOD]))
diff --git a/esphome/components/total_daily_energy/total_daily_energy.cpp b/esphome/components/total_daily_energy/total_daily_energy.cpp
index 1e60442ae7..83333acab7 100644
--- a/esphome/components/total_daily_energy/total_daily_energy.cpp
+++ b/esphome/components/total_daily_energy/total_daily_energy.cpp
@@ -20,7 +20,9 @@ void TotalDailyEnergy::setup() {
 
   this->parent_->add_on_state_callback([this](float state) { this->process_new_state_(state); });
 }
+
 void TotalDailyEnergy::dump_config() { LOG_SENSOR("", "Total Daily Energy", this); }
+
 void TotalDailyEnergy::loop() {
   auto t = this->time_->now();
   if (!t.is_valid())
@@ -37,6 +39,7 @@ void TotalDailyEnergy::loop() {
     this->publish_state_and_save(0);
   }
 }
+
 void TotalDailyEnergy::publish_state_and_save(float state) {
   this->total_energy_ = state;
   this->publish_state(state);
@@ -47,13 +50,29 @@ void TotalDailyEnergy::publish_state_and_save(float state) {
   this->last_save_ = now;
   this->pref_.save(&state);
 }
+
 void TotalDailyEnergy::process_new_state_(float state) {
   if (isnan(state))
     return;
   const uint32_t now = millis();
+  const float old_state = this->last_power_state_;
+  const float new_state = state;
   float delta_hours = (now - this->last_update_) / 1000.0f / 60.0f / 60.0f;
+  float delta_energy = 0.0f;
+  switch (this->method_) {
+    case TOTAL_DAILY_ENERGY_METHOD_TRAPEZOID:
+      delta_energy = delta_hours * (old_state + new_state) / 2.0;
+      break;
+    case TOTAL_DAILY_ENERGY_METHOD_LEFT:
+      delta_energy = delta_hours * old_state;
+      break;
+    case TOTAL_DAILY_ENERGY_METHOD_RIGHT:
+      delta_energy = delta_hours * new_state;
+      break;
+  }
+  this->last_power_state_ = new_state;
   this->last_update_ = now;
-  this->publish_state_and_save(this->total_energy_ + state * delta_hours);
+  this->publish_state_and_save(this->total_energy_ + delta_energy);
 }
 
 }  // namespace total_daily_energy
diff --git a/esphome/components/total_daily_energy/total_daily_energy.h b/esphome/components/total_daily_energy/total_daily_energy.h
index 123446c534..fd71b8decc 100644
--- a/esphome/components/total_daily_energy/total_daily_energy.h
+++ b/esphome/components/total_daily_energy/total_daily_energy.h
@@ -8,11 +8,18 @@
 namespace esphome {
 namespace total_daily_energy {
 
+enum TotalDailyEnergyMethod {
+  TOTAL_DAILY_ENERGY_METHOD_TRAPEZOID = 0,
+  TOTAL_DAILY_ENERGY_METHOD_LEFT,
+  TOTAL_DAILY_ENERGY_METHOD_RIGHT,
+};
+
 class TotalDailyEnergy : public sensor::Sensor, public Component {
  public:
   void set_min_save_interval(uint32_t min_interval) { this->min_save_interval_ = min_interval; }
   void set_time(time::RealTimeClock *time) { time_ = time; }
   void set_parent(Sensor *parent) { parent_ = parent; }
+  void set_method(TotalDailyEnergyMethod method) { method_ = method; }
   void setup() override;
   void dump_config() override;
   float get_setup_priority() const override { return setup_priority::DATA; }
@@ -29,11 +36,13 @@ class TotalDailyEnergy : public sensor::Sensor, public Component {
   ESPPreferenceObject pref_;
   time::RealTimeClock *time_;
   Sensor *parent_;
+  TotalDailyEnergyMethod method_;
   uint16_t last_day_of_year_{};
   uint32_t last_update_{0};
   uint32_t last_save_{0};
   uint32_t min_save_interval_{0};
   float total_energy_{0.0f};
+  float last_power_state_{0.0f};
 };
 
 }  // namespace total_daily_energy