From 084a622fe06b40f8818cb24214834fec6928d6d4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?L=2E=20L=C3=B3pez?= <ldotlopez@gmail.com>
Date: Wed, 21 Jun 2023 14:32:19 +0200
Subject: [PATCH] Rewrite statistics calculation

---
 custom_components/ideenergy/manifest.json |  4 +-
 custom_components/ideenergy/sensor.py     | 53 +++++++++++++++--------
 2 files changed, 36 insertions(+), 21 deletions(-)

diff --git a/custom_components/ideenergy/manifest.json b/custom_components/ideenergy/manifest.json
index 3fd8bc7..e38d52a 100644
--- a/custom_components/ideenergy/manifest.json
+++ b/custom_components/ideenergy/manifest.json
@@ -10,7 +10,7 @@
   "issue_tracker": "https://github.com/ldotlopez/ha-ideenergy/issues",
   "requirements": [
       "ideenergy==1.0.0",
-      "homeassistant-historical-sensor==1.0.0"
+      "homeassistant-historical-sensor==1.0.1"
     ],
   "version": "2.0.0"
-}
+}
\ No newline at end of file
diff --git a/custom_components/ideenergy/sensor.py b/custom_components/ideenergy/sensor.py
index aeda374..7e4f84d 100644
--- a/custom_components/ideenergy/sensor.py
+++ b/custom_components/ideenergy/sensor.py
@@ -27,8 +27,10 @@
 # https://github.com/home-assistant/core/blob/dev/homeassistant/components/sensor/__init__.py
 
 
+import itertools
 import logging
-from datetime import timedelta
+import statistics
+from datetime import datetime, timedelta
 from typing import Any, Callable, Dict, List, Optional
 
 from homeassistant.components.recorder.models import StatisticData, StatisticMetaData
@@ -97,31 +99,44 @@ def statatistic_id(self):
         return self.entity_id
 
     def get_statatistic_metadata(self) -> StatisticMetaData:
-        meta = super().get_statatistic_metadata()
-        meta["has_sum"] = True
+        meta = super().get_statatistic_metadata() | {"has_sum": True, "has_mean": True}
         return meta
 
     async def async_calculate_statistic_data(
         self, hist_states: List[HistoricalState], *, latest: Optional[dict]
     ) -> List[StatisticData]:
-        if latest:
-            cutoff = dtutil.utc_from_timestamp(latest.get("end") or 0)
-            accumulated = latest.get("sum") or 0
-            hist_states = [x for x in hist_states if x.dt >= cutoff]
-
-        else:
-            accumulated = 0
-
-        def calculate_statistics_from_accumulated(accumulated):
-            for x in hist_states:
-                accumulated = accumulated + x.state
-                # Shift start by 1 hour since DatedState.when represents the
-                # end of the period
-                yield StatisticData(
-                    start=x.dt - timedelta(hours=1), state=x.state, sum=accumulated
+        def hour_block_for_hist_state(hist_state: HistoricalState) -> datetime:
+            # XX:00:00 states belongs to previous hour block
+            if hist_state.dt.minute == 0 and hist_state.dt.second == 0:
+                dt = hist_state.dt - timedelta(hours=1)
+                return dt.replace(minute=0, second=0, microsecond=0)
+
+            else:
+                return hist_state.dt.replace(minute=0, second=0, microsecond=0)
+
+        total_accumulated = latest["sum"] if latest else 0
+        ret = []
+
+        # Group historical states by hour block
+        for dt, collection_it in itertools.groupby(
+            hist_states, key=hour_block_for_hist_state
+        ):
+            collection = list(collection_it)
+
+            hour_mean = statistics.mean([x.state for x in collection])
+            hour_accumulated = sum([x.state for x in collection])
+            total_accumulated = total_accumulated + hour_accumulated
+
+            ret.append(
+                StatisticData(
+                    start=dt,
+                    state=hour_accumulated,
+                    mean=hour_mean,
+                    sum=total_accumulated,
                 )
+            )
 
-        return list(calculate_statistics_from_accumulated(accumulated))
+        return ret
 
 
 class AccumulatedConsumption(RestoreEntity, IDeEntity, SensorEntity):