diff --git a/models/operation/model_ref.py b/models/operation/model_ref.py
index 9e01b41..2d860e0 100644
--- a/models/operation/model_ref.py
+++ b/models/operation/model_ref.py
@@ -111,47 +111,26 @@ def calc_battery_energy(self, grid_demand: np.array, pv_surplus: np.array):
 
                 if pv_surplus[i] > 0:
                     if bat_soc_start < capacity:
-                        if pv_surplus[i] <= max_charge_power:
-                            charge_amount = pv_surplus[i]
+                        bat_gap = capacity - bat_soc_start
+                        if pv_surplus[i] >= bat_gap:
+                            charge_amount = min(bat_gap, max_charge_power)
                         else:
-                            charge_amount = max_charge_power
+                            charge_amount = min(pv_surplus[i], max_charge_power)
+                        self.BatSoC[i] = (bat_soc_start + charge_amount * charge_efficiency)
+                        self.PV2Bat[i] = charge_amount
                         pv_surplus_after_battery[i] -= charge_amount
-                        self.BatSoC[i] = (
-                            bat_soc_start + charge_amount * charge_efficiency
-                        )
-                        self.PV2Bat[i] = charge_amount * charge_efficiency
-                        if self.BatSoC[i] > capacity:
-                            right_charge_amount = (
-                                capacity - bat_soc_start
-                            ) / charge_efficiency
-                            pv_surplus_after_battery[i] += (
-                                charge_amount - right_charge_amount
-                            )
-                            self.BatSoC[i] = capacity
-                            self.PV2Bat[i] = capacity - bat_soc_start
-                        else:
-                            pass
                     else:
                         self.BatSoC[i] = capacity
                         self.PV2Bat[i] = 0
                 else:
                     if bat_soc_start > 0:
-                        if bat_soc_start < max_discharge_power:
-                            discharge_amount = bat_soc_start
+                        if bat_soc_start * discharge_efficiency <= grid_demand[i]:
+                            discharge_amount = min(bat_soc_start, max_discharge_power)
                         else:
-                            discharge_amount = max_discharge_power
-                        grid_demand_after_battery[i] = (
-                            grid_demand[i] - discharge_amount * discharge_efficiency
-                        )
+                            discharge_amount = min(grid_demand[i], max_discharge_power)
+                        grid_demand_after_battery[i] = grid_demand[i] - discharge_amount * discharge_efficiency
                         self.BatSoC[i] = bat_soc_start - discharge_amount
                         self.Bat2Load[i] = discharge_amount * discharge_efficiency
-                        if grid_demand_after_battery[i] < 0:
-                            discharge_amount = grid_demand[i] / discharge_efficiency
-                            grid_demand_after_battery[i] = 0
-                            self.BatSoC[i] = bat_soc_start - discharge_amount
-                            self.Bat2Load[i] = grid_demand[i]
-                        else:
-                            pass
                     else:
                         grid_demand_after_battery[i] = grid_demand[i]
                         self.BatSoC[i] = 0