From 15e99c9c4a1983db0e3821a4c83e76cc40a9cb10 Mon Sep 17 00:00:00 2001
From: Chris Oloff <chris@uber5.com>
Date: Thu, 20 Nov 2025 08:46:12 +0200
Subject: [PATCH] sample2 improvements

---
 sample2.py | 34 ++++++++++++++++++++++++++++------
 1 file changed, 28 insertions(+), 6 deletions(-)

diff --git a/sample2.py b/sample2.py
index c58b6c2..d38a34c 100644
--- a/sample2.py
+++ b/sample2.py
@@ -18,9 +18,17 @@ class LoadForecast(BaseModel):
     async def update(self, now: datetime) -> None:
         pass
 
+class Output(BaseModel):
+    target_storage: float
+    target_geyser_temperature: float
+    storage: list[float]
+    grid_sell: list[float]
+    grid_buy: list[float]
+
 class SystemState(BaseModel):
     load_forecast: LoadForecast = LoadForecast()
     now: datetime = datetime.now()
+    output: Output | None = None
 
 
 state: SystemState = panel.rx(lambda s: SystemState.model_validate(s) if s else SystemState())(w42_state)
@@ -39,12 +47,14 @@ def get_label(at: datetime, slot_index: int) -> str:
 def load_fc_chart(state: SystemState) -> dict:
     slots = state.load_forecast.slots
     at = state.load_forecast.at
+    storage = state.output.storage if state.output else [0] * len(slots)
+    battery = [storage[ix] - storage[ix-1] if ix > 0 else 0 for ix in range(len(slots))] if state.output else [0] * len(slots)
     return {
         'title': {
             'text': '24hr Forecast'
         },
         'legend': {
-            'data': ['Load Forecast']
+            'data': ['Load Forecast', 'Battery Storage', 'Battery in/out']
         },
         'tooltip': {
             'trigger': 'axis'
@@ -56,11 +66,23 @@ def load_fc_chart(state: SystemState) -> dict:
         'yAxis': {
             'type': 'value'
         },
-        'series': [{
-            'name': 'Load Forecast',
-            'data': slots,
-            'type': 'line'
-        }]
+        'series': [
+            {
+                'name': 'Load Forecast',
+                'data': slots,
+                'type': 'line'
+            }, {
+                'name': 'Battery Storage',
+                'data': storage,
+                'type': 'line',
+                'color': 'orange'
+            }, {
+                'name': 'Battery in/out',
+                'data': battery,
+                'type': 'line',
+                'color': 'lightblue'
+            }
+        ]
     }
 
 load_fc_chart_rx = panel.rx(load_fc_chart)(state)