#!/usr/bin/env python3
"""Heat-spreading runs on one ET-SoC-1 card: heat chosen shires in on/off cycles while sampling every
temperature sensor with DM_CMD_GET_ASIC_PVT_SAMPLES (needs firmware with Roman's patch 0003).

The sampler holds the management node (/dev/etN_mgmt) for the whole session; the heater runs on the
ops node, one `timeout 10` process per heating phase. Run it under the card's lock:
  flock -n /run/lock/etsoc-shire0.lock python3 heatmovie.py --out DIR
"""
import argparse
import ctypes
import errno
import fcntl
import json
import os
import select
import struct
import subprocess
import sys
import threading
import time

MAGIC = 0xE7


def _ioc(direction, nr, size):
    return (direction << 30) | (size << 16) | (MAGIC << 8) | nr


class CmdDesc(ctypes.Structure):
    _fields_ = [("cmd", ctypes.c_void_p), ("size", ctypes.c_uint16),
                ("sq_index", ctypes.c_uint16), ("flags", ctypes.c_uint8)]


class RspDesc(ctypes.Structure):
    _fields_ = [("rsp", ctypes.c_void_p), ("size", ctypes.c_uint16), ("cq_index", ctypes.c_uint16)]


PUSH_SQ = _ioc(1, 6, ctypes.sizeof(CmdDesc))
POP_CQ = _ioc(3, 7, ctypes.sizeof(RspDesc))
CMD74 = 74
HOME = os.path.expanduser("~")
HEATER = HOME + "/nekko/build/sparsity/host/sparsity_host"
KERNEL = HOME + "/nekko/build/sparsity/kernel/sparsity.elf"

# name, shires, cores per shire, cycles. Shire numbers are the runtime's compute shires (= TS index).
WEST_HALF = [0, 8, 3, 24, 16, 4, 9, 1, 13, 25, 17, 14, 2, 10, 18, 11, 19, 27]  # die-frame columns c1-c3
EAST_HALF = [12, 20, 28, 21, 29, 5, 22, 30, 6, 26, 15, 7, 23, 31]                # columns c4-c6
QUAD_NW = [0, 8, 3, 24, 16, 4, 9, 1, 13]                                         # rows r0-r2, cols c1-c3
QUAD_SE = [22, 30, 6, 26, 15, 7, 23, 31]                                         # rows r3-r5, cols c4-c6
PLAN = [
    ("core-S14", [14], 1, 3),
    ("shire-S0", [0], 32, 4),
    ("shire-S31", [31], 32, 4),
    ("shire-S14", [14], 32, 4),
    ("block-2x2-centre", [13, 14, 21, 22], 32, 4),
    ("quadrant-NW", QUAD_NW, 32, 5),
    ("quadrant-SE", QUAD_SE, 32, 5),
    ("half-west", WEST_HALF, 32, 4),
    ("half-east", EAST_HALF, 32, 4),
]


# Session 2: the sensor map (every column, then every row of the die frame) and the two other corners.
COLS = {"col-c1": [0, 24, 9, 25, 2, 11], "col-c2": [8, 16, 1, 17, 10, 19], "col-c3": [3, 4, 13, 14, 18, 27],
        "col-c4": [12, 21, 22, 26], "col-c5": [20, 29, 30, 15, 23], "col-c6": [28, 5, 6, 7, 31]}
ROWS = {"row-r0": [0, 8, 3], "row-r1": [24, 16, 4, 12, 20, 28], "row-r2": [9, 1, 13, 21, 29, 5],
        "row-r3": [25, 17, 14, 22, 30, 6], "row-r4": [2, 10, 18, 26, 15, 7], "row-r5": [11, 19, 27, 23, 31]}
QUAD_NE = [12, 20, 28, 21, 29, 5]                                                 # rows r0-r2, cols c4-c6
QUAD_SW = [25, 17, 14, 2, 10, 18, 11, 19, 27]                                    # rows r3-r5, cols c1-c3
PLAN_SCAN = ([(k, v, 32, 3) for k, v in COLS.items()] + [(k, v, 32, 3) for k, v in ROWS.items()]
             + [("quadrant-NE", QUAD_NE, 32, 5), ("quadrant-SW", QUAD_SW, 32, 5)])


class Mgmt:
    def __init__(self, card):
        path = "/dev/et%d_mgmt" % card
        end = time.time() + 5
        while True:
            try:
                self.fd = os.open(path, os.O_RDWR | os.O_NONBLOCK)
                break
            except OSError as e:
                if e.errno in (errno.EPERM, errno.EBUSY) and time.time() < end:
                    time.sleep(0.05)
                    continue
                raise
        self.poller = select.poll()
        self.poller.register(self.fd, select.POLLIN)
        self.tag = 1
        self.rbuf = ctypes.create_string_buffer(1760)
        self.rdesc = RspDesc(ctypes.cast(self.rbuf, ctypes.c_void_p), 1760, 0)

    def pvt(self, group, first=0, count=255, timeout_s=1.0):
        tag = self.tag
        self.tag = (self.tag % 0xFFFE) + 1
        msg = struct.pack("<HHHH", 16, tag, CMD74, 0) + struct.pack("<BBB5x", group, first, count)
        sbuf = ctypes.create_string_buffer(msg, len(msg))
        fcntl.ioctl(self.fd, PUSH_SQ, CmdDesc(ctypes.cast(sbuf, ctypes.c_void_p), len(msg), 0, 0))
        end = time.time() + timeout_s
        while time.time() < end:
            self.poller.poll(5)
            try:
                n = fcntl.ioctl(self.fd, POP_CQ, self.rdesc)
            except OSError as e:
                if e.errno == errno.EAGAIN:
                    continue
                raise
            data = self.rbuf.raw[:n]
            _sz, rtag, rcmd, _fl = struct.unpack_from("<HHHH", data, 0)
            if rtag != tag or rcmd != CMD74:
                continue
            dev_us, status = struct.unpack_from("<Qi", data, 8)
            _g, _f, cnt, stride, _tot = struct.unpack_from("<BBBBB", data, 24)
            fault, = struct.unpack_from("<Q", data, 32)
            words = struct.unpack_from("<%dh" % (cnt * stride), data, 40)
            return status, dev_us, fault, [words[i * stride:(i + 1) * stride] for i in range(cnt)]
        raise TimeoutError("no reply to cmd 74")


class Sampler(threading.Thread):
    def __init__(self, mgmt, path, hz, stop_c):
        super().__init__(daemon=True)
        self.m, self.f, self.period, self.stop_c = mgmt, open(path, "w"), 1.0 / hz, stop_c
        self.running, self.max_c, self.n, self.errors = True, 0, 0, 0
        self.too_hot = threading.Event()

    def run(self):
        nxt = time.monotonic()
        while self.running:
            try:
                st, us, fault, ent = self.m.pvt(0, 0, 34)
                rec = {"t": round(time.time(), 4), "us": us, "st": st, "f": fault, "c": [e[0] for e in ent]}
                if self.n % 20 == 0:
                    st2, _, f2, io = self.m.pvt(4, 0, 1)
                    rec["io"] = io[0][0] if io and st2 == 0 else None
                    rec["hi"] = [e[1] for e in ent]
                    rec["lo"] = [e[2] for e in ent]
                self.f.write(json.dumps(rec) + "\n")
                self.n += 1
                mx = max(rec["c"]) if rec["c"] else 0
                self.max_c = max(self.max_c, mx)
                if mx >= self.stop_c:
                    self.too_hot.set()
            except Exception as ex:  # keep sampling; record the error
                self.errors += 1
                self.f.write(json.dumps({"t": round(time.time(), 4), "err": str(ex)[:200]}) + "\n")
                time.sleep(0.2)
            nxt += self.period
            delay = nxt - time.monotonic()
            if delay > 0:
                time.sleep(delay)
            else:
                nxt = time.monotonic()
        self.f.flush()
        self.f.close()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", required=True)
    ap.add_argument("--card", type=int, default=0)
    ap.add_argument("--hz", type=float, default=20)
    ap.add_argument("--on", type=float, default=8)
    ap.add_argument("--off", type=float, default=12)
    ap.add_argument("--baseline", type=float, default=20)
    ap.add_argument("--stop-c", type=int, default=85)
    ap.add_argument("--only", default="", help="comma-separated run names (default: the whole plan)")
    ap.add_argument("--cycles", type=int, default=0, help="override every run's cycle count (smoke tests)")
    ap.add_argument("--plan", default="main", choices=["main", "scan"])
    a = ap.parse_args()
    os.makedirs(a.out, exist_ok=True)
    ev = open(os.path.join(a.out, "events.jsonl"), "a")

    def event(**kw):
        kw["t"] = round(time.time(), 4)
        ev.write(json.dumps(kw) + "\n")
        ev.flush()
        print(json.dumps(kw), flush=True)

    plan = [p for p in (PLAN_SCAN if a.plan == "scan" else PLAN) if not a.only or p[0] in a.only.split(",")]
    if a.cycles:
        plan = [(n, sh, per, a.cycles) for n, sh, per, _ in plan]
    m = Mgmt(a.card)
    s = Sampler(m, os.path.join(a.out, "samples.jsonl"), a.hz, a.stop_c)
    s.start()
    event(ev="session_start", plan=[p[0] for p in plan], hz=a.hz, on=a.on, off=a.off)
    time.sleep(a.baseline)
    stopped = None
    for name, shires, per, cycles in plan:
        mask = 0
        for sh in shires:
            mask |= 1 << sh
        for k in range(cycles):
            if s.too_hot.is_set():
                stopped = "a sensor reached %d C" % a.stop_c
                break
            event(ev="heat_on", run=name, cycle=k, mask=hex(mask), per_shire=per, shires=shires)
            args = ["timeout", "10", HEATER, "--test", "fma", "--type", "fp32", "--pattern", "none",
                    "--values", "randn", "--shires", hex(mask), "--per-shire", str(per),
                    "--seconds", str(a.on), "--seed", "1", "--kernel", KERNEL]
            with open(os.path.join(a.out, "heater-%s-%d.out" % (name, k)), "w") as hout:
                rc = subprocess.run(args, stdout=hout, stderr=subprocess.STDOUT, stdin=subprocess.DEVNULL).returncode
            event(ev="heat_off", run=name, cycle=k, rc=rc)
            t_end = time.time() + a.off
            while time.time() < t_end:
                if s.too_hot.is_set():
                    break
                time.sleep(0.1)
        if stopped:
            break
    time.sleep(2)
    s.running = False
    s.join(timeout=5)
    event(ev="session_end", samples=s.n, sampler_errors=s.errors, max_c=s.max_c, stopped=stopped)
    os.close(m.fd)


if __name__ == "__main__":
    main()