"""Generic named-feature RLS learner, without Home Assistant."""
import random
import sys
import types
import unittest
from pathlib import Path

COMPONENTS = Path(__file__).resolve().parents[1] / 'custom_components' / 'heating_mpc'
sys.modules.setdefault('heating_mpc', types.ModuleType('heating_mpc'))
sys.modules['heating_mpc'].__path__ = [str(COMPONENTS)]
from heating_mpc import learning as lm  # noqa: E402

FEATURES = ['loss', 'heat', 'bias', 'internal']
TRUE = {'loss': 0.004, 'heat': 0.12, 'bias': 0.0005, 'internal': 0.15}


def synthetic(learner, steps=3000, seed=1, noise=0.003, gaming=True):
    """Samples from dT = sum(theta*x) + noise with realistic, varied inputs."""
    rng = random.Random(seed)
    for k in range(steps):
        x = {'loss': -20 + 20 * rng.random(), 'heat': (k // 7) % 3 / 2,
             'bias': 1.0, 'internal': 0.4 if gaming and (k // 40) % 5 == 0 else 0.0}
        y = sum(TRUE[f] * x[f] for f in TRUE) + rng.gauss(0, noise)
        learner.learn([x[f] for f in learner.features], y, command=x['heat'] > 0)
    return learner


class FeatureLearnerTests(unittest.TestCase):
    def test_recovers_coefficients_and_reports_confidence(self):
        learner = synthetic(lm.FeatureLearner(FEATURES))
        for name, value in TRUE.items():
            with self.subTest(feature=name):
                self.assertAlmostEqual(learner.coefficient(name), value, delta=0.1 * value + 0.001)
        info = learner.coefficients()
        self.assertTrue(info['internal']['identified'])
        self.assertGreater(info['internal']['confidence'], 0.8)
        # Standard error in °C/5 min per kW: small against the coefficient.
        self.assertLess(info['internal']['std'], 0.01)
        self.assertAlmostEqual(learner.noise_var ** .5, 0.003, delta=0.001)
        self.assertEqual(learner.count + learner.rejected_count, 3000)

    def test_unexercised_feature_stays_at_prior_without_confidence(self):
        learner = synthetic(lm.FeatureLearner(FEATURES), gaming=False)
        info = learner.coefficients()['internal']
        self.assertEqual((info['value'], info['samples'], info['identified']), (0.0, 0, False))
        self.assertEqual(info['confidence'], 0)
        # Windup bound: the unexcited variance never exceeds the prior.
        i = learner.features.index('internal')
        self.assertLessEqual(learner.p[i][i], lm.spec_for('internal').variance)

    def test_bounds_and_linked_directional_wind(self):
        learner = lm.FeatureLearner(['internal'])
        for _ in range(200):
            learner.learn([1.0], -0.2)  # cooling with more power: unphysical
        self.assertEqual(learner.coefficient('internal'), 0.0)
        wind = lm.FeatureLearner(['vent_wind', 'vent_wind_dir'])
        for k in range(500):
            x = [-50.0, -50.0 * (1 if k % 2 else -1)]
            wind.learn(x, 0.02 * x[1])  # only directional signal, of the "wrong" sign
        self.assertLessEqual(abs(wind.coefficient('vent_wind_dir')), wind.coefficient('vent_wind') + 1e-12)

    def test_invalid_samples_rejected(self):
        learner = lm.FeatureLearner(['loss', 'bias'])
        for x, y in (([1.0], 0.0), ([float('nan'), 1.0], 0.0), ([1.0, 1.0], 0.9),
                     ([1.0, 1.0], float('inf'))):
            self.assertFalse(learner.learn(x, y))
        self.assertEqual((learner.count, learner.rejected_count), (0, 4))
        with self.assertRaises(ValueError):
            lm.FeatureLearner(['loss', 'loss'])
        with self.assertRaises(KeyError):
            lm.FeatureLearner(['unknown'])

    def test_round_trip_and_feature_remapping(self):
        learner = synthetic(lm.FeatureLearner(FEATURES), steps=800)
        same = lm.FeatureLearner.from_dict(learner.as_dict(), FEATURES)
        self.assertEqual((same.theta, same.p, same.count, same.active, same.noise_var),
                         (learner.theta, learner.p, learner.count, learner.active, learner.noise_var))
        # Adding a neighbour keeps every learned coefficient; the new one starts at its prior.
        grown = lm.FeatureLearner.from_dict(learner.as_dict(), FEATURES + ['couple:bedroom'])
        for name in FEATURES:
            self.assertEqual(grown.coefficient(name), learner.coefficient(name))
        self.assertEqual(grown.coefficient('couple:bedroom'), 0.0)
        self.assertEqual(grown.p[4][4], lm.spec_for('couple').variance)
        self.assertEqual(grown.p[0][1], learner.p[0][1])
        # Removing the heat source keeps the rest, including their covariance.
        shrunk = lm.FeatureLearner.from_dict(learner.as_dict(), ['loss', 'heat', 'bias'])
        self.assertEqual(shrunk.theta, learner.theta[:3])
        self.assertEqual(shrunk.p[2][0], learner.p[2][0])

    def test_malformed_saved_data_starts_fresh(self):
        good = synthetic(lm.FeatureLearner(['loss', 'bias']), steps=50).as_dict()
        fresh = lm.FeatureLearner(['loss', 'bias'])
        for change in ({'features': ['loss', 'loss']}, {'theta': [0.1]}, {'p': [[1, 0], [0, -1]]},
                       {'p': [[1, 0.5], [0, 1]]}, {'count': -1}, {'active': [1]},
                       {'theta': [float('nan'), 0]}, {'on_count': 99}, {'residual_sq': -1},
                       {'residual_weight': 2}):
            with self.subTest(change=change):
                data = dict(good)
                data.update(change)
                restored = lm.FeatureLearner.from_dict(data, ['loss', 'bias'])
                self.assertEqual((restored.theta, restored.count), (fresh.theta, 0))
        for garbage in (None, [], 'x', {}):
            self.assertEqual(lm.FeatureLearner.from_dict(garbage, ['loss']).count, 0)


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