#!/usr/bin/env python3
"""SAFETY-94: the REAL Controller on a lagging battery whose reading is late.

    python3 tools/safety94_rig.py          # prints every scenario's counts

Each battery ramps at 800 W/s, then follows a first-order lag (`tau`, 10 s
as on the SAFETY-85 rig). Its reading (register 0x0518) is published every
10 s, at a different phase per unit, and repeats its value between publishes,
as the live boards do. B reads 100 W below what it gives (its standby draw:
-100 W at idle, live 2026-10-03). The meter updates every `tick` s; the add-on
reads every 1 s and cycles on a new meter value, as `run_control()` does.

A frozen register stops publishing at `freeze_at`; a refusing battery (`cut`)
gives 0 W and reads -4 W plus its offset (B: -104 W), with `noise` W of
seeded jitter (SAFETY-97). Uses only the Controller API release/1.0 02ed120
has, so the same file runs on the tree before SAFETY-94.
"""

import asyncio
import math
import os
import random
import sys
from datetime import datetime, timedelta, timezone

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import app.main as m                                      # noqa: E402
from app.main import Controller                           # noqa: E402

BASE = datetime(2026, 10, 3, 10, 50, tzinfo=timezone.utc)   # drain from 600 s
PUBLISH_S = 10.0
PHASE = {"a": 6.0, "b": 1.6}       # live: A publishes at :x6, B at :x1.6
OFFSET = {"a": 0.0, "b": -100.0}   # B's standby draw, live 2026-10-03
CUT_W = -4.0
LIVE = {"gain": 0.6, "slew_w": 400, "target_grid_w": -10, "deadband_w": 15,
        "step_w": 10, "heartbeat_s": 10, "stale_input_s": 15,
        "auto_start": True, "meter_entity": "sensor.grid",
        "meter_source": "off"}
MAINT = {"maintenance_enabled": True, "maintenance_start_hour": 11,
         "maintenance_discharge_w": 2500, "maintenance_hold_min": 30}


class _Store:
    def __init__(self, data=None):
        from app.store import DEFAULTS
        self.data = dict(DEFAULTS, **(data or {}))

    def set(self, k, v):
        self.data[k] = v

    def get_time(self, k):
        v = self.data.get(k)
        return datetime.fromisoformat(v) if v else None

    def set_time(self, k, when):
        self.data[k] = when.isoformat() if when else None


class _Pub:
    enabled = False

    def publish(self, *a, **k):
        pass

    def close(self):
        pass


class Plant:
    """The batteries, their registers and the meter, in simulated seconds."""

    def __init__(self, units, load, tau, tick, soc, cut, freeze_at,
                 publish_s=PUBLISH_S, dither=False, noise=0.0, seed=0):
        self.units, self.load, self.tau, self.tick = units, load, tau, tick
        self.noise, self.rng = noise, random.Random(seed)
        self.publish_s, self.dither, self.n_pub = publish_s, dither, 0
        self.soc_fn, self.cut, self.freeze_at = soc, cut, freeze_at
        self.cmd = {u: 0.0 for u in units}
        self.ramp = {u: 0.0 for u in units}
        self.real = {u: 0.0 for u in units}
        self.reg = {u: round(OFFSET[u]) for u in units}
        self.grid_v = float(load(0.0))
        self.t = 0.0

    def advance(self, t):
        """Integrate to `t` in 0.1 s steps; publish and sample on schedule."""
        while self.t < t - 1e-9:
            self.t = round(self.t + 0.1, 6)
            for u in self.units:
                want = 0.0 if self.cut(u, self.t) else self.cmd[u]
                d = want - self.ramp[u]
                self.ramp[u] += math.copysign(min(abs(d), 80.0), d)
                self.real[u] += ((self.ramp[u] - self.real[u])
                                 * (1 - math.exp(-0.1 / self.tau)))
                frozen = (self.freeze_at is not None
                          and self.freeze_at(u) is not None
                          and self.t >= self.freeze_at(u))
                ph = (self.t - PHASE[u] % self.publish_s) / self.publish_s
                if not frozen and abs(ph - round(ph)) < 1e-6:
                    # `dither`: the +-1 W a settled register shows (live).
                    self.n_pub += 1
                    jit = (self.n_pub * 7 % 3 - 1) if self.dither else 0
                    self.reg[u] = (
                        float(round(CUT_W + OFFSET[u] + self.rng.uniform(
                            -self.noise, self.noise)))
                        if self.cut(u, self.t)
                        else float(round(self.real[u] + OFFSET[u]) + jit))
            if abs(self.t / self.tick - round(self.t / self.tick)) < 1e-6:
                self.grid_v = float(round(self.load(self.t)
                                          - sum(self.real.values())))


class _Hass:
    def __init__(self, plant):
        self.p = plant

    async def limits(self, e):
        return (-5000.0, 5000.0, 1.0)

    async def call(self, *a):
        return True

    async def number(self, e, invert=False):
        if e == "sensor.grid":
            return self.p.grid_v
        u = e[-1]
        if e.startswith("sensor.soc"):
            return self.p.soc_fn(u, self.p.t)
        if e.startswith("sensor.batt"):
            return self.p.reg[u]
        return None

    async def set_number(self, e, v):
        self.p.cmd[e[-1]] = float(v)
        return True


def options(n, extra=None):
    o = dict(LIVE)
    if n == 1:
        o.update({"setpoint_entity": "number.sp_a",
                  "soc_entity": "sensor.soc_a", "batt_entity": "sensor.batt_a",
                  "max_w": 4000, "bms_soc_floor": 10})
    else:
        o["inverters"] = [
            {"name": f"inverter-{u}", "setpoint_entity": f"number.sp_{u}",
             "soc_entity": f"sensor.soc_{u}", "batt_entity": f"sensor.batt_{u}",
             "max_w": 2500, "bms_soc_floor": 10} for u in "ab"[:n]]
    o.update(extra or {})
    return o


def run(n, load, dur, *, tau=10.0, tick=1.0, soc=None, cut=None,
        freeze_at=None, extra=None, floor_s=15.0, publish_s=PUBLISH_S,
        dither=False, noise=0.0, seed=0, store=None):
    """Returns a dict: per-unit stale trips [(t, unit)], refusals
    [(t, unit or 'site')], rows [(t, grid, {u: cmd}, {u: real}, {u: reg},
    phases)] and the controller. `store`: saved state to start from."""
    units = tuple("ab"[:n])
    soc = soc or (lambda u, t: 50.0)
    cut = cut or (lambda u, t: False)
    plant = Plant(units, load, tau, tick, soc, cut, freeze_at, publish_s,
                  dither, noise, seed)

    async def go():
        c = Controller(options(n, extra), _Hass(plant), _Store(store), _Pub(),
                       stop=asyncio.Event())
        for i in c.inverters:
            i.dev_min, i.dev_max = -5000.0, 5000.0
        c.auto = True
        real_dt = m.datetime
        clock = [0.0]

        class _DT(real_dt):
            @classmethod
            def now(cls, tz=None):
                return BASE + timedelta(seconds=clock[0])

        m.datetime = _DT
        out = {"stale": [], "refused": [], "rows": [], "c": c}
        was = {i.name: False for i in c.inverters}
        was_ref = {i.name: 0 for i in c.inverters}
        was_site = 0
        last_grid, last_cycle = None, None
        try:
            t = 0.0
            while t < dur:
                plant.advance(t)
                clock[0] = t
                await c.read_inputs(now_mono=t)
                for i in c.inverters:
                    if i.batt_fresh.stale and not was[i.name]:
                        out["stale"].append((t, i.name))
                    was[i.name] = i.batt_fresh.stale
                cycled = (c.grid != last_grid or last_cycle is None
                          or t - last_cycle >= floor_s)
                if cycled:
                    last_grid, last_cycle = c.grid, t
                    await c.cycle()
                # A refusal can start on a read (a gate trip) or in a cycle.
                for i in c.inverters:
                    if i.refused and not was_ref[i.name]:
                        out["refused"].append((t, i.name))
                    was_ref[i.name] = i.refused
                if c.batt_refused and not was_site:
                    out["refused"].append((t, "site"))
                was_site = c.batt_refused
                if cycled:
                    out["rows"].append((
                        t, plant.grid_v, dict(plant.cmd),
                        {u: round(v) for u, v in plant.real.items()},
                        dict(plant.reg), [mm.phase for mm in c.maints]))
                t += 1.0
        finally:
            m.datetime = real_dt
        return out
    return asyncio.run(go())


# --- the scenarios ---------------------------------------------------------
def rest_then(step_at, w):
    """A house at 0 W (the batteries rest) until `step_at`, then `w`."""
    return lambda t: w if t >= step_at else 0.0


def reversal(at_s, charge_w, house_w):
    """Solar surplus (the batteries charge), then a house load from `at_s`."""
    return lambda t: house_w if t >= at_s else charge_w


def drain_soc(floor_at):
    """A drains (maintenance) and reaches its floor at `floor_at`; B at 50 %."""
    def soc(u, t):
        if u != "a":
            return 50.0
        return 50.0 if t < floor_at else 10.0
    return soc


SCENARIOS = {
    # first command after a long rest (the live 17:39 shape)
    "rest_step": dict(load=rest_then(600.0, 2400.0), dur=900.0),
    # maintenance drain starts after a rest (the live 19:00 shape)
    "drain_start": dict(load=lambda t: 300.0, dur=900.0,
                        extra=MAINT, soc=drain_soc(1e9)),
    # A's drain ends and its charge starts (the live 21:21 shape)
    "drain_charge": dict(load=lambda t: 300.0, dur=1200.0, extra=MAINT,
                         soc=drain_soc(900.0)),
    # a charge reverses into a discharge through the lag (SAFETY-94)
    "reversal": dict(load=reversal(300.0, -2500.0, 800.0), dur=600.0),
}


def counts(n, name, **kw):
    sc = dict(SCENARIOS[name], **kw)
    r = run(n, **sc)
    return len(r["stale"]), len(r["refused"]), r


if __name__ == "__main__":
    for n in (1, 2):
        for name in SCENARIOS:
            for tick in (1.0, 2.0, 5.0):
                for pub, dith in ((10.0, False), (1.0, True)):
                    s, f, r = counts(n, name, tick=tick, publish_s=pub,
                                     dither=dith)
                    print(f"N={n} {name:13s} tick {tick:.0f} s publish "
                          f"{pub:2.0f} s{' dither' if dith else '       '}: "
                          f"stale {s} {r['stale'][:2]}, refusals {f} "
                          f"{r['refused'][:2]}")