hoa64/report.py

"""JSON spatial report — compact payload for other models / tools."""

from __future__ import annotations

import json
import time
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any, Optional, Sequence, Union

import numpy as np

from .analysis import angular_error_deg, doa_from_intensity, field_energy, peak_direction
from .audio_io import ensure_hoa_channels, read_wav
from .basis import MAX_ORDER, N_CHANNELS, channel_names
from .stream import (
    FrameAnalysis,
    SourceSpec,
    analyze_hoa_frames,
    analyze_hoa_stft_bands,
    encode_mono_plane_wave,
    encode_scene,
    hoa_rms,
)

PathLike = Union[str, Path]

REPORT_SCHEMA_VERSION = "spatial-hoa.report.v1"


@dataclass
class SpatialReport:
    """Machine-readable spatial summary for LLM / diffusion conditioning."""

    schema: str = REPORT_SCHEMA_VERSION
    kind: str = "spatial_field"
    sample_rate: int = 0
    n_samples: int = 0
    duration_sec: float = 0.0
    max_order: int = MAX_ORDER
    n_channels: int = N_CHANNELS
    # Global broadband
    energy: float = 0.0
    doa_az_deg: float = 0.0
    doa_el_deg: float = 0.0
    peak_az_deg: float = 0.0
    peak_el_deg: float = 0.0
    peak_power: float = 0.0
    # Optional detail
    bands: list[dict] = field(default_factory=list)
    frames: list[dict] = field(default_factory=list)
    sources_hint: list[dict] = field(default_factory=list)
    channel_rms: list[float] = field(default_factory=list)
    notes: list[str] = field(default_factory=list)
    meta: dict = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)

    def to_json(self, *, indent: int | None = 2) -> str:
        return json.dumps(self.to_dict(), indent=indent)

    def save(self, path: PathLike) -> None:
        Path(path).write_text(self.to_json() + "\n", encoding="utf-8")

    def one_liner(self) -> str:
        """Short string suitable for tool_result / system hints."""
        return (
            f"spatial: DOA az={self.doa_az_deg:.1f}° el={self.doa_el_deg:.1f}° "
            f"peak=({self.peak_az_deg:.1f},{self.peak_el_deg:.1f}) "
            f"E={self.energy:.4g} T={self.duration_sec:.2f}s"
        )


def _frames_to_dicts(frames: Sequence[FrameAnalysis]) -> list[dict]:
    return [
        {
            "t_sec": f.t_center_sec,
            "energy": f.energy,
            "doa_az_deg": f.doa_az_deg,
            "doa_el_deg": f.doa_el_deg,
            "peak_az_deg": f.peak_az_deg,
            "peak_el_deg": f.peak_el_deg,
            "order1_energy": f.order1_energy,
        }
        for f in frames
    ]


def report_from_hoa(
    hoa: np.ndarray,
    sample_rate: int,
    *,
    max_order: int = MAX_ORDER,
    frame_ms: float = 40.0,
    hop_ms: float = 20.0,
    include_frames: bool = True,
    include_bands: bool = True,
    include_peak_map: bool = True,
    max_frames: int = 64,
    meta: Optional[dict] = None,
    notes: Optional[list[str]] = None,
) -> SpatialReport:
    """Build a SpatialReport from HOA coefficients (C,T) or (C,)."""
    a = np.asarray(hoa, dtype=np.float64)
    if a.ndim == 1:
        a = a.reshape(-1, 1)
    a = ensure_hoa_channels(a, max_order=max_order)
    n_samples = int(a.shape[1])
    duration = n_samples / float(sample_rate) if sample_rate else 0.0

    # Broadband DOA must be AC-safe (time-mean of coeffs → 0 for audio).
    # 1) global intensity from products W*X, W*Y, W*Z
    # 2) energy-weighted blend of per-frame DOAs
    from .basis import az_el_from_unit, unit_vector

    W, Yc, Zc, Xc = a[0], a[1], a[2], a[3]
    I_global = np.array(
        [
            float(np.mean(W * Xc)),
            float(np.mean(W * Yc)),
            float(np.mean(W * Zc)),
        ]
    )
    nI = float(np.linalg.norm(I_global))
    if nI > 1e-18:
        az, el = az_el_from_unit(I_global / nI, degrees=True)
        az, el = float(az), float(el)
    else:
        az, el = 0.0, 0.0

    frames = analyze_hoa_frames(
        a,
        sample_rate,
        frame_ms=frame_ms,
        hop_ms=hop_ms,
        max_order=max_order,
        peak_grid=False,
    )
    if frames:
        vecs = []
        for f in frames:
            if f.energy <= 1e-18:
                continue
            u = unit_vector(f.doa_az_deg, f.doa_el_deg, degrees=True)
            vecs.append(u * f.energy)
        if vecs:
            v = np.sum(vecs, axis=0)
            n = float(np.linalg.norm(v))
            if n > 1e-15:
                az_f, el_f = az_el_from_unit(v / n, degrees=True)
                # Prefer frame blend when global intensity is weak
                if nI < 1e-12:
                    az, el = float(az_f), float(el_f)
                else:
                    # average unit vectors of global + frames
                    ug = unit_vector(az, el, degrees=True)
                    uf = unit_vector(float(az_f), float(el_f), degrees=True)
                    um = ug + uf
                    nm = float(np.linalg.norm(um))
                    if nm > 1e-15:
                        az, el = az_el_from_unit(um / nm, degrees=True)
                        az, el = float(az), float(el)

    energy = float(np.mean(np.sum(a * a, axis=0)))

    if include_peak_map:
        # Pseudo-static field from RMS * sign(corr with W) for map peak
        rms = np.sqrt(np.mean(a * a, axis=1) + 1e-30)
        sign = np.sign(np.mean(a * a[0:1, :], axis=1) + 1e-30)
        pseudo = rms * sign
        paz, pel, pv = peak_direction(
            pseudo, n_azi=72, n_el=36, max_order=min(max_order, 3)
        )
    else:
        paz, pel, pv = az, el, energy

    bands: list[dict] = []
    if include_bands and n_samples >= 256:
        bands = analyze_hoa_stft_bands(a, sample_rate)

    frame_dicts: list[dict] = []
    if include_frames and frames:
        step = max(1, len(frames) // max_frames)
        frame_dicts = _frames_to_dicts(frames[::step][:max_frames])

    return SpatialReport(
        sample_rate=int(sample_rate),
        n_samples=n_samples,
        duration_sec=duration,
        max_order=max_order,
        n_channels=(max_order + 1) ** 2,
        energy=energy,
        doa_az_deg=float(az),
        doa_el_deg=float(el),
        peak_az_deg=float(paz),
        peak_el_deg=float(pel),
        peak_power=float(pv),
        bands=bands,
        frames=frame_dicts,
        channel_rms=[float(x) for x in hoa_rms(a)[:8]],  # first 8 for brevity
        notes=list(notes or []),
        meta=dict(meta or {}),
    )


def report_from_mono_wav(
    path: PathLike,
    azimuth_deg: float,
    elevation_deg: float = 0.0,
    *,
    max_order: int = MAX_ORDER,
    **kwargs,
) -> SpatialReport:
    """Load mono (or mixdown) WAV, encode as plane wave from known direction, report."""
    audio, sr = read_wav(path)
    if audio.shape[0] == 1:
        mono = audio[0]
    else:
        mono = np.mean(audio, axis=0)
        notes_extra = ["mixed multi-channel input down to mono before plane-wave encode"]
    hoa = encode_mono_plane_wave(mono, azimuth_deg, elevation_deg, max_order=max_order)
    notes = list(kwargs.pop("notes", []) or [])
    if audio.shape[0] > 1:
        notes.extend(notes_extra)
    notes.append(
        f"encoded mono as plane wave az={azimuth_deg}° el={elevation_deg}°"
    )
    meta = dict(kwargs.pop("meta", None) or {})
    meta.update(
        {
            "source_path": str(path),
            "encode": "mono_plane_wave",
            "encode_az_deg": azimuth_deg,
            "encode_el_deg": elevation_deg,
            "created_unix": time.time(),
        }
    )
    return report_from_hoa(
        hoa, sr, max_order=max_order, notes=notes, meta=meta, **kwargs
    )


def report_from_ambix_wav(
    path: PathLike,
    *,
    max_order: int = MAX_ORDER,
    **kwargs,
) -> SpatialReport:
    """Load multi-channel Ambix WAV (ACN order) and analyze."""
    audio, sr = read_wav(path)
    hoa = ensure_hoa_channels(audio, max_order=max_order)
    notes = list(kwargs.pop("notes", []) or [])
    notes.append(f"loaded Ambix-style multi-channel WAV with {audio.shape[0]} ch")
    meta = dict(kwargs.pop("meta", None) or {})
    meta.update(
        {
            "source_path": str(path),
            "encode": "ambix_wav",
            "input_channels": int(audio.shape[0]),
            "created_unix": time.time(),
        }
    )
    return report_from_hoa(
        hoa, sr, max_order=max_order, notes=notes, meta=meta, **kwargs
    )


def report_from_scene(
    sources: Sequence[SourceSpec],
    sample_rate: int,
    *,
    max_order: int = MAX_ORDER,
    **kwargs,
) -> SpatialReport:
    """Analyze a synthetic multi-source scene."""
    hoa = encode_scene(sources, max_order=max_order)
    hints = [
        {
            "label": s.label or f"src{i}",
            "azimuth_deg": s.azimuth_deg,
            "elevation_deg": s.elevation_deg,
            "rms": float(np.sqrt(np.mean(np.asarray(s.signal) ** 2) + 1e-30)),
        }
        for i, s in enumerate(sources)
    ]
    notes = list(kwargs.pop("notes", []) or [])
    notes.append(f"synthetic scene with {len(sources)} plane-wave source(s)")
    meta = dict(kwargs.pop("meta", None) or {})
    meta["encode"] = "synthetic_scene"
    rep = report_from_hoa(
        hoa, sample_rate, max_order=max_order, notes=notes, meta=meta, **kwargs
    )
    rep.sources_hint = hints
    return rep


def load_report(path: PathLike) -> dict:
    return json.loads(Path(path).read_text(encoding="utf-8"))