import unittest

import pandas as pd

from methods.battery_optimizer import DT_H, optimize_battery_plan


class BatteryCostOptimizerTest(unittest.TestCase):
    def plan(self, load, pv, import_prices=None, export_prices=None, **overrides):
        index = pd.date_range("2026-09-29T00:00:00", periods=len(load), freq="5min")
        config = {
            "batt_capacity_kwh": 2.0,
            "batt_power_kw": 1.0,
            "batt_min_soc": 0.0,
            "batt_max_soc": 100.0,
            "batt_charge_efficiency": 1.0,
            "batt_discharge_efficiency": 1.0,
            "batt_soc_percent": 0.0,
            "batt_grid_charging_enabled": False,
            "batt_degradation_chf_kwh": 0.0,
            "tarif_bezug": "dynamic",
            "tarif_einspeisung": "dynamic",
            "tarif_peak_fest": 0.0,
        }
        config.update(overrides)
        data = {
            "config": config,
            "df_fut": pd.DataFrame({
                "import_price": import_prices or [0.30] * len(index),
                "export_price": export_prices or [0.10] * len(index),
            }, index=index),
            "current_soc_perc": config["batt_soc_percent"],
            "current_month_peak_kw": overrides.get("current_month_peak_kw", 0.0),
        }
        return index, optimize_battery_plan(data, dict(zip(index, pv)), dict(zip(index, load)))

    def test_grid_charging_is_opt_in(self):
        index, plan = self.plan(
            [500.0] * 4,
            [0.0] * 4,
            import_prices=[0.05, 0.05, 0.50, 0.50],
        )
        self.assertTrue(all(plan["battery"][t] <= 1e-6 for t in index))

    def test_cheap_grid_energy_is_shifted_to_expensive_period(self):
        index, plan = self.plan(
            [500.0] * 4,
            [0.0] * 4,
            import_prices=[0.05, 0.05, 0.50, 0.50],
            batt_grid_charging_enabled=True,
        )
        self.assertGreater(plan["battery"][index[0]], 0.0)
        self.assertLess(plan["battery"][index[-1]], 0.0)
        self.assertGreater(plan["grid"][index[0]], 500.0)
        self.assertAlmostEqual(plan["grid"][index[-1]], 0.0, places=5)

    def test_high_feed_in_value_prefers_export(self):
        index, plan = self.plan(
            [0.0, 1000.0],
            [1000.0, 0.0],
            import_prices=[0.20, 0.20],
            export_prices=[0.60, 0.60],
        )
        self.assertAlmostEqual(plan["grid"][index[0]], -1000.0, places=5)
        self.assertAlmostEqual(plan["grid"][index[1]], 1000.0, places=5)

    def test_peak_tariff_prevents_grid_charge_above_existing_peak(self):
        index, plan = self.plan(
            [1000.0] * 6,
            [0.0] * 6,
            import_prices=[0.05] * 3 + [0.50] * 3,
            batt_grid_charging_enabled=True,
            tarif_peak_fest=20.0,
            current_month_peak_kw=1.0,
        )
        self.assertLessEqual(plan["planned_peak_kw"], 1.0 + 1e-7)
        self.assertTrue(all(plan["grid"][t] <= 1000.0 + 1e-5 for t in index))

    def test_nearly_empty_battery_is_not_discharged_further_at_low_value(self):
        index, plan = self.plan(
            [1000.0] * 4,
            [0.0] * 4,
            import_prices=[0.01] * 4,
            export_prices=[0.0] * 4,
            batt_soc_percent=5.0,
            batt_economic_reserve_soc_percent=10.0,
            batt_terminal_value_chf_kwh=0.0,
        )
        self.assertTrue(all(plan["battery"][timestamp] >= -1e-6 for timestamp in index))
        self.assertTrue(all(abs(plan["grid"][timestamp] - 1000.0) <= 1e-5 for timestamp in index))
        self.assertAlmostEqual(plan["economic_min_soc_percent"], 5.0)

    def test_profitable_export_discharge_remains_allowed_above_reserve(self):
        index, plan = self.plan(
            [100.0, 100.0],
            [1000.0, 1000.0],
            import_prices=[0.20, 0.20],
            export_prices=[0.80, 0.80],
            batt_soc_percent=100.0,
            batt_economic_reserve_soc_percent=10.0,
            batt_degradation_chf_kwh=0.03,
            batt_terminal_value_chf_kwh=0.0,
        )
        self.assertTrue(any(plan["battery"][timestamp] < -1e-6 for timestamp in index))
        self.assertTrue(all(plan["grid"][timestamp] <= -899.0 for timestamp in index))
        self.assertAlmostEqual(plan["economic_min_soc_percent"], 10.0)

    def test_power_and_soc_limits_hold(self):
        load = [0.0] * 12 + [1000.0] * 12
        pv = [1000.0] * 12 + [0.0] * 12
        index, plan = self.plan(
            load,
            pv,
            batt_capacity_kwh=1.0,
            batt_power_kw=0.5,
            batt_min_soc=20.0,
            batt_max_soc=80.0,
            batt_soc_percent=20.0,
        )
        soc = 0.2
        for timestamp in index:
            target = plan["battery"][timestamp]
            self.assertLessEqual(abs(target), 500.0 + 1e-6)
            if target >= 0.0:
                soc += target * DT_H / 1000.0
            else:
                soc += target * DT_H / 1000.0
            self.assertGreaterEqual(soc, 0.2 - 1e-8)
            self.assertLessEqual(soc, 0.8 + 1e-8)


if __name__ == "__main__":
    unittest.main()