"""Runnable check for SAFETY-91: a held battery is drawn only after the others
are full. `python3 test_safety91.py`

At N>=2 a battery at or under its peak reserve was held back only while the
peak guard asked for power. The self-consumption loop drew it at its full
rating while a draining battery stayed at its drain rate (SAFETY-87 rig test
F1), and at a quarter's edge in a HOLD the held battery gave 2500 W for one
cycle (F2). Now a held battery weighs 0 whenever the site gives, a draining
battery rises above its drain rate before it, and the edge is graded on the
next quarter. SAFETY-25 option A stands: once the others are full the loop
may still draw it. Drives the REAL Controller on a lagging battery (10 s, as
on the rig). DECISIONS/SAFETY-91-held-battery-last.md.

1. B draining, A held, house over B's rating: B rises to full before A gives.
2. A held reaching its cut-out: nothing is placed on it, B gives its rating.
3. HOLD across a quarter edge: the held battery gives nothing at the edge.
4. Idle, A held: B carries the house first; both held or B refusing: as before;
   a SoC reading flipping across the reserve does not swap the split.
5. The rules at their boundaries (claim lists, N=3, a site stop).
6. next_quarter() at its boundaries.
7. N=1 is untouched.

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

import logging
import math
import os
import sys
import time
import types
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 safety84_identity as s84                                # noqa: E402
from app import peakguard                                      # noqa: E402
from app.arbiter import (Claim, P_MAINTENANCE, P_SAFETY,        # noqa: E402
                         resolve)
from app.main import Controller                                # noqa: E402
from app.maintenance import DRAIN, HOLD, IDLE, MaintResult     # noqa: E402
from app.p1 import make_sample, split_signed                   # noqa: E402
from app.store import unit_key                                 # noqa: E402
from app.tariff import QUARTER_S, Projection, next_quarter     # noqa: E402

fails = []
A, B = "inverter-a", "inverter-b"
FULL = 2500.0
CUT = 10.0              # bms_soc_floor of both test batteries


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


def lagging(fn, lag=10.0, tick=1.0, lags=None):
    """`fn` on plantsim's battery (800 W/s ramp, `lag` s first-order lag, or
    `lags[u]`), a 1 s meter tick, and a BMS that stops discharge at its
    cut-out."""
    real_p, real_t = rig._Hass.power, rig.TICK

    def power(self, u):
        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 / (lags or {}).get(u, lag)))
        st[u] = (a, tgt, self.t)
        return 0.0 if self.soc(u, self.t) <= CUT and a > 0 else float(round(a))
    rig._Hass.power, rig.TICK = power, tick
    try:
        return fn()
    finally:
        rig._Hass.power, rig.TICK = real_p, real_t


def store(phase_of, n=2):
    """`phase_of`: {name: phase}; the others have just done their cycle."""
    s = rig._Store()
    for u in (A, B)[:n]:
        key = (lambda k: k) if n == 1 else (lambda k: unit_key(u, k))
        if u in phase_of:
            s.data[key("phase")] = phase_of[u]
            s.data[key("phase_started")] = rig.BASE.isoformat()
        else:
            s.data[key("last_completed")] = rig.BASE.isoformat()
    return s


def run(house, phase_of, socs, dur, late=0.0, kwh=15.36, n=2, read=None,
        lags=None):
    """Rows and controller. Levels follow each battery's measured power; the
    P1 telegram read at each cycle is `late` s old. `read(u, t, level)` is
    the SoC the controller reads (default the level)."""
    soc, last, levels = dict(socs), [0.0], {}
    grid = []

    def before(c, t):
        grid.append((t, c.hass.grid()))
        for u in soc:
            soc[u] = max(CUT, soc[u] - c.hass.power(u) * (t - last[0])
                         / 3600.0 / (kwh * 10.0))
        last[0], levels[t] = t, dict(soc)
        imp, exp = split_signed(c.hass.grid())
        c.p1.submit(make_sample("ha_signed", imp, exp, phases=1,
                                ingest_ts=rig.BASE + timedelta(seconds=t - late),
                                ingest_mono=time.monotonic()))
    rows, c = lagging(lambda: s84.run(
        n, load=house if callable(house) else (lambda t: house), cheap=False,
        store=store(phase_of, n), before=before, dur=dur,
        soc=(lambda u, t: soc[u]) if read is None
        else (lambda u, t: read(u, t, soc[u])),
        extra={"maintenance_discharge_w": 1000,
               "maintenance_hold_min": 600}), lags=lags)
    c.levels, c.grid = levels, grid
    return rows, c


def without_fix(fn):
    """release/1.0 8c438ab: held only for the peak, a drain stays pinned, the
    quarter's edge is not rolled over."""
    real = (Controller._held_back, Controller._drain_rise)
    import app.main as m
    real_nq = m.next_quarter
    Controller._held_back = lambda self: self._peak_reserve()
    Controller._drain_rise = lambda self, claims, results: None
    m.next_quarter = lambda p, max_hold_s: p
    try:
        return fn()
    finally:
        Controller._held_back, Controller._drain_rise = real
        m.next_quarter = real_nq


def cmd(r):
    return dict(r[2])


def seg(rows, lo, hi):
    """Mean A and B commands over [lo, hi) s."""
    rs = [r for r in rows if lo <= r[0] < hi]
    return (round(sum(cmd(r)["a"] for r in rs) / len(rs)),
            round(sum(cmd(r)["b"] for r in rs) / len(rs)))


def a_before_b(rows, t0=0.0):
    """Cycles where held A is asked to give while B is not at full power."""
    return sum(1 for r in rows if r[0] >= t0
               and cmd(r)["a"] > 0.0 and cmd(r)["b"] < FULL)


def unmet(rows, lo, hi):
    return sum(1 for r in rows if lo <= r[0] < hi and r[7] and "UNMET" in r[7])


# ============================================================================ #
print("1. B draining at 1000 W, A held at 15 %, house over B's rating")
STEPS = (800.0, 4000.0, 5500.0)


def steps(t):
    return STEPS[min(int(t // 300), 2)]


DR = run(steps, {B: DRAIN}, {"a": 15.0, "b": 60.0}, 900.0)
DR_OLD = without_fix(lambda: run(steps, {B: DRAIN}, {"a": 15.0, "b": 60.0},
                                 900.0))
for tag, (rows, c) in (("after", DR), ("before", DR_OLD)):
    print(f"    {tag}: house 800 A/B {seg(rows, 60, 300)}, 4000 "
          f"{seg(rows, 360, 600)}, 5500 {seg(rows, 660, 900)}; A asked while "
          f"B not full: {a_before_b(rows)} cycles; UNMET "
          f"{unmet(rows, 0, 900)}")
check("⭐ A (held) is never asked to give while B has room: every cycle that "
      "puts power on A has B at its full 2500 W",
      a_before_b(DR[0]) == 0 and DR[1].maints[1].phase == DRAIN)
check("...house 4000: B rises from its 1000 W drain rate to 2500 W and A "
      "gives only the 1500 W rest (before: A 2500, B 1000)",
      seg(DR[0], 360, 600)[1] >= 2480 and seg(DR[0], 360, 600)[0] <= 1520
      and seg(DR_OLD[0], 360, 600)[0] >= 2480
      and seg(DR_OLD[0], 360, 600)[1] <= 1020)
check("...house 800: B still drains at its 1000 W rate (the rest is "
      "exported, as before) and A gives nothing",
      all(cmd(r)["b"] >= 1000.0 and cmd(r)["a"] == 0.0
          for r in DR[0] if 60 <= r[0] < 300))
check("...house 5500 (over both ratings): option A stands, A is drawn as a "
      "last resort once B is full, and nothing is left UNMET",
      seg(DR[0], 660, 900) == (2500, 2500) and unmet(DR[0], 0, 900) == 0)
check("non-vacuous: before the fix A was asked while B had room",
      a_before_b(DR_OLD[0]) >= 100)

# ============================================================================ #
print("\n2. the held battery reaches its cut-out (house 5500)")
CO = run(5500.0, {B: DRAIN}, {"a": 11.0, "b": 60.0}, 1200.0)
CO_OLD = without_fix(lambda: run(5500.0, {B: DRAIN}, {"a": 11.0, "b": 60.0},
                                 1200.0))


def after_cut(rows, c):
    """Rows once A reads its cut-out (its level follows its power)."""
    return [r for r in rows if c.levels[r[0]]["a"] <= CUT]


_ac, _ac_old = after_cut(*CO), after_cut(*CO_OLD)
print(f"    after A's cut-out: B {seg(_ac, 0, 9e9)[1]} W over {len(_ac)} "
      f"cycles (before: {seg(_ac_old, 0, 9e9)[1]} W); A soc "
      f"{CO[1].inverters[0].soc:.2f} %")
check("⭐ at its cut-out nothing is placed on A (0 W on every cycle), B "
      "gives its full 2500 W and nothing is left UNMET on the empty A",
      CO[1].inverters[0].soc <= CUT and len(_ac) >= 100
      and all(cmd(r)["a"] == 0.0 and cmd(r)["b"] == FULL for r in _ac)
      and unmet(CO[0], 600, 1200) == 0)
check("non-vacuous: before the fix B stayed pinned under its rating on most "
      "of those cycles (the rig saw ~2422 W)",
      len(_ac_old) >= 100
      and sum(1 for r in _ac_old if cmd(r)["b"] < FULL) > len(_ac_old) / 2)

# ============================================================================ #
print("\n3. HOLD across a quarter edge (A full holding, B held at 15 %)")
EDGE = run(5500.0, {A: HOLD}, {"a": 99.0, "b": 15.0}, 1000.0, late=0.5)
EDGE_OLD = without_fix(lambda: run(5500.0, {A: HOLD}, {"a": 99.0, "b": 15.0},
                                   1000.0, late=0.5))


def at_edge(rows):
    return [(r[0], cmd(r)["a"], cmd(r)["b"]) for r in rows
            if QUARTER_S - 1 <= r[0] <= QUARTER_S + 1]


print(f"    edge after {at_edge(EDGE[0])}, before {at_edge(EDGE_OLD[0])}")
check("⭐ the telegram at the edge comes late (0.5 s): B gives nothing on any "
      "cycle (the first cycle, before any quarter is measured, aside), A "
      "keeps giving and stays in its HOLD",
      all(cmd(r)["b"] == 0.0 for r in EDGE[0] if r[0] > 0)
      and all(a > 2000.0 for _t, a, _b in at_edge(EDGE[0]))
      and EDGE[1].maints[0].phase == HOLD)
check("non-vacuous: before the fix B gave over 2000 W at the edge (F2)",
      any(b > 2000.0 for _t, _a, b in at_edge(EDGE_OLD[0])))

# ============================================================================ #
print("\n4. idle, house 2000 W")
ID = run(2000.0, {}, {"a": 15.0, "b": 60.0}, 300.0)
ID_OLD = without_fix(lambda: run(2000.0, {}, {"a": 15.0, "b": 60.0}, 300.0))
print(f"    A held: after {seg(ID[0], 120, 300)}, before "
      f"{seg(ID_OLD[0], 120, 300)}")
check("⭐ A held, B able: B carries the house alone; A is asked only on the "
      "cycles where the lagging B is already at its full command",
      a_before_b(ID[0]) == 0 and seg(ID[0], 120, 300)[0] == 0
      and seg(ID[0], 120, 300)[1] >= 1950)
check("non-vacuous: before the fix the weights gave A a share",
      seg(ID_OLD[0], 120, 300)[0] > 50)
BOTH = run(2000.0, {}, {"a": 15.0, "b": 14.0}, 300.0)
BOTH_OLD = without_fix(lambda: run(2000.0, {}, {"a": 15.0, "b": 14.0}, 300.0))
check("...both held: nothing to give first, the split is as before",
      [r[:3] for r in BOTH[0]] == [r[:3] for r in BOTH_OLD[0]])


def held_back(socs, refused=(0, 0), live=(A, B)):
    c = ID[1]
    for inv, s, r in zip(c.inverters, socs, refused):
        inv.soc, inv.refused = s, r
    c.live_units = lambda: live
    try:
        return c._held_back()
    finally:
        del c.live_units
        for inv in c.inverters:
            inv.refused = 0


check("...held back only while another can give: A held beside B able; "
      "nothing with B refusing to give, B dropped, or both under",
      held_back((15.0, 60.0)) == {A} and held_back((15.0, 60.0), (0, 1))
      == frozenset() and held_back((15.0, 60.0), live=(A,)) == frozenset()
      and held_back((15.0, 14.0)) == frozenset()
      and held_back((15.0, 60.0), (0, -1)) == {A})

# A's SoC READING flips 20/21 % across its 20 % reserve every 10 s, and B
# follows its command faster (3 s) than A (10 s), as two different inverters.
def _flip(u, t, level):
    return (20.0 + int(t // 10) % 2) if u == "a" else level


def _a_jumps(rows, t0):
    rs = [r for r in rows if r[0] >= t0]
    return sum(1 for p, q in zip(rs, rs[1:])
               if abs(cmd(p)["a"] - cmd(q)["a"]) > 300.0)


def _grid_rms(c, t0):
    g = [w + 10.0 for t, w in c.grid if t >= t0]
    return math.sqrt(sum(w * w for w in g) / len(g))


_lags = {"a": 10.0, "b": 3.0}
FL = run(2000.0, {}, {"a": 20.0, "b": 60.0}, 600.0, read=_flip, lags=_lags)
_real_was = Controller.__dict__["_held_back"]


def _no_hyst(self):
    self._held_was = frozenset()
    return _real_was(self)


Controller._held_back = _no_hyst
try:
    FL_OLD = run(2000.0, {}, {"a": 20.0, "b": 60.0}, 600.0, read=_flip,
                 lags=_lags)
finally:
    Controller._held_back = _real_was
print(f"    A's reading flipping 20/21 %: A jumps {_a_jumps(FL[0], 120)}, "
      f"grid rms {_grid_rms(FL[1], 120):.0f} W (no hysteresis: "
      f"{_a_jumps(FL_OLD[0], 120)}, {_grid_rms(FL_OLD[1], 120):.0f} W)")
check("⭐ a SoC reading flipping across the reserve does not swap the split: "
      "A stays held back (0 W) and the grid stays near its target",
      _a_jumps(FL[0], 120) == 0 and seg(FL[0], 120, 600)[0] == 0
      and _grid_rms(FL[1], 120) < 30.0)
check("non-vacuous: without the 1-point hysteresis the split swaps on "
      "every flip and the grid swings",
      _a_jumps(FL_OLD[0], 120) >= 20 and _grid_rms(FL_OLD[1], 120) > 100.0)
REL = run(2000.0, {}, {"a": 20.0, "b": 60.0}, 300.0,
          read=lambda u, t, lv: (20.0 if t < 100 else 21.0 if t < 200
                                 else 22.0) if u == "a" else lv)
check("...held back at 20 %, still at 21 %, released at 22 % (more than 1 "
      "point above its reserve)",
      seg(REL[0], 40, 100)[0] == 0 and seg(REL[0], 140, 200)[0] == 0
      and seg(REL[0], 240, 300)[0] > 500)
ID[1]._held_was = frozenset({A, B})
check("...a battery kept held back by the hysteresis is not one that can "
      "give: A at 20 % and B kept at 21 % hold nobody back",
      held_back((20.0, 21.0)) == frozenset())

# ============================================================================ #
print("\n5. the rules at their boundaries")


def stub(names, held, live=None):
    """A Controller stand-in for `_drain_rise()`, on SIMULATED units."""
    s = types.SimpleNamespace(
        inverters=[types.SimpleNamespace(name=n) for n in names],
        actuator_of=lambda src, inv=None: inv.name,
        live_units=lambda: tuple(live or names),
        _held_back=lambda: frozenset(held))
    return s


def rise(names, held, target, drains, extra=(), live=None):
    """`_drain_rise()` then the split: (claims kinds, shares)."""
    cl = [Claim.limit("supervised", P_SAFETY, -FULL, FULL, actuator=n)
          for n in names]
    cl += [Claim.set("maintenance", P_MAINTENANCE, w, DRAIN, actuator=n)
           for n, w in drains.items()]
    cl += list(extra)
    res = [MaintResult(DRAIN, True, drains[n], False, []) if n in drains
           else MaintResult(IDLE, False, None, False, []) for n in names]
    Controller._drain_rise(stub(names, held, live), cl, res)
    r = resolve(cl + [Claim.set("loop", 10, target)], units=tuple(names),
                live=set(live or names), reserve=held)
    kinds = {c.actuator: c.kind for c in cl if c.source == "maintenance"}
    return kinds, {u: round(v) for u, v in r.shares.items()}


_n3 = rise(("x", "y", "z"), {"z"}, 1500.0, {"x": 1000.0})
print(f"    N=3, x draining 1000, y able, z held, 1500 W: {_n3[1]}")
check("⭐ N=3: the draining battery keeps its drain rate as a floor (x >= "
      "1000 W) while sharing with the able one; the held one gives nothing",
      _n3[0] == {"x": "limit"} and _n3[1]["x"] >= 1000 and _n3[1]["z"] == 0
      and sum(_n3[1].values()) == 1500)
_stop = rise((A, B), {A}, 3000.0, {B: 1000.0},
             (Claim.set("safety", P_SAFETY, 0.0, "stopped"),))
check("⭐ a site stop decides alone: the drain stays a pin, nothing rises",
      _stop[0] == {B: "set"} and _stop[1] == {A: 0, B: 0})
_hd = rise((A, B), {A}, 5000.0, {A: 1000.0})
check("...a held battery's own drain stays a pin (it is not raised for the "
      "loop): A 1000 W, B 2500 W",
      _hd[0] == {A: "set"} and _hd[1] == {A: 1000, B: 2500})
_nh = rise((A, B), set(), 3000.0, {B: 1000.0})
check("...with nothing held the drain stays a pin, as before (SAFETY-78)",
      _nh[0] == {B: "set"} and _nh[1] == {A: 2000, B: 1000})
_dead = rise((A, B), {A}, 3000.0, {B: 1000.0}, live=(A,))
check("...a draining battery that is not live is not touched",
      _dead[0] == {B: "set"})

# ============================================================================ #
print("\n6. next_quarter()")
T0 = rig.BASE


def proj(elapsed, covered=1.0, held=3000.0, start=T0):
    return Projection(projected_w=2000.0, elapsed_w=1900.0, elapsed_s=elapsed,
                      covered_frac=covered, complete=covered >= 1.0,
                      block_start=start, held_w=held)


_in = proj(899.0)
_past = next_quarter(proj(900.4, covered=899.9 / 900.4), 30.0)
check("⭐ inside its block a projection is unchanged (the same object)",
      next_quarter(_in, 30.0) is _in)
check("⭐ past its end with a fresh held reading: the next quarter, "
      "projected from that reading",
      _past.block_start == T0 + timedelta(seconds=QUARTER_S)
      and abs(_past.elapsed_s - 0.4) < 1e-9 and _past.projected_w == 3000.0
      and _past.covered_frac == 1.0
      and peakguard.blindness(_past, 30.0, False) is None)
_old = proj(960.0, covered=880.0 / 960.0)
_edge29 = proj(901.0, covered=871.5 / 901.0)     # 29.5 s uncovered
_edge31 = proj(901.0, covered=870.5 / 901.0)     # 30.5 s uncovered
check("...the held reading is credited up to max_hold_s and no further "
      "(29.5 s uncovered: rolled; 30.5 s: unchanged)",
      next_quarter(_edge29, 30.0) is not _edge29
      and next_quarter(_edge31, 30.0) is _edge31)
check("...a held reading too old to be credited, no block, or an export: "
      "unchanged, or projected at 0 W offtake",
      next_quarter(_old, 30.0) is _old
      and next_quarter(proj(950.0, start=None), 30.0).block_start is None
      and next_quarter(proj(900.5, held=-800.0), 30.0).projected_w == 0.0)

# ============================================================================ #
print("\n7. N=1")
for _name, _args in (("HOLD, late telegram at the edge",
                      (5500.0, {A: HOLD}, {"a": 99.0}, 1000.0, 0.5)),
                     ("drain under the reserve, house 4000",
                      (4000.0, {A: DRAIN}, {"a": 15.0}, 600.0, 0.5))):
    _new = run(*_args, n=1)
    _old = without_fix(lambda: run(*_args, n=1))
    print(f"    {_name}: A at the edge "
          f"{[cmd(r)['a'] for r in _new[0] if 899 <= r[0] <= 901]}")
    check(f"⭐ N=1, {_name}: every row (claims, commands, the guard) is "
          "identical with and without the fix",
          _new[1].maints[0].phase == _args[1][A] and len(_new[0]) > 500
          and _new[0] == _old[0])
_one = types.SimpleNamespace(inverters=[object()])
_cl = [Claim.set("maintenance", P_MAINTENANCE, 1000.0, DRAIN)]
check("⭐ at one inverter nothing is held back and no pin moves (the N=1 "
      "identity digests in test_safety78/79/83/84/86/87 stay pinned)",
      Controller._held_back(_one) == frozenset()
      and Controller._drain_rise(types.SimpleNamespace(
          inverters=[object()], _held_back=lambda: frozenset()), _cl,
          [MaintResult(DRAIN, True, 1000.0, False, [])]) is None
      and _cl[0].kind == "set")

print()
if fails:
    print(f"{len(fails)} FAILED: " + ", ".join(fails))
    sys.exit(1)
print("all SAFETY-91 checks passed")