"""Runnable check for SAFETY-119: no "Battery is not keeping up" when the
loop takes the wire back from another owner. `python3 test_safety119.py`

During a maintenance charge with no usable peak forecast, the measured peak
guard steps in during a quarter's last seconds and hands back. On the cycle
the loop gets the wire back, its integrator is re-seeded from the other
owner's command and kicked by a meter reading taken before the battery moved.
It crossed its rail for one or two cycles, the loop's ask beat the one-sided
limit, so the cycle counted as the loop's own and raised
`goodwe_integrator_bound`. The battery gave what it was told. Now the loop's
bound does not count until the plant's settle (SHORT_CYCLES cycles and
SHORT_SPAN_S) has passed since the hand-back.
DECISIONS/SAFETY-119-bound-after-hand-back.md.

1. The rig's shape (1 s meter that lags the battery, 1 s battery lag, house
   800 W, rating 2500 W, limit 2500 W, margin 50, charge 2000 W, no
   forecast), 2 h, three seeds: no notification, the cap held, every command
   the same as without the fix; without the fix it warns.
2. A genuine integrator bound still warns: the loop in charge with the house
   over the rating (at once, as before), and the loop at its rail after a
   hand-back (after the settle, not lost).
3. With a usable forecast: the same commands and the same alarms as without
   the fix.
4. The rule at its boundaries (`Controller._settling`).

Mutations (MUTATION.md, purged copies) are listed in the pull request.
"""

import logging
import math
import os
import random
import sys
import time
from datetime import timedelta

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
sys.path.insert(0, os.path.join(HERE, "tools"))
logging.disable(logging.WARNING)

import opt11_identity as rig                                   # noqa: E402
import safety79_identity as s79                                # noqa: E402
import safety84_identity as s84                                # noqa: E402
import app.main as app_main                                    # noqa: E402
from app import messages                                       # noqa: E402
from app.main import Controller                                # noqa: E402
from app.p1 import make_sample, split_signed                   # noqa: E402

fails = []
OPT0 = dict(s84.OPT)
CAP = 2500.0
QUARTER = 900.0


def check(name, cond):
    print(f"  {'ok  ' if cond else 'FAIL'} {name}")
    if not cond:
        fails.append(name)


async def _call(self, *a):
    """The rig's Home Assistant, recording each notification (time, id)."""
    if len(a) >= 3 and a[1] == "create" and isinstance(a[2], dict):
        self.__dict__.setdefault("notes", []).append(
            (self.t, a[2].get("notification_id")))
    return True


def _power(self, u, lag=1.0):
    """The TEST-23 plant's battery: 800 W/s ramp, 1 s first-order lag."""
    st = self.__dict__.setdefault("_plant", {})
    a, tgt, tl = st.get(u, (0.0, 0.0, self.t))
    for _ in range(int(round((self.t - tl) / 0.1))):
        d = self.cmd[u] - tgt
        tgt += math.copysign(min(abs(d), 80.0), d)
        a += (tgt - a) * (1 - math.exp(-0.1 / lag))
    st[u] = (a, tgt, self.t)
    if self.cut(u, self.t) and self.cmd[u] >= -4.0:
        return -4.0
    return float(round(a))


_real_number = rig._Hass.number


async def _number(self, e, invert=False):
    """sensor.fc: a usable demand forecast of 800 W (part 3)."""
    if e == "sensor.fc":
        return 800.0
    return await _real_number(self, e, invert)


rig._Hass.call, rig._Hass.power, rig._Hass.number = _call, _power, _number
rig.TICK = 1.0
DELAY = [1]


def _feed(c, t):
    """The P1 meter `DELAY` ticks behind the battery, as on the rig (the
    meter reached HA after the battery reading). Records the true grid."""
    h = c.hass.__dict__.setdefault("grid_true", [])
    h.append((t, c.hass.grid()))
    g = h[max(0, len(h) - 1 - DELAY[0])][1]
    imp, exp = split_signed(g)
    c.p1.submit(make_sample("ha_signed", imp, exp, phases=1,
                            ingest_ts=rig.BASE + timedelta(seconds=t),
                            ingest_mono=time.monotonic()))


s79._feed = _feed


def house(seed, base=800.0, step=None):
    """800 W +- 30 W, a new value each second; `step` = (t, W) from then."""
    rnd = random.Random(seed)
    noise = [rnd.uniform(-30.0, 30.0) for _ in range(4000)]

    def load(t):
        if step is not None and t >= step[0]:
            return step[1]
        return base + noise[int(t) % 4000]
    return load


def run(load, *, dur=7200.0, delay=1, extra=None, cheap=True, cut=None,
        charging=True):
    """One N=1 run; `charging=False`: no maintenance, the loop in charge."""
    extra = dict(extra or {})
    s84.OPT = dict(OPT0, peak_cap_w=int(CAP), peak_margin_w=50,
                   maintenance_charge_w=2000)
    if not charging:
        extra.update(maintenance_enabled=False, peak_cap_w=20000)
    for k in [k for k in extra if k in s84.OPT]:
        s84.OPT[k] = extra.pop(k)
    DELAY[0] = delay
    try:
        return s84.run(1, load=load, cheap=cheap,
                       store=(s79.store_charging(1) if charging
                              else rig._Store()),
                       soc=lambda u, t: 60.0, dur=dur, cut=cut,
                       extra=dict({"max_w": 2500, "slew_w": 0}, **extra))
    finally:
        s84.OPT = OPT0


def patched(obj, name, value):
    def wrap(fn):
        real = getattr(obj, name)
        setattr(obj, name, value)
        try:
            return fn()
        finally:
            setattr(obj, name, real)
    return wrap


no_fix = patched(Controller, "_settling", lambda self, owns, now: False)


def notes(c):
    return [t for t, i in getattr(c.hass, "notes", [])
            if i == messages.ID_LOOP_BOUND]


def quarters(c):
    """Each whole quarter's true average import (W)."""
    q = {}
    for t, g in c.hass.grid_true:
        q.setdefault(int(t // QUARTER), []).append(g)
    return [sum(v) / len(v) for v in q.values() if len(v) == QUARTER]


def commands(rows):
    return [(r[0], r[1], r[2]) for r in rows]


def guard_steps(c):
    return sum(1 for e in c.events if "PEAK_GUARD:" in e)


# ============================================================================ #
print("1. a maintenance charge with no forecast, the guard steps in at each "
      "quarter's end, 2 h")
NEW, OLD = {}, {}
for seed, delay in ((1, 1), (2, 1), (3, 2)):
    NEW[seed] = run(house(seed), delay=delay)
    OLD[seed] = no_fix(lambda: run(house(seed), delay=delay))
    print(f"    seed {seed} meter {delay} s behind: notifications "
          f"{notes(NEW[seed][1])} (without the fix {notes(OLD[seed][1])}), "
          f"highest quarter {max(quarters(NEW[seed][1])):.0f} W, guard "
          f"step-ins in the last 40 log rows {guard_steps(NEW[seed][1])}")
check("⭐ no 'Battery is not keeping up' notification, cycle or episode",
      all(notes(c) == [] and c.clamp_wd.bound == 0
          and c.clamp_wd.bound_episodes == 0 for _r, c in NEW.values()))
check("⭐ the peak cap is held: every quarter's average is under 2500 W",
      all(len(quarters(c)) == 8 and max(quarters(c)) < CAP
          for _r, c in NEW.values()))
check("⭐ control is unchanged: every command the same as without the fix",
      all(commands(NEW[k][0]) == commands(OLD[k][0]) for k in NEW))
check("...the charge ran and the guard stepped in near each quarter's end",
      all(any(r[1] == -2000.0 for r in rows) and any(
          r[-1][1] and r[-1][1][0] and r[0] % QUARTER >= QUARTER - 40
          for r in rows) for rows, _c in NEW.values()))
check("non-vacuous: without the fix every run warns, at a quarter's edge",
      all(len(notes(c)) >= 1 for _r, c in OLD.values())
      and all(min(t % QUARTER, QUARTER - t % QUARTER) <= 40
              for _r, c in OLD.values() for t in notes(c)))

# ============================================================================ #
print("\n2. a genuine integrator bound still warns")
_rail = run(house(1, base=4000.0), dur=600.0,
            charging=False)
_rail_old = no_fix(lambda: run(house(1, base=4000.0), dur=600.0,
                                 charging=False))
print(f"    loop in charge, house 4000 W over the 2500 W rating: "
      f"{notes(_rail[1])}, without the fix {notes(_rail_old[1])}")
check("⭐ the loop in charge at its rail warns (log and notification), at "
      "the same cycle as before",
      len(notes(_rail[1])) >= 1 and notes(_rail[1]) == notes(_rail_old[1])
      and _rail[1].clamp_wd.bound == _rail_old[1].clamp_wd.bound
      and any("OUTPUT_CLAMP episode" in e for e in _rail[1].events))

# The charge sets the wire until 600 s, then a solar surplus of 6000 W: the
# loop takes the wire back and stands at its -2500 W rail.
_hb = run(house(1, step=(600.0, -6000.0)), dur=1200.0,
          extra={"peak_cap_w": 20000})
_hb_old = no_fix(lambda: run(house(1, step=(600.0, -6000.0)), dur=1200.0,
                             extra={"peak_cap_w": 20000}))
_first = notes(_hb[1])[:1]
print(f"    hand-back to a loop at its rail at 600 s: first notification "
      f"{_first}, without the fix {notes(_hb_old[1])[:1]}; command at the "
      f"end {_hb[0][-1][1]}")
check("⭐ the loop at its rail after a hand-back still warns, once the "
      "settle has passed (12 cycles, 22 s): delay, not loss",
      _first and 600.0 + app_main.SHORT_SPAN_S <= _first[0] <= 640.0
      and _hb[0][-1][1] == -2500.0
      and _hb[1].clamp_wd.bound > 100)
check("...without the fix it warned at once (non-vacuous)",
      notes(_hb_old[1])[:1] and notes(_hb_old[1])[0] < 600.0
      + app_main.SHORT_SPAN_S)
# The same charge with no forecast; at 1795 s, under the guard's step-in at
# the quarter's end, the house goes to 4000 W, over the 2500 W rating.
_up = run(house(1, step=(1795.0, 4000.0)), dur=1900.0)
_up_old = no_fix(lambda: run(house(1, step=(1795.0, 4000.0)), dur=1900.0))
print(f"    guard step-in, then the house over the rating from 1795 s: "
      f"{notes(_up[1])}, without the fix {notes(_up_old[1])}")
check("⭐ a genuine bound in the rig's own scenario still warns, within the "
      "settle of the guard's hand-back",
      notes(_up[1]) and 1795.0 <= notes(_up[1])[0] <= 1795.0 + 40.0
      and any("OUTPUT_CLAMP episode" in e for e in _up[1].events))

# ============================================================================ #
print("\n3. with a usable forecast (800 W): unchanged")
FC = {"peak_forecast_entity": "sensor.fc"}
_fc = run(house(1), extra=FC)
_fc_old = no_fix(lambda: run(house(1), extra=FC))
print(f"    notifications {notes(_fc[1])} (without the fix "
      f"{notes(_fc_old[1])}), highest quarter "
      f"{max(quarters(_fc[1])):.0f} W, charge "
      f"{min(r[1] for r in _fc[0][:300]):.0f} W")
check("⭐ the same commands and the same alarms as without the fix, and "
      "the cap held",
      commands(_fc[0]) == commands(_fc_old[0])
      and notes(_fc[1]) == notes(_fc_old[1])
      and _fc[1].clamp_wd.engagements == _fc_old[1].clamp_wd.engagements
      and max(quarters(_fc[1])) < CAP)
check("...the forecast was used (the charge is shaped under -2000 W)",
      -2000.0 < min(r[1] for r in _fc[0][:300]) < 0.0)

# ============================================================================ #
print("\n4. the rule at its boundaries")


class _C:
    _settling = Controller._settling

    def __init__(self):
        self._wire_lost, self._handback = False, None


def seq(c, steps):
    return [c._settling(o, t) for o, t in steps]


N, SPAN = app_main.SHORT_CYCLES, app_main.SHORT_SPAN_S
_c = _C()
check("from the start, owned every cycle: never settling (no hand-back)",
      seq(_c, [(True, float(i)) for i in range(30)]) == [False] * 30)
_c = _C()
_out = seq(_c, [(False, 0.0)] + [(True, 2.0 * i) for i in range(1, N + 3)])
check(f"after another owner: settling for {N - 1} owned cycles, then not "
      f"(12 cycles at 2 s span {SPAN:.0f} s)",
      _out == [False] + [True] * (N - 1) + [False] * 3)
_c = _C()
_out = seq(_c, [(False, 0.0)] + [(True, 0.5 * i) for i in range(1, 60)])
_end = _out.index(False, 1)
check(f"at a 0.5 s meter the span decides: settling until {SPAN:.0f} s",
      _end == 1 + int(SPAN / 0.5) and all(_out[1:_end]))
_c = _C()
_out = seq(_c, [(False, 0.0), (True, 2.0), (True, 4.0), (False, 6.0),
                (True, 8.0), (True, 10.0)])
check("losing the wire again restarts the window",
      _out == [False, True, True, False, True, True]
      and _c._handback == (8.0, 2))
_c = _C()
check("no monotonic clock: the cycle count alone decides",
      seq(_c, [(False, None)] + [(True, None)] * (N + 1))
      == [False] + [True] * (N - 1) + [False] * 2)

# An over-rating (a command over `max_w`) is a bug, never settled away.
_ctl = _hb[1]
_ctl.loop_clamped = (True, "x", True)
_ctl._wire_lost, _ctl._handback = True, None


class _Res:
    winner, bound_by, target_w = "loop", "", 0.0
    unit_bounds = {}


class _D:
    held, target_w = False, 0.0


_ctl._loop_decision, _ctl.allocation = _D(), None
_ctl._loop_alarm_input(_Res())
check("an over-rating on a hand-back cycle still counts (bug=True kept)",
      _ctl.loop_clamped == (True, "x", True) and _ctl._handback is not None)
_ctl.loop_clamped = (True, "y", False)
_ctl._loop_alarm_input(_Res())
check("...while the integrator bound on the next cycle is not counted",
      _ctl.loop_clamped[0] is False)

_ctl._wire_lost, _ctl._handback = True, (1.0, 3)
_ctl._loop_decision = None                    # stopped or blind
_ctl._loop_alarm_input(_Res())
check("a stop clears it: the next start is not a hand-back",
      (_ctl._wire_lost, _ctl._handback) == (False, None))

print()
if fails:
    print(f"FAILED: {len(fails)} check(s)")
    for f in fails:
        print(f"  - {f}")
    sys.exit(1)
print("all SAFETY-119 checks passed")