# SPDX-License-Identifier: Apache-2.0
# Copyright 2026 ayutaz
"""KoeAct-JP inference from the released files: 16 kHz speech -> gated Action JSON.

The release directory holds ``model.safetensors`` (Ralomi's input layer and blocks 1..4 as
released, the fine-tuned blocks 5..8, the CTC classifier and the decoder) and ``config.json``
(the ``p_act`` threshold). Nothing is downloaded. The outputs equal ``koeact.infer``'s torch
engine (``tests/test_runtime.py``); this module imports no training or evaluation code, so it
runs from the bundled ``code/`` directory alone (``koeact.release``).

    python -m koeact.runtime --model <release dir> a.wav b.wav
    python -m koeact.runtime --model <release dir> --mic

``from_device_line`` turns a CoreS3 ``ACTION ... p_act=... ids=...`` line (the firmware stops
at the decoder's piece ids) into the same gated JSON; ``follow`` does it for every recognition in
a stream of serial lines (``--serial <port>``, needs pyserial; touch the screen to record 4 s).
"""

import argparse
import json
import math
import re
from pathlib import Path

import numpy as np

from koeact.action import parse_output, to_json
from koeact.features.logmel import HOP, MAX_SECONDS, SAMPLE_RATE, WINDOW, log_mel
from koeact.model import pieces

FORMAT = "koeact-runtime-1"
FIRST_BLOCK = 4  # the decoder reads layer 4; blocks 5..8 belong to the fine-tuned head
MIN_SAMPLES = WINDOW + 3 * HOP  # Ralomi needs four log-Mel frames
CLEAN_RMS_DBFS = -23.0
CLEAN_LEAD, CLEAN_TAIL = 0.5, 0.3  # seconds of silence before and after
PEAK = 0.99
PAIR = re.compile(r"(\w+)=(-?\d+(?:\.\d+)?)")


def build_net(ft_blocks: int):
    """Ralomi's lower half (input layer, blocks 1..4) and the plan B head over it."""
    from torch import nn

    from koeact.model.plan_b import DecoderB
    from koeact.model.ralomi_torch import RalomiTorch

    class KoeActNet(nn.Module):
        def __init__(self) -> None:
            super().__init__()
            ralomi = RalomiTorch()
            self.subsampling_projection = ralomi.subsampling_projection
            self.input_normalization = ralomi.input_normalization
            self.lower = nn.ModuleList(ralomi.blocks[:FIRST_BLOCK])
            self.head = DecoderB(ralomi, ft_blocks)

    return KoeActNet()


def _set_level(x: np.ndarray, dbfs: float) -> np.ndarray:
    rms = float(np.sqrt(np.mean(np.square(x, dtype=np.float64)))) if len(x) else 0.0
    if not 0 < rms < math.inf:  # silence stays silence; NaN or inf is left as it is
        return x
    y = x * (10 ** ((dbfs - 20 * math.log10(rms)) / 20))
    peak = float(np.abs(y).max())
    return y * (PEAK / peak) if peak > PEAK else y


def clean(audio: np.ndarray) -> np.ndarray:
    """``koeact.features.augment.clean``: -23 dBFS RMS, 0.5 s of silence before, 0.3 s after."""
    lead = np.zeros(round(CLEAN_LEAD * SAMPLE_RATE), np.float32)
    tail = np.zeros(round(CLEAN_TAIL * SAMPLE_RATE), np.float32)
    return np.concatenate([lead, _set_level(audio, CLEAN_RMS_DBFS), tail]).astype(np.float32)


_clean = clean  # ``KoeAct.__call__`` has a ``clean`` parameter


def prepare(audio: np.ndarray) -> np.ndarray:
    """Cut to 15 s and pad short audio to four log-Mel frames (as the evaluations did)."""
    x = audio[: MAX_SECONDS * SAMPLE_RATE]
    if len(x) < MIN_SAMPLES:
        x = np.pad(x, (0, MIN_SAMPLES - len(x)))
    return x


def resample(x: np.ndarray, sr: int) -> np.ndarray:
    if sr == SAMPLE_RATE:
        return x
    from math import gcd

    from scipy.signal import resample_poly

    g = gcd(SAMPLE_RATE, sr)
    return resample_poly(x, SAMPLE_RATE // g, sr // g).astype(np.float32)


def load_audio(path: Path) -> np.ndarray:
    """Mono float32 at 16 kHz (other rates are resampled; channels are averaged)."""
    import soundfile as sf

    x, sr = sf.read(path, dtype="float32", always_2d=True)
    return resample(x.mean(axis=1), sr)


class KoeAct:
    def __init__(self, net, threshold: float) -> None:
        self.net, self.threshold = net.eval(), threshold
        self.grammar = pieces.ActionGrammar()

    @classmethod
    def from_dir(cls, path: Path | str, threads: int = 0) -> "KoeAct":
        import torch
        from safetensors.torch import load_file

        if threads:
            torch.set_num_threads(threads)
        path = Path(path)
        config = json.loads((path / "config.json").read_text("utf-8"))
        if config.get("format") != FORMAT:
            raise ValueError(f"{path}: not a {FORMAT} release ({config.get('format')!r})")
        net = build_net(config["ft_blocks"])
        net.load_state_dict(load_file(path / "model.safetensors"))
        return cls(net, config["threshold"])

    def _frames(self, features: np.ndarray):
        """``[1, steps, 2, 320]``: layer 4 (rounded to fp16, as the training cache) and zeros
        for layer 8, which the head recomputes from layer 4."""
        import torch

        from koeact.model.ralomi_torch import SUBSAMPLING

        x = torch.from_numpy(features)[None]
        steps = x.shape[1] // SUBSAMPLING
        h = x[:, : steps * SUBSAMPLING].reshape(1, steps, -1)
        h = torch.relu(self.net.input_normalization(self.net.subsampling_projection(h)))
        for block in self.net.lower:
            h = block(h)
        h4 = h[0].numpy()
        frames = np.stack([h4, np.zeros_like(h4)], axis=1).astype(np.float16)
        return torch.from_numpy(frames.astype(np.float32))[None]

    def __call__(self, audio: np.ndarray, clean: bool = True) -> dict:
        """One utterance (float32, 16 kHz) -> ``output`` (gated JSON), ``raw``, ``p_act``.

        ``clean`` applies the training preprocessing first (v0.1.0 did not)."""
        import torch

        from koeact.model import decoder_a2

        if clean:
            audio = _clean(audio)
        with torch.no_grad():
            x = self._frames(log_mel(prepare(audio)))
            mask = torch.ones(x.shape[:2], dtype=torch.bool)
            d = decoder_a2.greedy(self.net.head, x, mask, self.grammar)[0]
        return gate(to_json(d.calls), round(d.p_act, 6), self.threshold)


def gate(raw: str, p_act: float, threshold: float) -> dict:
    accepted = p_act >= threshold  # below the threshold the answer is "[]"
    return {"output": raw if accepted else "[]", "raw": raw, "p_act": p_act,
            "accepted": accepted}  # fmt: skip


def from_device_line(line: str, threshold: float) -> dict:
    """A CoreS3 ``ACTION``/``A2`` line (``p_act=`` and the piece ids after ``ids=``) -> JSON."""
    head, _, ids = line.partition(" ids=")
    fields = dict(PAIR.findall(head))
    calls = parse_output(pieces.decode_text([int(i) for i in ids.split()])).calls or []
    return gate(to_json(calls), float(fields["p_act"]), threshold)


def follow(lines, threshold: float):
    """Gated JSON for each microphone recognition (``A2 ... real=1 ... ids=``) or ``ACTION`` line;
    ``real=0`` lines are timings over stand-in frames and are skipped."""
    for line in lines:
        if " ids=" not in line:
            continue
        if line.startswith("ACTION ") or (line.startswith("A2 ") and " real=1 " in line):
            yield from_device_line(line, threshold)


def serial_lines(port: str):
    import serial

    with serial.Serial(port, 115200, timeout=1) as s:
        while True:
            raw = s.readline()
            if raw:
                yield raw.decode("utf-8", "replace").rstrip()


def record_utterance() -> np.ndarray:
    """One utterance from the default microphone: Enter starts, Enter stops."""
    import sounddevice as sd

    sr = int(sd.query_devices(kind="input")["default_samplerate"])
    chunks: list[np.ndarray] = []
    with sd.InputStream(samplerate=sr, channels=1, dtype="float32",
                        callback=lambda data, *_: chunks.append(data[:, 0].copy())):  # fmt: skip
        input("  ● 話してください（終わったら Enter） / speak, then press Enter")
    return resample(np.concatenate(chunks) if chunks else np.zeros(0, np.float32), sr)


def main() -> None:
    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    p.add_argument("wav", nargs="*", type=Path)
    p.add_argument("--model", type=Path, default=Path("."), help="the release directory")
    p.add_argument("--threads", type=int, default=0)
    p.add_argument("--mic", action="store_true", help="Enter starts and stops; q quits")
    p.add_argument(
        "--no-clean",
        action="store_true",
        help="feed the audio as is (v0.1.0); the default levels it and pads silence",
    )
    p.add_argument("--serial", help="a CoreS3 with the KoeAct firmware (reads config.json only)")
    args = p.parse_args()
    if args.serial:
        threshold = json.loads((args.model / "config.json").read_text("utf-8"))["threshold"]
        for r in follow(serial_lines(args.serial), threshold):
            print(json.dumps(r, ensure_ascii=False), flush=True)
        return
    model = KoeAct.from_dir(args.model, args.threads)
    if args.mic:
        while input("Enter で録音（q で終了） / Enter to record, q to quit > ").strip() != "q":
            r = model(record_utterance(), clean=not args.no_clean)
            print(f"  {r['output']}  (p_act {r['p_act']:.3f})")
        return
    for w in args.wav:
        x = load_audio(w)
        r = model(x, clean=not args.no_clean)
        print(json.dumps({"file": str(w), "output": r["output"], "p_act": r["p_act"],
                          "accepted": r["accepted"]}, ensure_ascii=False), flush=True)  # fmt: skip


if __name__ == "__main__":
    main()