import unittest

from components.pid.temp_controller_smith import TempController
from components.pid.temp_controller_fsm import States
from components.plant.pot import Pot


DT = 1.0
AMBIENT = 20.0
PLANT_PARAMS = {"M": 27.96, "C": 3403.43, "L": 0.2, "Td": 17}

OUTER_PARAMS = {"kp": 0.6, "ki": 0.0, "kd": 0.0, "kt": 0.0}
INNER_HEAT_PARAMS = {"kp": 0.08, "ki": 0.02, "kd": 0.0, "kt": 1.5}
INNER_HOLD_CLAMPED = {"kp": 0.08, "ki": 0.02, "kd": 0.0, "kt": 1.5, "yi_max": 0.3}
INNER_HOLD_UNCLAMPED = dict(INNER_HEAT_PARAMS)


def make_controller(inner_hold_params):
    ctrl = TempController(DT)
    ctrl.set_params({
        "beta": 0.9,
        "Outer": dict(OUTER_PARAMS),
        "Inner": {
            "Heat": dict(INNER_HEAT_PARAMS),
            "Hold": dict(inner_hold_params),
            "Cool": dict(INNER_HEAT_PARAMS),
        },
    })
    ctrl.set_model_plant_params(PLANT_PARAMS)
    ctrl.set_ambient_temperature(AMBIENT)
    ctrl.set_enabled(True)
    return ctrl


def make_plant(initial_temp):
    plant = Pot(DT)
    plant.set_plant_params(PLANT_PARAMS)
    plant.set_ambient_temperature(AMBIENT)
    plant.initial(initial_temp)
    return plant


def tick(ctrl, plant, temp_soll, heatrate_soll):
    plant.process()
    ctrl.set_theta_ist(plant.get_temperature())
    ctrl.set_theta_soll(temp_soll)
    ctrl.set_heatrate_soll(heatrate_soll)
    ctrl.process()
    power = 3500.0 * max(0.0, ctrl.get_power())
    plant.set_power(power)
    ctrl.set_model_power(power)


class TestHoldDisturbanceOvershoot(unittest.TestCase):
    """Reproduces the cold-water-during-HOLD incident from
    docs/overshoot_hold_windup.md and checks that Inner.Hold's yi_max
    measurably reduces the resulting overshoot."""

    def _run_disturbance(self, inner_hold_params):
        ctrl = make_controller(inner_hold_params)
        plant = make_plant(30.0)

        # Settle at HOLD, 30 degC.
        for _ in range(60):
            tick(ctrl, plant, 30.0, 1.0)
        self.assertEqual(ctrl.state, States.HOLD)

        # Cold-water disturbance: knock plant.temp down ~0.5 degC.
        plant.temp -= 0.5

        peak = plant.get_temperature()
        for _ in range(600):
            tick(ctrl, plant, 30.0, 1.0)
            peak = max(peak, plant.get_temperature())
            # Disturbance must stay within HOLD the whole time, matching
            # the incident (diff never reached HoldHeat).
            self.assertEqual(ctrl.state, States.HOLD)

        return peak

    def test_clamped_overshoot_smaller_than_unclamped(self):
        peak_clamped = self._run_disturbance(INNER_HOLD_CLAMPED)
        peak_unclamped = self._run_disturbance(INNER_HOLD_UNCLAMPED)

        overshoot_clamped = peak_clamped - 30.0
        overshoot_unclamped = peak_unclamped - 30.0

        self.assertLess(overshoot_clamped, overshoot_unclamped)


class TestRealRampStillReachesTarget(unittest.TestCase):
    """Guards against reintroducing the rejected flat-clamp regression:
    Inner.Heat has no yi_max, so a genuine ramp must still be able to
    reach its commanded heat rate."""

    def test_ramp_reaches_commanded_heatrate(self):
        ctrl = make_controller(INNER_HOLD_CLAMPED)
        plant = make_plant(20.0)

        heatrate_soll = 1.5
        temp_soll = 40.0

        max_heatrate_ist = 0.0
        reached_heat = False
        for _ in range(600):
            tick(ctrl, plant, temp_soll, heatrate_soll)
            if ctrl.state == States.HEAT:
                reached_heat = True
                max_heatrate_ist = max(max_heatrate_ist, ctrl.get_heatrate_ist())

        self.assertTrue(reached_heat)
        self.assertGreater(max_heatrate_ist, 1.4)


class TestHoldHeatHoldTransitionIsBumpless(unittest.TestCase):
    """Per-state yi_max swap must not itself introduce a discontinuity in y
    at a HOLD<->HEAT transition - the same Pid instance/state carries over,
    only the yi_max ceiling changes going forward."""

    def test_no_bump_at_transitions(self):
        ctrl = make_controller(INNER_HOLD_CLAMPED)
        plant = make_plant(20.0)

        y_prev = None
        state_prev = None
        max_bump = 0.0
        for _ in range(1000):
            tick(ctrl, plant, 40.0, 1.5)
            y = ctrl.get_power()
            if state_prev is not None and ctrl.state != state_prev and \
                    {state_prev, ctrl.state} <= {States.HOLD, States.HEAT}:
                max_bump = max(max_bump, abs(y - y_prev))
            y_prev = y
            state_prev = ctrl.state

        # A "bump" here would be a y-step far larger than what one tick of
        # normal PID evolution produces elsewhere in the same run - not
        # literally zero, since heatrate_err itself keeps evolving tick to
        # tick regardless of the state transition. Note: a HEAT->HOLD
        # transition can retroactively clamp yi if a sustained ramp pushed
        # it past Inner.Hold's yi_max before HeatHold's threshold fired,
        # producing a small (not literally bumpless) step - bounded well
        # below the full-freeze/thaw magnitude this replaces.
        self.assertLess(max_bump, 0.1)


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