"""Compact wire format for /api/analyze results (``?format=columnar``).

The object format repeats every key per node/link and writes each float as
~18 characters of JSON (9.5 MB for a 2,000-node, 25-step network). Here the
time series are packed per variable into one little-endian float64 array per
element kind, base64-encoded (~10.7 bytes per value), so full double precision
is kept -- the Native-vs-WASM comparison tolerances (rel 1e-8) still apply.

Layout of each packed series: element-major, i.e. element ``i``'s values are
``data[i * steps:(i + 1) * steps]`` in the order of ``ids``.
src/api.ts (decodeColumnarResult) turns it back into the object format.
"""
import base64
import sys
from array import array

FORMAT = "spipe-columnar-1"
NODE_SERIES = ("demand", "head", "pressure", "quality")
LINK_SERIES = ("flow", "velocity", "headloss", "quality")
PUMP_FIELDS = ("energy", "efficiency", "totalKwh", "peakKw", "avgKw", "utilization")


def _pack(items: dict, ids: list, key: str, steps: int):
    """Packed base64 series, or None when any element lacks a full series
    (e.g. no water-quality run) -- the decoder then leaves it empty."""
    flat = array("d")
    for element_id in ids:
        series = items[element_id]["timeSeries"].get(key) or []
        if len(series) != steps:
            return None
        flat.extend(series)
    if sys.byteorder != "little":
        flat.byteswap()
    return base64.b64encode(flat.tobytes()).decode("ascii")


def to_columnar(result: dict) -> dict:
    steps = len(result.get("timeSteps") or [])
    nodes, links = result["nodes"], result["links"]
    node_ids, link_ids = list(nodes), list(links)

    # Pump series and the pump's scalar fields (initial energy/efficiency, kWh
    # totals) use distinct keys -- they share names in the object format.
    pumps = {
        lid: {
            "energySeries": links[lid]["timeSeries"].get("energy") or [],
            "efficiencySeries": links[lid]["timeSeries"].get("efficiency") or [],
            **{f: links[lid].get(f) for f in PUMP_FIELDS if f in links[lid]},
        }
        for lid in link_ids
        if "totalKwh" in links[lid]
    }

    return {
        "format": FORMAT,
        **{k: v for k, v in result.items() if k not in ("nodes", "links")},
        "nodes": {
            "ids": node_ids,
            "types": [nodes[i]["type"] for i in node_ids],
            "coordinates": [nodes[i]["coordinates"] for i in node_ids],
            "series": {k: _pack(nodes, node_ids, k, steps) for k in NODE_SERIES},
        },
        "links": {
            "ids": link_ids,
            "types": [links[i]["type"] for i in link_ids],
            "fromNodes": [links[i]["fromNode"] for i in link_ids],
            "toNodes": [links[i]["toNode"] for i in link_ids],
            "vertices": [links[i]["vertices"] for i in link_ids],
            "series": {k: _pack(links, link_ids, k, steps) for k in LINK_SERIES},
            "pumps": pumps,
        },
    }