# SPDX-License-Identifier: Apache-2.0
"""
PlotSpec -> JSON.

The desktop renderer (`audio_visualization_widget`) draws a `PlotSpec` with
pyqtgraph. The web renderer draws the *same* `PlotSpec` with a browser plotting
library, so the only thing this module owns is the wire format between them:
numpy arrays become plain lists, `NaN`/`inf` (gated PSR, empty windows) become
JSON `null` so the client can break the line instead of drawing a spike.

Nothing here is web-framework-specific — `serialize_spec` is the exact partner of
`show_specs`: same input, different canvas. Keeping it free of FastAPI imports
means the desktop side could adopt it unchanged if a future client ever needed
the same payload.

Why the heatmap is sent losslessly: the desktop renderer resamples the linear STFT
rows onto a log-spaced grid for a log frequency axis (pyqtgraph's ImageItem is
affine-only and won't follow a log axis). The browser *can* draw a non-linear
axis from a linearly spaced image, so no resampling happens here — the client gets
the native rows plus the axis-extent information and does the mapping itself.
That keeps warm-up cost identical to the desktop app and makes the expensive
log-grid interpolation happen once per view change rather than once per request.
"""

from __future__ import annotations

import math
from typing import Any, Optional

import numpy as np

from plotspec import (
  PlotSpec, RefLineProps, ViewState, X_RELATIVE,
)


def _numbers(values) -> list:
  """numpy array -> JSON-safe list, with non-finite floats mapped to None."""
  if values is None:
    return []
  arr = np.asarray(values, dtype=float).ravel()
  out: list[Optional[float]] = []
  for v in arr.tolist():
    out.append(None if (v is None or not math.isfinite(v)) else float(v))
  return out


def _ints(values) -> list[float]:
  """numpy array -> JSON-safe list of floats (no None mapping — index axes)."""
  return [float(v) for v in np.asarray(values, dtype=float).ravel().tolist()]


def _bounds(values) -> Optional[list[float]]:
  """[min, max] of the finite values, or None when everything is non-finite."""
  if values is None:
    return None
  arr = np.asarray(values, dtype=float).ravel()
  arr = arr[np.isfinite(arr)]
  if not arr.size:
    return None
  return [float(arr.min()), float(arr.max())]


def _serialize_axes(axes) -> dict[str, Any]:
  return {
    "x_label": axes.x_label,
    "y_label": axes.y_label,
    "y_log": bool(axes.y_log),
    "x_log": bool(axes.x_log),
    "y_range": list(axes.y_range) if axes.y_range else None,
    "x_range": list(axes.x_range) if axes.x_range else None,
    "y_log_allowed": bool(axes.y_log_allowed),
    "x_log_allowed": bool(axes.x_log_allowed),
  }


def _serialize_heatmap(hm) -> dict[str, Any]:
  y = np.asarray(hm.y, dtype=float)
  z = np.asarray(hm.z, dtype=np.float32)
  return {
    "x": _numbers(hm.x),
    "y": _numbers(y),
    "z": [[float(v) for v in row] for row in z.tolist()],
    "z_min": float(hm.z_min),
    "z_max": float(hm.z_max),
    "cmap": hm.cmap,
    "label": hm.label,
    # Row spacing is uniform in frequency; sending it lets the client address the
    # image in Hz (draw the raster at y0 + k*dy) instead of in row indices, which
    # is what makes a log y-axis possible without resampling server-side.
    "y0": float(y[0]) if y.size else 0.0,
    "dy": float(y[1] - y[0]) if y.size > 1 else 1.0,
    "deg": False,  # z is already dB or dBFS, never angle
  }


def serialize_spec(spec: PlotSpec) -> dict[str, Any]:
  """One dataset, one metric, ready to draw."""
  payload = {
    "title": spec.title,
    "axes": _serialize_axes(spec.axes),
    "hlines": [
      {
        "y": float(hl.y),
        "label": hl.label,
        "color": hl.color,
        "style": hl.style,
        "width": float(hl.width),
      }
      for hl in spec.hlines
    ],
    "annotations": list(spec.annotations),
    "segments": [],
    "series": [],
    "heatmap": None,
    "y_bounds": None,
  }

  for band in spec.bands:
    lo = np.broadcast_to(np.asarray(band.lo, dtype=float), np.asarray(band.x).shape)
    payload["segments"].append({
      "kind": "band",
      "x": _numbers(band.x),
      "lo": _numbers(lo),
      "hi": _numbers(band.hi),
      "label": band.label,
      "color": band.color,
    })

  for curve in spec.curves:
    payload["series"].append({
      "kind": "line",
      "x": _numbers(curve.x),
      "y": _numbers(curve.y),
      "label": curve.label,
      "width": float(curve.width),
      "color": curve.color,
    })

  if spec.heatmap is not None:
    payload["heatmap"] = _serialize_heatmap(spec.heatmap)

  # Value range the browser should frame to, before any user zoom. For bands the
  # envelope is lo..hi, not just hi — RMS fills from zero, waveform spans ±peak.
  values = [np.asarray(c.y, dtype=float) for c in spec.curves]
  for band in spec.bands:
    values.append(np.asarray(band.lo, dtype=float))
    values.append(np.asarray(band.hi, dtype=float))
  if not spec.axes.y_range:
    finite = [v[np.isfinite(v)] for v in values]
    finite = [v for v in finite if v.size]
    if finite:
      payload["y_bounds"] = [float(min(v.min() for v in finite)),
                             float(max(v.max() for v in finite))]
  return payload


def serialize_view(view: ViewState) -> dict[str, Any]:
  return {"y_log": view.y_log, "x_log": view.x_log, "x_mode": view.x_mode}


def serialize_ref_lines(lines: list[RefLineProps]) -> list[dict[str, Any]]:
  return [
    {"value": float(l.value), "color": l.color, "style": l.style, "label": l.label}
    for l in lines
  ]


def relative_x_mode(view: ViewState) -> bool:
  return view.x_mode == X_RELATIVE