"""
test_adaptive_pid.py — Unit tests for pid_controller.py and weight_estimator.py
(Issue #136).  No ROS2 runtime required — pure Python.
"""

import math
import pytest

from saltybot_adaptive_pid.pid_controller import (
    GainSet,
    GainScheduler,
    InstabilityDetector,
    PIDController,
)
from saltybot_adaptive_pid.weight_estimator import (
    CalibrationRoutine,
    CalibrationState,
    WeightEstimator,
)


# ─── GainSet ─────────────────────────────────────────────────────────────────

class TestGainSet:

    def test_fields(self):
        g = GainSet(kp=1.0, ki=0.1, kd=0.5)
        assert g.kp == pytest.approx(1.0)
        assert g.ki == pytest.approx(0.1)
        assert g.kd == pytest.approx(0.5)

    def test_clamp_within_bounds_unchanged(self):
        g = GainSet(5.0, 1.0, 2.0).clamp(0, 10, 0, 5, 0, 5)
        assert g.kp == pytest.approx(5.0)
        assert g.ki == pytest.approx(1.0)
        assert g.kd == pytest.approx(2.0)

    def test_clamp_below_min(self):
        g = GainSet(-1.0, -0.5, -2.0).clamp(0, 10, 0, 5, 0, 5)
        assert g.kp == pytest.approx(0.0)
        assert g.ki == pytest.approx(0.0)
        assert g.kd == pytest.approx(0.0)

    def test_clamp_above_max(self):
        g = GainSet(200.0, 50.0, 100.0).clamp(0, 10, 0, 5, 0, 5)
        assert g.kp == pytest.approx(10.0)
        assert g.ki == pytest.approx(5.0)
        assert g.kd == pytest.approx(5.0)

    def test_clamp_returns_new_instance(self):
        g = GainSet(5.0, 1.0, 2.0)
        g2 = g.clamp(0, 10, 0, 5, 0, 5)
        assert g2 is not g


# ─── PIDController ───────────────────────────────────────────────────────────

class TestPIDController:

    def _make(self, kp=1.0, ki=0.0, kd=0.0, **kwargs):
        return PIDController(kp=kp, ki=ki, kd=kd, **kwargs)

    def test_p_only(self):
        pid = self._make(kp=2.0)
        assert pid.step(1.0, 0.01) == pytest.approx(2.0)

    def test_p_negative_error(self):
        pid = self._make(kp=3.0)
        assert pid.step(-1.0, 0.01) == pytest.approx(-3.0)

    def test_i_accumulates(self):
        pid = self._make(kp=0.0, ki=1.0)   # isolate I term
        # After 3 steps of error=1, dt=0.1: integral = 0.3, output = ki*integral = 0.3
        pid.step(1.0, 0.1)
        pid.step(1.0, 0.1)
        out = pid.step(1.0, 0.1)
        assert out == pytest.approx(0.3, abs=1e-9)

    def test_d_zero_on_first_step(self):
        pid = self._make(kp=0.0, kd=10.0)   # isolate D term
        out = pid.step(1.0, 0.01)
        assert out == pytest.approx(0.0)   # D is zero on first step

    def test_d_on_rate_change(self):
        pid = self._make(kp=0.0, kd=1.0)   # isolate D term
        pid.step(0.0, 0.1)          # first step, D=0
        out = pid.step(1.0, 0.1)   # error changed by 1 over dt=0.1 → D=1/0.1=10
        assert out == pytest.approx(10.0)

    def test_output_min_clamp(self):
        # Negative error → large negative output → clamped at output_min
        pid = self._make(kp=100.0, output_min=-5.0)
        out = pid.step(-1.0, 0.01)
        assert out == pytest.approx(-5.0)

    def test_output_max_clamp(self):
        pid = self._make(kp=100.0, output_max=5.0)
        out = pid.step(1.0, 0.01)
        assert out == pytest.approx(5.0)

    def test_anti_windup_clamps_integral(self):
        pid = self._make(ki=1.0, i_clamp=0.5)
        for _ in range(100):
            pid.step(1.0, 0.1)    # integrate a lot
        assert abs(pid.integral) <= 0.5 + 1e-9

    def test_reset_clears_integral(self):
        pid = self._make(ki=1.0)
        for _ in range(10):
            pid.step(1.0, 0.1)
        assert pid.integral > 0.0
        pid.reset()
        assert pid.integral == pytest.approx(0.0)

    def test_reset_clears_derivative_state(self):
        pid = self._make(kp=0.0, kd=5.0)   # isolate D term
        pid.step(1.0, 0.1)         # sets prev_error
        pid.reset()
        out = pid.step(1.0, 0.1)  # prev_error=None after reset → D=0
        assert out == pytest.approx(0.0)

    def test_update_gains_takes_effect(self):
        pid = self._make(kp=1.0)
        pid.update_gains(kp=3.0, ki=0.0, kd=0.0)
        assert pid.step(1.0, 0.01) == pytest.approx(3.0)

    def test_zero_dt_returns_zero(self):
        pid = self._make(kp=5.0)
        assert pid.step(1.0, 0.0) == pytest.approx(0.0)

    def test_negative_dt_returns_zero(self):
        pid = self._make(kp=5.0)
        assert pid.step(1.0, -0.1) == pytest.approx(0.0)

    def test_zero_error_no_pid_output(self):
        pid = self._make(kp=5.0, ki=1.0, kd=1.0)
        pid.step(0.0, 0.1)
        out = pid.step(0.0, 0.1)
        assert out == pytest.approx(0.0)

    def test_gain_properties(self):
        pid = self._make(kp=2.0, ki=0.5, kd=0.3)
        assert pid.kp == pytest.approx(2.0)
        assert pid.ki == pytest.approx(0.5)
        assert pid.kd == pytest.approx(0.3)


# ─── GainScheduler ───────────────────────────────────────────────────────────

def _make_scheduler(**kwargs) -> GainScheduler:
    table = {
        "empty": GainSet(10.0, 0.5, 1.0),
        "light": GainSet(15.0, 0.7, 1.5),
        "heavy": GainSet(20.0, 1.0, 2.0),
    }
    thresholds = {"empty_max": 0.5, "light_max": 2.0}
    defaults = dict(blend_alpha=0.05, kp_min=0.0, kp_max=100.0,
                    ki_min=0.0, ki_max=10.0, kd_min=0.0, kd_max=10.0)
    defaults.update(kwargs)
    return GainScheduler(table, thresholds, **defaults)


class TestGainScheduler:

    def test_classify_empty(self):
        s = _make_scheduler()
        assert s.classify(0.0) == "empty"
        assert s.classify(0.5) == "empty"

    def test_classify_light(self):
        s = _make_scheduler()
        assert s.classify(0.51) == "light"
        assert s.classify(2.0) == "light"

    def test_classify_heavy(self):
        s = _make_scheduler()
        assert s.classify(2.01) == "heavy"
        assert s.classify(10.0) == "heavy"

    def test_classify_boundary_empty_light(self):
        s = _make_scheduler()
        assert s.classify(0.5) == "empty"
        assert s.classify(0.500001) == "light"

    def test_initial_class_is_empty(self):
        s = _make_scheduler()
        assert s.active_class == "empty"

    def test_update_blends_toward_target(self):
        s = _make_scheduler(blend_alpha=0.1)
        initial_kp = s.active_gains.kp
        # Push toward "heavy"
        for _ in range(50):
            gains, _, _ = s.update(5.0)
        # Should have moved significantly toward heavy (kp=20)
        assert gains.kp > initial_kp + 1.0

    def test_blend_does_not_overshoot(self):
        s = _make_scheduler(blend_alpha=0.5)
        target_kp = 20.0  # heavy
        for _ in range(200):
            gains, _, _ = s.update(5.0)
        assert gains.kp <= target_kp + 1e-6

    def test_class_changed_flag_true_on_change(self):
        s = _make_scheduler()
        s.update(0.1)   # empty
        _, _, changed = s.update(5.0)   # → heavy
        assert changed is True

    def test_class_changed_flag_false_same_class(self):
        s = _make_scheduler()
        s.update(0.1)
        _, _, changed = s.update(0.2)   # still empty
        assert changed is False

    def test_override_gains_applied(self):
        s = _make_scheduler()
        s.override_gains(kp=7.0, ki=0.3, kd=0.8)
        assert s.active_gains.kp == pytest.approx(7.0)
        assert s.active_gains.ki == pytest.approx(0.3)
        assert s.override_active is True

    def test_override_bypasses_scheduling(self):
        s = _make_scheduler()
        s.override_gains(kp=7.0, ki=0.3, kd=0.8)
        for _ in range(20):
            gains, _, _ = s.update(5.0)   # heavy weight, should not change
        assert gains.kp == pytest.approx(7.0)

    def test_override_gains_clamped(self):
        s = _make_scheduler(kp_max=10.0)
        s.override_gains(kp=999.0, ki=0.5, kd=1.0)
        assert s.active_gains.kp == pytest.approx(10.0)

    def test_clear_override_resumes_scheduling(self):
        s = _make_scheduler()
        s.override_gains(kp=7.0, ki=0.3, kd=0.8)
        s.clear_override()
        assert s.override_active is False

    def test_revert_to_defaults_snaps_immediately(self):
        s = _make_scheduler(blend_alpha=0.001)
        for _ in range(5):
            s.update(5.0)   # start blending toward heavy (kp=20)
        s.revert_to_defaults(0.0)   # snap back to empty (kp=10)
        assert s.active_gains.kp == pytest.approx(10.0)
        assert s.active_class == "empty"

    def test_revert_clears_override(self):
        s = _make_scheduler()
        s.override_gains(kp=7.0, ki=0.3, kd=0.8)
        s.revert_to_defaults(0.0)
        assert s.override_active is False

    def test_safety_bounds_clamp_table_gains(self):
        s = _make_scheduler(kp_max=12.0)   # cap below heavy table value (20)
        for _ in range(200):
            gains, _, _ = s.update(5.0)    # heavy
        assert gains.kp <= 12.0 + 1e-9


# ─── InstabilityDetector ─────────────────────────────────────────────────────

class TestInstabilityDetector:

    def _make(self, **kwargs):
        defaults = dict(tilt_threshold_rad=0.4, oscillation_count=6, window_size=50)
        defaults.update(kwargs)
        return InstabilityDetector(**defaults)

    def test_stable_small_tilts(self):
        det = self._make()
        for _ in range(50):
            assert det.update(0.1) is False

    def test_too_few_samples_always_stable(self):
        det = self._make()
        assert det.update(1.0) is False   # only 1 sample
        assert det.update(1.0) is False
        assert det.update(1.0) is False

    def test_unstable_tilt_exceeds_threshold(self):
        det = self._make(tilt_threshold_rad=0.3, window_size=10)
        for _ in range(6):
            det.update(0.5)   # > 0.3, majority
        assert det.update(0.5) is True

    def test_stable_below_threshold(self):
        det = self._make(tilt_threshold_rad=0.5)
        for _ in range(20):
            result = det.update(0.3)
        assert result is False

    def test_oscillation_triggers_unstable(self):
        det = self._make(oscillation_count=4, window_size=50)
        # Inject 8 alternating signs → 7 reversals ≥ 4
        tilts = [0.1, -0.1] * 8
        results = [det.update(t) for t in tilts]
        assert any(results)

    def test_reset_clears_window(self):
        det = self._make(tilt_threshold_rad=0.3, window_size=10)
        for _ in range(10):
            det.update(0.5)
        det.reset()
        # After reset, too few samples → stable
        assert det.update(0.5) is False


# ─── WeightEstimator ─────────────────────────────────────────────────────────

def _make_estimator(**kwargs) -> WeightEstimator:
    defaults = dict(
        idle_current_a=0.5, current_per_kg_a=0.3,
        window_s=1.0, control_rate_hz=10.0,  # window_size=10
        steady_tilt_threshold_rad=0.05, change_threshold_kg=0.3,
    )
    defaults.update(kwargs)
    return WeightEstimator(**defaults)


class TestWeightEstimator:

    def test_insufficient_samples_returns_none(self):
        est = _make_estimator()
        result, _ = est.update(0.5, tilt_rad=0.0)
        assert result is None

    def test_sufficient_samples_returns_estimate(self):
        est = _make_estimator()
        for _ in range(10):
            est.update(0.8, tilt_rad=0.0)   # 0.8A → (0.8-0.5)/0.3 = 1.0 kg
        wt, _ = est.update(0.8, tilt_rad=0.0)
        assert wt == pytest.approx(1.0, abs=0.05)

    def test_non_steady_samples_rejected(self):
        est = _make_estimator()
        for _ in range(20):
            est.update(0.8, tilt_rad=0.1)   # |0.1| > 0.05 → rejected
        result, _ = est.update(0.8, tilt_rad=0.1)
        assert result is None

    def test_known_current_gives_correct_weight(self):
        # 1.1A, idle=0.5, per_kg=0.3 → (1.1-0.5)/0.3 = 2.0 kg
        est = _make_estimator()
        for _ in range(15):
            est.update(1.1, tilt_rad=0.0)
        wt, _ = est.update(1.1, tilt_rad=0.0)
        assert wt == pytest.approx(2.0, abs=0.05)

    def test_negative_estimate_clamped_to_zero(self):
        # current below idle → weight should be 0
        est = _make_estimator()
        for _ in range(15):
            est.update(0.1, tilt_rad=0.0)
        wt, _ = est.update(0.1, tilt_rad=0.0)
        assert wt == pytest.approx(0.0)

    def test_weight_changed_flag_on_first_estimate(self):
        est = _make_estimator()
        # Find the first non-None estimate — it must have changed=True
        for _ in range(15):
            wt, changed = est.update(0.8, tilt_rad=0.0)
            if wt is not None:
                assert changed is True
                break

    def test_weight_unchanged_below_threshold(self):
        est = _make_estimator(change_threshold_kg=0.5)
        for _ in range(15):
            est.update(0.8, tilt_rad=0.0)
        # Now feed very similar current (tiny perturbation)
        for _ in range(15):
            wt, changed = est.update(0.81, tilt_rad=0.0)
        # delta is ~0.03 kg << threshold 0.5
        assert changed is False

    def test_reset_clears_samples_and_estimate(self):
        est = _make_estimator()
        for _ in range(15):
            est.update(0.8, tilt_rad=0.0)
        est.reset()
        assert est.last_estimate is None
        result, _ = est.update(0.8, tilt_rad=0.0)
        assert result is None   # need to re-fill window


# ─── CalibrationRoutine ───────────────────────────────────────────────────────

def _make_cal(**kwargs) -> CalibrationRoutine:
    defaults = dict(
        amplitude_rad=0.05, frequency_hz=1.0,
        rock_cycles=2, settle_s=1.0,
        settle_tilt_threshold_rad=0.03,
        idle_current_a=0.5, current_per_kg_a=0.3,
    )
    defaults.update(kwargs)
    return CalibrationRoutine(**defaults)


class TestCalibrationRoutine:

    def test_initial_state_idle(self):
        cal = _make_cal()
        assert cal.state == CalibrationState.IDLE
        assert cal.is_active is False

    def test_start_transitions_to_rocking(self):
        cal = _make_cal()
        assert cal.start() is True
        assert cal.state == CalibrationState.ROCKING
        assert cal.is_active is True

    def test_start_returns_false_when_active(self):
        cal = _make_cal()
        cal.start()
        assert cal.start() is False

    def test_start_succeeds_from_done(self):
        cal = _make_cal(rock_cycles=1, settle_s=0.01)
        cal.start()
        dt = 0.01
        for _ in range(200):
            cal.step(0.8, tilt_rad=0.0, dt=dt)
        # Should be DONE or FAILED at this point
        assert cal.state in (CalibrationState.DONE, CalibrationState.FAILED)
        assert cal.start() is True   # can restart

    def test_abort_returns_to_idle(self):
        cal = _make_cal()
        cal.start()
        cal.abort()
        assert cal.state == CalibrationState.IDLE
        assert cal.is_active is False

    def test_rocking_output_sinusoidal(self):
        # freq=1Hz, cycles=2 → duration=2s; use 110 steps at dt=0.01 to cover >0.5 period
        cal = _make_cal(amplitude_rad=0.1, frequency_hz=1.0)
        cal.start()
        outputs = [cal.step(0.5, 0.0, 0.01) for _ in range(110)]
        # Sinusoidal over >1 s: must cross zero and include both signs
        assert any(o > 0.0 for o in outputs)
        assert any(o < 0.0 for o in outputs)

    def test_rocking_output_bounded_by_amplitude(self):
        cal = _make_cal(amplitude_rad=0.05)
        cal.start()
        for _ in range(200):
            out = cal.step(0.5, 0.0, 0.005)
            assert abs(out) <= 0.05 + 1e-9

    def test_idle_output_is_zero(self):
        cal = _make_cal()
        assert cal.step(0.5, 0.0, 0.01) == pytest.approx(0.0)

    def test_rocking_transitions_to_settling(self):
        # rock_cycles=2, freq=1Hz → rock_duration=2s; dt=0.1 → 20 steps
        cal = _make_cal(rock_cycles=2, frequency_hz=1.0)
        cal.start()
        for _ in range(22):
            cal.step(0.5, 0.0, 0.1)
        assert cal.state == CalibrationState.SETTLING

    def test_settling_output_is_zero(self):
        cal = _make_cal(rock_cycles=1, frequency_hz=2.0)
        cal.start()
        # Skip through rocking
        for _ in range(15):
            cal.step(0.5, 0.0, 0.1)
        # Now in settling
        assert cal.state == CalibrationState.SETTLING
        out = cal.step(0.5, 0.0, 0.1)
        assert out == pytest.approx(0.0)

    def test_done_with_enough_settled_samples(self):
        # freq=2Hz, cycles=1 → rock_duration=0.5s; dt=0.05 → 10 steps
        # settle_s=0.5s → 10 more steps
        cal = _make_cal(rock_cycles=1, frequency_hz=2.0, settle_s=0.5)
        cal.start()
        for _ in range(100):
            cal.step(0.8, tilt_rad=0.0, dt=0.05)
        assert cal.state == CalibrationState.DONE
        assert cal.weight_estimate is not None
        assert cal.weight_estimate >= 0.0

    def test_weight_estimate_known_current(self):
        # 0.8A, idle=0.5, per_kg=0.3 → 1.0 kg
        cal = _make_cal(rock_cycles=1, frequency_hz=2.0, settle_s=0.5)
        cal.start()
        for _ in range(100):
            cal.step(0.8, tilt_rad=0.0, dt=0.05)
        assert cal.state == CalibrationState.DONE
        assert cal.weight_estimate == pytest.approx(1.0, abs=0.1)

    def test_failed_when_no_settled_samples(self):
        # Robot never settles (large tilt)
        cal = _make_cal(
            rock_cycles=1, frequency_hz=2.0, settle_s=0.5,
            settle_tilt_threshold_rad=0.01
        )
        cal.start()
        for _ in range(100):
            cal.step(0.8, tilt_rad=0.1, dt=0.05)   # tilt > threshold → no samples
        assert cal.state == CalibrationState.FAILED
        assert cal.weight_estimate is None

    def test_rocking_duration(self):
        """Rocking lasts exactly rock_cycles / frequency_hz seconds."""
        freq = 1.0
        cycles = 3
        cal = _make_cal(rock_cycles=cycles, frequency_hz=freq, settle_s=100.0)
        cal.start()
        rock_dur = cycles / freq   # 3 s
        dt = 0.01
        t = 0.0
        while cal.state == CalibrationState.ROCKING and t < rock_dur + 0.1:
            cal.step(0.5, 0.0, dt)
            t += dt
        assert t == pytest.approx(rock_dur, abs=0.02)

    def test_weight_clamped_to_zero_below_idle(self):
        # Current below idle → should clamp to 0, not negative
        cal = _make_cal(rock_cycles=1, frequency_hz=2.0, settle_s=0.5)
        cal.start()
        for _ in range(100):
            cal.step(0.2, tilt_rad=0.0, dt=0.05)  # 0.2 < 0.5 idle
        if cal.state == CalibrationState.DONE:
            assert cal.weight_estimate == pytest.approx(0.0)


# ─── Integration ─────────────────────────────────────────────────────────────

class TestIntegration:

    def test_pid_with_gain_scheduler_round_trip(self):
        """GainScheduler feeds gains into PID; both operate consistently."""
        sched = _make_scheduler()
        pid = PIDController(kp=10.0, ki=0.5, kd=1.0,
                            output_min=-50.0, output_max=50.0)

        # Simulate 1 s at 100 Hz, weight=0.2 kg (empty class)
        dt = 0.01
        total_effort = 0.0
        for _ in range(100):
            gains, _, _ = sched.update(0.2)
            pid.update_gains(gains.kp, gains.ki, gains.kd)
            error = 0.1   # constant tilt error
            total_effort += pid.step(error, dt)

        assert total_effort > 0.0   # robot should be correcting

    def test_instability_reverts_pid_gains(self):
        sched = _make_scheduler()
        det = InstabilityDetector(tilt_threshold_rad=0.3, window_size=10)
        pid = PIDController(kp=10.0, ki=0.0, kd=0.0)

        # Inject large oscillations
        tilts = [0.5, -0.5] * 10
        for t in tilts:
            det.update(t)
            if det.update(t):
                sched.revert_to_defaults(0.0)
                gains = sched.active_gains
                pid.update_gains(gains.kp, gains.ki, gains.kd)
                pid.reset()
                break

        assert sched.active_class == "empty"
        assert pid.integral == pytest.approx(0.0)

    def test_calibration_updates_weight_estimator_chain(self):
        """Full: calibration → weight estimate → gain scheduling."""
        cal = _make_cal(rock_cycles=1, frequency_hz=2.0, settle_s=0.5)
        sched = _make_scheduler()

        # Run calibration
        cal.start()
        for _ in range(100):
            cal.step(0.8, tilt_rad=0.0, dt=0.05)

        assert cal.state == CalibrationState.DONE
        wt = cal.weight_estimate
        assert wt is not None

        # Feed weight into scheduler
        for _ in range(5):
            gains, cls, _ = sched.update(wt)

        # ~1 kg → "light" class
        assert cls == "light"
