#!/usr/bin/env python3
"""OPT-11 P3: replay days of recorded house load and real prices through the
REAL Controller (read_inputs + cycle, the arbiter, the allocator, the peak
guard on a P1 meter), with two modelled batteries, and report energy, cycles
and money, plus the P3 invariants on every tick.

    python3 tools/opt11p3_replay.py --days 16 --variant p3_ls
    python3 tools/opt11p3_replay.py --captured          # 2026-09-09/10

The same file runs on release/1.0 7a00f12 (the P2 baseline) and on this
branch: it uses only the Controller API 7a00f12 has, and sets the batteries'
choices through the store keys P1 defined (`pack_sell` was reserved there).

REAL: the house (grid + battery power recorded 2026-09-02..17 on the live
install, scenarios/loop-2026-09-02_2026-09-17: the house's own net demand,
solar included), the prices (EPEX BE day-ahead, hourly, same weekday one year
earlier (364 days), or the captured quarter-hourly energy-charts bodies for
2026-09-09/10), the decision code.
MODELLED: the batteries (a first-order lag, 3 s, 800 W/s ramp, SoC
integrated, 95 % each way, a reading that jitters by a few watts, and
optionally a device limit or a refusal to charge), the all-in price (the EPEX file's own stated
assumption: (wholesale + 0.02 + 0.10) x 1.06), export paid at the wholesale
price, tomorrow's prices known from 13:00 local. The absolute money is
therefore an estimate; the comparison between variants is the point.
"""

import argparse
import asyncio
import gzip
import json
import math
import os
import sys
import time
from datetime import datetime, timedelta, timezone

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

import app.main as m                                      # noqa: E402
from app.main import Controller                           # noqa: E402
from app.p1 import make_sample, split_signed               # noqa: E402
from app.price import PricePoint, PriceProduct, PriceRead, brussels  # noqa
from app.static_tariff import TariffPlan                  # noqa: E402
from app.store import unit_key                            # noqa: E402

LOOP = os.path.join(HERE, "scenarios", "loop-2026-09-02_2026-09-17")
EPEX = os.path.join(HERE, "scenarios",
                    "epex-be-dayahead-2025-09-01_2026-09-01.json")
QH = os.path.join(HERE, "scenarios",
                  "energy-charts-be-price-2026-09-09_2026-09-10-twoday.json")
TICK = 5.0
CAP_KWH = 10.0
RATING = 2500.0
PEAK_CAP = 3500


def all_in(wholesale_kwh):
    return round((wholesale_kwh + 0.02 + 0.10) * 1.06, 6)


def load_house(days):
    """[(t_utc_epoch, house_w)] from the recording: grid + battery power."""
    out = []
    for d in days:
        def rd(kind):
            f = os.path.join(LOOP, f"{d}-{kind}.jsonl.gz")
            rows = []
            with gzip.open(f, "rt") as fh:
                for line in fh:
                    r = json.loads(line)
                    try:
                        v = float(r["v"])
                    except ValueError:
                        continue
                    rows.append((datetime.fromisoformat(r["t"]).timestamp(),
                                 v))
            return rows
        g, b = rd("grid"), rd("batt")
        j = 0
        for t, gw in g:
            while j + 1 < len(b) and b[j + 1][0] <= t:
                j += 1
            bw = b[j][1] if b and b[j][0] <= t else 0.0
            out.append((t, gw + bw))
    out.sort()
    return out


def epex_points(shift_days=364):
    d = json.load(open(EPEX))["prices"]
    pts = []
    for r in d:
        s = datetime.fromisoformat(r["start"]) + timedelta(days=shift_days)
        pts.append((s.astimezone(timezone.utc), 3600,
                    float(r["eur_per_kwh"])))
    return pts


def qh_points():
    d = json.load(open(QH))
    return [(datetime.fromtimestamp(t, timezone.utc), 900, p / 1000.0)
            for t, p in zip(d["unix_seconds"], d["price"])
            if p is not None]


class Plant:
    def __init__(self, soc, refuse_charge=False):
        self.cmd, self.p, self.soc = 0.0, 0.0, float(soc)
        self.refuse_charge = refuse_charge

    def step(self, dt, floor):
        c = self.cmd
        if c > 0 and self.soc <= floor:
            c = 0.0
        if c < 0 and (self.soc >= 100.0 or self.refuse_charge):
            c = 0.0
        want = self.p + (c - self.p) * (1 - math.exp(-dt / 3.0))
        dp = max(-800.0 * dt, min(800.0 * dt, want - self.p))
        self.p += dp
        e = self.p * dt / 3600.0 / 1000.0            # kWh, + = out
        eff = 1 / 0.95 if e > 0 else 0.95
        self.soc -= e * eff / CAP_KWH * 100.0
        self.soc = max(0.0, min(100.0, self.soc))


class Hass:
    def __init__(self, units, house_at, soc0=None, refuse=(), limits=None):
        self.units = units
        self.lim = limits or {}
        self.plants = {u: Plant((soc0 or {}).get(u, 50.0), u in refuse)
                       for u in units}
        self.house_at = house_at
        self.t = 0.0

    def grid(self):
        return self.house_at(self.t) - sum(p.p for p in self.plants.values())

    async def limits(self, e):
        lo, hi = self.lim.get(e[-1], (-5000.0, 5000.0))
        return (lo, hi, 1.0)

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

    async def number(self, e, invert=False):
        if e == "sensor.grid":
            return self.grid()
        u = e[-1]
        if e.startswith("sensor.soc"):
            return round(self.plants[u].soc, 1)
        if e.startswith("sensor.batt"):
            # A real reading jitters by a few watts; a dead-flat one is a
            # stale register to the add-on (SAFETY-37), which is not modelled.
            j = (ord(u) * 31 + int(self.t)) % 7 - 3
            return float(round(self.plants[u].p) + j)
        return None

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


class _Store:
    def __init__(self):
        from app.store import DEFAULTS
        self.data = dict(DEFAULTS)

    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


VARIANTS = {
    # (modes on, arbitrage_enabled, {unit: (use, mode, sell)})
    "modes_off": (False, False, {}),
    "arb_only": (False, True, {}),
    "p2_sf_ls": (True, False, {"a": (True, "solar_first", False),
                               "b": (True, "leftover_solar", False)}),
    "p2_arb": (True, True, {"a": (True, "solar_first", False),
                            "b": (True, "leftover_solar", False)}),
    "p3_ls": (True, False, {"a": (True, "solar_first", False),
                            "b": (True, "leftover_solar", False)}),
    "p3_ch_sell": (True, False, {"a": (True, "solar_first", False),
                                 "b": (True, "cheap_hours", True)}),
    "p3_both_ch": (True, False, {"a": (True, "cheap_hours", False),
                                 "b": (True, "cheap_hours", True)}),
}


def options(modes, arb, n=2):
    o = {"meter_entity": "sensor.grid", "meter_source": "ha_signed",
         "p1_net_entity": "sensor.net", "meter_phases": 1,
         "peak_cap_w": PEAK_CAP, "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,
         "battery_modes_enabled": modes, "arbitrage_enabled": arb,
         "cycle_cost_threshold_eur": 0.10,
         "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": int(RATING), "bms_soc_floor": 10,
                        "capacity_kwh": CAP_KWH} for u in "ab"[:n]]}
    return o


def plan_at(now, pts):
    """The firm day-ahead points a site holds at `now`: today, and tomorrow
    once published (13:00 local)."""
    loc = brussels(now)
    day0 = loc.replace(hour=0, minute=0, second=0, microsecond=0)
    end = day0 + timedelta(days=2 if loc.hour >= 13 else 1)
    sel = [PricePoint(start=brussels(s), eur_per_kwh=all_in(p),
                      provisional=False)
           for s, span, p in pts if day0 <= s < end]
    if len(sel) < 2:
        return None
    return TariffPlan(PriceRead.series(sel, PriceProduct.DAY_AHEAD,
                                       "generic_dayahead", now), None,
                      "provider")


def run(variant, house, pts, start, dur, *, check=None, stride=1, n=2,
        soc0=None, refuse=(), extra=None, choice=None, patch=None,
        limits=None):
    modes, arb, var_choice = VARIANTS[variant]
    choice = var_choice if choice is None else choice
    units = tuple("ab"[:n])
    import bisect
    ts = [h[0] for h in house]

    def house_at(t):
        i = bisect.bisect_right(ts, start.timestamp() + t) - 1
        return house[max(0, i)][1]

    wholesale = {s: (span, p) for s, span, p in pts}
    starts = sorted(wholesale)

    def price_at(now):
        i = bisect.bisect_right(starts, now) - 1
        if i < 0:
            return None
        s = starts[i]
        span, p = wholesale[s]
        return p if now < s + timedelta(seconds=span) else None

    store = _Store()
    for u, (use, mode, sell) in choice.items():
        store.data[unit_key(f"inverter-{u}", "pack_use")] = use
        store.data[unit_key(f"inverter-{u}", "pack_mode")] = mode
        store.data[unit_key(f"inverter-{u}", "pack_sell")] = sell
    stats = dict(imp_kwh=0.0, exp_kwh=0.0, cost=0.0, exp_rev=0.0,
                 out_kwh={u: 0.0 for u in "ab"[:n]},
                 in_kwh={u: 0.0 for u in "ab"[:n]},
                 ticks=0, quarters={}, trade_ticks=0)

    async def go():
        hass = Hass(units, house_at, soc0, refuse, limits)
        import tempfile
        db = os.path.join(tempfile.mkdtemp(prefix="p3r-"), "tariff.db")
        c = Controller(dict(options(modes, arb, n), **(extra or {})), hass,
                       store, _Pub(), stop=asyncio.Event(), peak_db_path=db)
        for i in c.inverters:
            i.dev_min, i.dev_max = (limits or {}).get(i.name[-1],
                                                      (-5000.0, 5000.0))
        if patch is not None:
            patch(c)
        c.auto = True
        real_dt = m.datetime
        clock = [start]

        class _DT(real_dt):
            @classmethod
            def now(cls, tz=None):
                return clock[0] if tz is None else clock[0].astimezone(tz)

        m.datetime = _DT
        plan_hour = None
        try:
            t = 0.0
            while t < dur:
                now = start + timedelta(seconds=t)
                hass.t = t
                clock[0] = now
                if plan_hour != now.replace(minute=0, second=0,
                                            microsecond=0):
                    plan_hour = now.replace(minute=0, second=0,
                                            microsecond=0)
                    c.tariff_plan, c.tariff_at = plan_at(now, pts), now
                else:
                    c.tariff_at = now
                imp, exp = split_signed(hass.grid())
                c.p1.submit(make_sample("ha_signed", imp, exp, phases=1,
                                        ingest_ts=now,
                                        ingest_mono=time.monotonic()))
                await c.read_inputs(now_mono=t)
                await c.cycle()
                for p in hass.plants.values():
                    p.step(TICK, 10.0)
                g = hass.grid()
                pr = price_at(now)
                dt_h = TICK / 3600.0
                if g > 0:
                    stats["imp_kwh"] += g / 1000 * dt_h
                    stats["cost"] += g / 1000 * dt_h * (all_in(pr) if pr
                                                         is not None else 0)
                else:
                    stats["exp_kwh"] += -g / 1000 * dt_h
                    stats["exp_rev"] += -g / 1000 * dt_h * (pr or 0.0)
                for u, p in hass.plants.items():
                    if p.p > 0:
                        stats["out_kwh"][u] += p.p / 1000 * dt_h
                    else:
                        stats["in_kwh"][u] -= p.p / 1000 * dt_h
                q = int(now.timestamp() // 900)
                acc = stats["quarters"].setdefault(q, [0.0, 0, False])
                acc[0] += max(0.0, g)
                acc[1] += 1
                if (getattr(c.arb, "action", None) or "").endswith("CHARGE"):
                    acc[2] = True
                    stats["trade_ticks"] += 1
                stats["ticks"] += 1
                if check is not None:
                    check(c, hass, now, pr)
                t += TICK
        finally:
            m.datetime = real_dt
        return c
    c = asyncio.run(go())
    return stats, c


def summary(stats):
    qs = stats["quarters"]
    worst = max((a[0] / a[1] for a in qs.values()), default=0.0)
    worst_trade = max((a[0] / a[1] for a in qs.values() if a[2]),
                      default=0.0)
    cyc = {u: round(v / CAP_KWH, 2) for u, v in stats["out_kwh"].items()}
    return {"import_kwh": round(stats["imp_kwh"], 2),
            "export_kwh": round(stats["exp_kwh"], 2),
            "import_eur": round(stats["cost"], 2),
            "export_eur": round(stats["exp_rev"], 2),
            "net_eur": round(stats["cost"] - stats["exp_rev"], 2),
            "cycles": cyc, "worst_quarter_w": round(worst, 1),
            "worst_charge_quarter_w": round(worst_trade, 1),
            "charge_ticks": stats["trade_ticks"]}


LOOP_DAYS = [f"2026-09-{d:02d}" for d in range(2, 18)]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--variant", default="p3_ls")
    ap.add_argument("--days", type=int, default=16)
    ap.add_argument("--captured", action="store_true")
    a = ap.parse_args()
    if a.captured:
        days = ["2026-09-09", "2026-09-10"]
        pts = qh_points()
    else:
        days = LOOP_DAYS[:a.days]
        pts = epex_points()
    house = load_house(days)
    start = datetime.fromisoformat(days[0] + "T00:00:00+02:00") \
        .astimezone(timezone.utc)
    dur = 86400.0 * len(days)
    t0 = time.time()
    stats, _c = run(a.variant, house, pts, start, dur)
    out = summary(stats)
    out["variant"], out["days"] = a.variant, len(days)
    out["wall_s"] = round(time.time() - t0, 1)
    print(json.dumps(out))


if __name__ == "__main__":
    main()