"""Layered zone model: fallbacks, contributions and the compiled optimizer step."""
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, physics, zone_model as zm  # noqa: E402

VENT = {'opening': 1.0, 'facing': -0.5, 'other_opening': 0.5, 'recovery': 0.0, 't_ref': 21.0,
        'wind': 4.0, 'gust': 7.0}


def trained(features, theta):
    learner = learning.FeatureLearner(features)
    learner.theta = list(theta)
    learner.count, learner.on_count, learner.off_count = 300, 100, 100
    return learner


def models():
    living = zm.ZoneModel(physics.Learner(), trained(zm.extended_features(True, ('bed',)),
                                                     (0.004, 0.12, 0.0005, 0.15, 0.01)))
    bed = zm.ZoneModel(physics.Learner(), trained(zm.extended_features(False, ('living',)),
                                                  (0.006, 0.09, 0.0, 0.01)),
                       trained(['vent_natural', 'recovery', 'vent_wind', 'vent_wind_dir', 'vent_cross'],
                               (0.01, 0.15, 0.001, 0.0006, 0.002)))
    return {'living': living, 'bed': bed}


class ZoneModelTests(unittest.TestCase):
    def test_layers_and_fallbacks(self):
        self.assertEqual(zm.extended_features(False, ()), [])
        self.assertEqual(zm.extended_features(True, ('b',)), ['loss', 'heat', 'bias', 'internal', 'couple:b'])
        living = models()['living']
        inputs = {'outside': 5.0, 'lag': 0.5, 'internal': 0.4, 'temps': {'living': 21.0, 'bed': 20.0}}
        self.assertAlmostEqual(living.delta(21.0, inputs),
                               0.004 * -16 + 0.06 + 0.0005 + 0.06 + 0.01 * -1)
        # A missing input (electrical heat or a neighbour) falls back to the base model.
        base = living.base.predict_delta(21.0, 5.0, 0.5)
        self.assertAlmostEqual(living.delta(21.0, {**inputs, 'internal': None}), base)
        self.assertAlmostEqual(living.delta(21.0, {**inputs, 'temps': {'living': 21.0}}), base)
        # Until the extended model is as trained as the base needs to be, the base is used.
        living.extended.count = 10
        self.assertFalse(living.uses_extended)
        self.assertAlmostEqual(living.delta(21.0, inputs), base)

    def test_contributions_add_up_to_the_change(self):
        bed = models()['bed']
        inputs = {'outside': 5.0, 'lag': 0.5, 'temps': {'living': 21.0, 'bed': 19.0}, 'ventilation': VENT}
        terms = bed.contributions(19.0, inputs)
        self.assertEqual(set(terms), {'loss', 'heat', 'bias', 'neighbors', 'ventilation', 'recovery'})
        self.assertAlmostEqual(sum(terms.values()), 12 * bed.delta(19.0, inputs), places=2)
        self.assertLess(terms['ventilation'], 0)
        # Unknown window state: no ventilation term rather than a guess.
        self.assertAlmostEqual(bed.ventilation_delta(19.0, {**inputs, 'ventilation': {**VENT, 'opening': None}}), 0)

    def test_compiled_step_matches_direct_model_exactly(self):
        zone_models = models()
        steps = 24
        ventilation = [{**VENT, 'opening': 1.0 if k < 3 else 0.0, 'recovery': 0.0 if k < 3 else 0.6}
                       for k in range(steps)]

        def inputs_at(key, k):
            return {'outside': 5.0, 'internal': 0.4 if key == 'living' else None,
                    'ventilation': ventilation[k] if key == 'bed' else {}}
        fast = zm.compile_step_model(zone_models, ['living', 'bed'], steps, inputs_at)
        rng = random.Random(5)
        for _ in range(300):
            state = {'living': rng.uniform(15, 26), 'bed': rng.uniform(15, 26)}
            lag, k = rng.random(), rng.randrange(steps + 3)  # beyond the horizon: last step repeats
            direct = {key: value + zone_models[key].delta(value, {
                **inputs_at(key, min(k, steps - 1)), 'lag': lag, 'temps': state}) for key, value in state.items()}
            compiled = fast(state, lag, k)
            for key in state:
                self.assertAlmostEqual(compiled[key], direct[key], places=12)
        # With a zone missing, its neighbour falls back to its base model, also when compiled.
        solo = zm.compile_step_model(zone_models, ['living'], steps, inputs_at)
        base = zone_models['living'].base.predict_delta(20.0, 5.0, 0.3)
        self.assertAlmostEqual(solo({'living': 20.0}, 0.3, 0)['living'], 20.0 + base, places=12)


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