# Copyright 2026 Kemal Yaylali
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Compensators and the Kalman bank.

The inverse-property tests here exist because getting that algebra wrong has
already cost this project twice: once on temperature, where a mismatched
simulator injected 1.2 C of phantom noise floor, and once on humidity, where
the correction ran the wrong way against a reference hygrometer.
"""

from __future__ import annotations

import numpy as np
import pytest

from ashvale.estimation import HumidityCompensator, KalmanCV, ThermalCompensator
from ashvale.physics import dew_point, saturation_vapour_pressure

# ---------------------------------------------------------------- thermal

def test_thermal_forward_model_is_the_exact_inverse_of_the_compensator():
    """T_raw = (T + k*T_cpu)/(1+k) must invert T = T_raw - k(T_cpu - T_raw)."""
    for k, t_true, t_cpu in [(0.55, 19.0, 40.0), (0.26, 24.4, 40.2), (1.0, 5.0, 30.0)]:
        c = ThermalCompensator(k0=k, k_min=0.0, k_max=2.0)
        t_raw = (t_true + k * t_cpu) / (1.0 + k)
        assert c.compensate(t_raw, t_cpu) == pytest.approx(t_true, abs=1e-9)


def test_thermal_calibration_moves_k_toward_the_truth():
    c = ThermalCompensator(k0=0.30, k_min=0.05, k_max=1.5)
    k_true, t_true, t_cpu = 0.62, 19.0, 41.0
    t_raw = (t_true + k_true * t_cpu) / (1.0 + k_true)
    before = abs(c.k - k_true)
    c.calibrate(t_raw, t_cpu, t_true)
    assert abs(c.k - k_true) < before


def test_thermal_clamp_survives_a_mistyped_reference():
    c = ThermalCompensator(k0=0.55, k_min=0.15, k_max=1.20)
    for _ in range(50):
        c.calibrate(25.0, 40.0, -300.0)      # absurd reference
    assert c.k_min <= c.k <= c.k_max


def test_thermal_compensation_is_a_noop_without_a_gradient():
    c = ThermalCompensator(k0=0.8)
    assert c.compensate(21.0, 21.0) == pytest.approx(21.0)
    # and never amplifies when the CPU is cooler than the sensor
    assert c.compensate(21.0, 15.0) == pytest.approx(21.0)


# ---------------------------------------------------------------- humidity

def test_humidity_psychrometric_round_trip():
    """The simulator's forward model must invert the compensator exactly."""
    rh_true, t_true, t_raw = 62.0, 19.0, 25.6
    rh_sensor = rh_true * float(saturation_vapour_pressure(t_true) /
                                saturation_vapour_pressure(t_raw))
    hc = HumidityCompensator(psychrometric=True)
    assert hc.compensate(rh_sensor, t_raw, t_true) == pytest.approx(rh_true, abs=1e-6)


def test_humidity_psychrometric_preserves_dew_point():
    """Vapour pressure is the conserved quantity, so dew point must not move."""
    rh_sensor, t_raw, t_true = 60.0, 25.6, 19.0
    hc = HumidityCompensator(psychrometric=True)
    out = hc.compensate(rh_sensor, t_raw, t_true)
    assert float(dew_point(t_true, out)) == pytest.approx(float(dew_point(t_raw, rh_sensor)),
                                                          abs=1e-6)


def test_humidity_psychrometric_disabled_by_default():
    hc = HumidityCompensator()
    assert hc.compensate(60.0, 25.6, 19.0) == pytest.approx(60.0)


def test_humidity_offset_converges_on_a_reference():
    """The measured case: board reads 75.35% where the truth is 50.4%."""
    hc = HumidityCompensator()
    errors = []
    for _ in range(6):
        hc.calibrate(75.35, 27.94, 24.86, 50.4)
        errors.append(abs(hc.compensate(75.35, 27.94, 24.86) - 50.4))
    assert errors[-1] < errors[0]
    assert errors[-1] < 0.5


def test_humidity_offset_is_clamped():
    hc = HumidityCompensator()
    for _ in range(50):
        hc.calibrate(50.0, 20.0, 20.0, 100.0)
    assert hc.off_min <= hc.offset <= hc.off_max


def test_humidity_output_stays_in_range():
    hc = HumidityCompensator(offset=30.0)
    assert 0.0 <= hc.compensate(95.0, 20.0, 20.0) <= 100.0
    hc2 = HumidityCompensator(offset=-30.0)
    assert 0.0 <= hc2.compensate(5.0, 20.0, 20.0) <= 100.0


def test_humidity_state_round_trips_through_dict():
    hc = HumidityCompensator(offset=-24.2, psychrometric=True)
    hc.calibrate(70.0, 25.0, 21.0, 50.0)
    back = HumidityCompensator.from_dict(hc.to_dict())
    assert back.offset == pytest.approx(hc.offset)
    assert back.psychrometric is hc.psychrometric
    assert back.n_calibrations == hc.n_calibrations


# ---------------------------------------------------------------- kalman

def test_kalman_covariance_stays_symmetric_and_psd():
    """Joseph form exists precisely so this holds over a long run."""
    kf = KalmanCV(q=1e-6, r=0.05)
    rng = np.random.default_rng(7)
    for _ in range(20000):
        kf.update(20.0 + 0.05 * rng.normal(), 2.0)
    P = np.asarray(kf.P, dtype=float)
    assert np.allclose(P, P.T, atol=1e-12)
    assert np.all(np.linalg.eigvalsh(P) > -1e-12)


def test_kalman_tracks_a_constant_and_reports_zero_rate():
    kf = KalmanCV(q=1e-8, r=0.01)
    for _ in range(2000):
        kf.update(15.0, 2.0)
    assert kf.level == pytest.approx(15.0, abs=1e-3)
    assert kf.rate == pytest.approx(0.0, abs=1e-5)


def test_kalman_recovers_a_known_ramp_rate():
    kf = KalmanCV(q=1e-4, r=0.01)
    true_rate = 0.5 / 3600.0          # 0.5 units per hour
    for i in range(6000):
        kf.update(10.0 + true_rate * i * 2.0, 2.0)
    assert kf.rate * 3600.0 == pytest.approx(0.5, rel=0.05)


def test_kalman_ignores_non_finite_measurements():
    kf = KalmanCV(q=1e-6, r=0.05)
    kf.update(20.0, 2.0)
    lvl_before = kf.level
    kf.update(float("nan"), 2.0)
    assert kf.level == pytest.approx(lvl_before)


def test_kalman_nis_is_near_one_when_noise_matches_the_model():
    """NIS is the honest self-check: consistent filter, NIS about 1."""
    r = 0.04
    kf = KalmanCV(q=1e-7, r=r)
    rng = np.random.default_rng(11)
    nis = []
    for i in range(4000):
        kf.update(18.0 + np.sqrt(r) * rng.normal(), 2.0)
        if i > 500:
            nis.append(kf.nis)
    assert 0.5 < float(np.mean(nis)) < 2.0


def test_kalman_state_round_trips_through_dict():
    kf = KalmanCV(q=1e-6, r=0.05)
    for _ in range(50):
        kf.update(12.0, 2.0)
    back = KalmanCV.from_dict(kf.to_dict())
    assert back.level == pytest.approx(kf.level)
    assert back.rate == pytest.approx(kf.rate)


# ---------------------------------------------------------------- thermostat

def test_thermostat_reversion_is_first_order_and_preserves_dew_point():
    """A heated room is a closed loop, and heating adds no moisture.

    Two properties, both easy to get wrong. The temperature must close the gap
    to the setpoint exponentially rather than jumping or drifting, and the
    implied humidity change must leave the dew point exactly where it was: RH
    falls only because es(T) rose, which is why a heated house in winter is dry.
    """
    import math

    from ashvale.config import load_config
    from ashvale.physics import dew_point, saturation_vapour_pressure
    from ashvale.station import Station

    cfg = load_config()
    cfg.site.heating = True
    cfg.site.heating_setpoint_c = 23.0
    cfg.site.thermal_time_constant_h = 1.5
    st = Station(cfg)
    st.live = {"temp_smooth": 18.0}

    tau = 1.5 * 3600.0
    for h in (900, 3600, 10800, 86400):
        expected = (23.0 - 18.0) * (1.0 - math.exp(-h / tau))
        assert st._setpoint_delta("temperature", h, 18.0) == pytest.approx(expected, rel=1e-9)

    # monotonic toward the setpoint, never past it
    deltas = [st._setpoint_delta("temperature", h, 18.0)
              for h in (900, 3600, 10800, 21600, 86400)]
    assert all(a < b for a, b in zip(deltas, deltas[1:]))
    assert deltas[-1] <= 5.0 + 1e-9

    # dew point invariant
    t0, rh0 = 18.0, 55.0
    d_t = st._setpoint_delta("temperature", 86400, t0)
    d_rh = st._setpoint_delta("humidity", 86400, rh0)
    assert float(dew_point(t0 + d_t, rh0 + d_rh)) == pytest.approx(
        float(dew_point(t0, rh0)), abs=1e-6)
    assert d_rh < 0.0, "warming a room at constant moisture must lower RH"
    assert float(saturation_vapour_pressure(t0 + d_t)) > float(
        saturation_vapour_pressure(t0))

    # a thermostat cannot move the synoptic field
    assert st._setpoint_delta("pressure", 86400, 1013.0) == 0.0

    # and off, the member is exactly persistence
    cfg.site.heating = False
    assert st._setpoint_delta("temperature", 86400, 18.0) == 0.0
    assert st._setpoint_delta("humidity", 86400, 55.0) == 0.0


def test_forecast_head_migrates_state_from_before_the_setpoint_member():
    """An old save has three weights where there are now four."""
    from ashvale.models.nowcast import MEMBERS, ForecastHead

    h = ForecastHead(target="temperature", horizon_s=900, n_features=4)
    state = h.to_dict()
    state["weights"] = [0.2, 0.3, 0.5]          # a pre-setpoint save
    state["member_mae"] = [0.4, 0.5, 0.6]
    back = ForecastHead.from_dict(state)
    assert back.weights.size == len(MEMBERS)
    assert float(back.weights.sum()) == pytest.approx(1.0)
    # member_mae must migrate too. Missing it did not fail on load, it failed
    # later inside learn() on a broadcast error, which is a worse place to
    # discover a migration bug.
    assert back.member_mae.size == len(MEMBERS)
    back.learn(np.zeros(4), 20.0, 20.5, 0.1, 0.2)      # must not raise


# ------------------------------------------------- dual-thermometer fusion

def _bare_board():
    from ashvale.sensors import SD_HTS221, SD_LPS25HB, SenseBoard, _ChannelNoise
    b = SenseBoard.__new__(SenseBoard)
    b._noise_h = _ChannelNoise(SD_HTS221)
    b._noise_p = _ChannelNoise(SD_LPS25HB)
    b._gradient = None
    b._gradient_lam = 0.9967
    return b


def _two_channels(n=4000, seed=5):
    from ashvale.sensors import K_HTS221, K_LPS25HB
    rng = np.random.default_rng(seed)
    cpu = 43.0 + 0.5 * np.sin(np.arange(n) / 500.0)
    th = (24.0 + K_HTS221 * cpu) / (1 + K_HTS221) + 0.049 * rng.normal(size=n)
    tp = (24.0 + K_LPS25HB * cpu) / (1 + K_LPS25HB) + 0.007 * rng.normal(size=n)
    return th, tp


def test_fusion_does_not_move_the_mean():
    """The whole point of removing the gradient first.

    The two chips stand about 1.3 C apart, so weighting them by variance drags
    temp_raw onto the quieter one. k was fitted against the mean of the two, and
    after the 1.55x gain of the inverse model that shift becomes about a degree
    of silent bias on every reading downstream.
    """
    th, tp = _two_channels()
    board = _bare_board()
    fused = np.array([board._fuse(th[i], tp[i])[0] for i in range(th.size)])
    avg = (th + tp) / 2.0
    w = slice(1000, None)
    assert abs(fused[w].mean() - avg[w].mean()) < 0.01, "fusion shifted the calibration"


def test_fusion_is_quieter_than_the_average():
    th, tp = _two_channels()
    board = _bare_board()
    fused = np.array([board._fuse(th[i], tp[i])[0] for i in range(th.size)])
    avg = (th + tp) / 2.0
    w = slice(1000, None)

    def wn(x):
        return np.std(np.diff(x)) / np.sqrt(2)

    assert wn(fused[w]) < wn(avg[w]) / 2.0, "fusion did not halve the noise"


def test_fusion_survives_one_dead_channel():
    board = _bare_board()
    value, var = board._fuse(float("nan"), 29.5)
    assert value == 29.5, "a dead HTS221 must not poison the reading"
    value, var = board._fuse(30.5, float("nan"))
    assert value == 30.5
    value, var = board._fuse(float("nan"), float("nan"))
    assert not np.isfinite(value)


def test_kalman_rate_is_physical_in_a_still_room():
    """The tuning failure this guards against.

    On a real station the temperature filter reported a median rate of
    12.4 C/h while the room moved 0.37 C/h. Process noise was set to track
    perhaps a hundred times faster than any of these signals actually move.
    """
    from ashvale.config import CONFIG
    from ashvale.estimation import KalmanCV

    dt = CONFIG.sensor.sample_period_s
    rng = np.random.default_rng(3)
    n = 6000
    truth = 24.0 + 0.4 * np.arange(n) * dt / 3600.0      # a real 0.4 C/h drift
    z = truth + 0.0877 * rng.normal(size=n)              # measured input noise

    kf = KalmanCV(CONFIG.sensor.kalman_q_temp, CONFIG.sensor.kalman_r_temp)
    rates = [kf.update(z[i], dt)[1] * 3600.0 for i in range(n)]
    settled = np.abs(np.array(rates[600:]))

    assert np.median(settled) < 3.0, (
        f"median |rate| {np.median(settled):.1f} C/h in a room drifting 0.4 C/h")
    assert np.percentile(settled, 95) < 10.0


def test_retuning_q_survives_a_reload():
    """Tuning lives in config, not in the state file.

    q and r were persisted and restored, so a retune deployed to a running
    station did nothing: the service restarted and the filters carried on with
    whatever tuning was in force when the state was last written. The symptom
    is a config change that appears to work and does not, which is the worst
    kind.
    """
    from ashvale.config import load_config
    from ashvale.estimation import SignalTracker

    cfg = load_config()
    cfg.sensor.kalman_q_temp = 2.0e-6          # an old, badly tuned state file
    old = SignalTracker(cfg)
    for i in range(50):
        old.step(1.7554e9 + i * 2.0, 24.0, 50.0, 1013.0, 43.0)
    saved = old.to_dict()
    assert saved["filters"]["temperature"]["q"] == 2.0e-6

    cfg.sensor.kalman_q_temp = 1.0e-9          # the retune
    fresh = SignalTracker(cfg)
    fresh.load_dict(saved)
    assert fresh.filters["temperature"].q == 1.0e-9, \
        "the state file overrode the configured tuning"
    # the estimate itself must still be carried across
    assert fresh.filters["temperature"].initialised
    assert fresh.filters["temperature"].x[0] == pytest.approx(
        old.filters["temperature"].x[0])