"""Kuroko: of two speakers on separate channels, does the sound a speaker is making claim the floor?

Kuroko hears both speakers as log-mel frames and answers 25 times a second for each of them (`ANSWERS`). An
answer rests on the audio up to its own moment and on nothing later, but for the 0.6 ms a resampler reaches ahead
at 48 kHz, so the answers can be used while a conversation unfolds. The write-up says how the answers become
end-of-turn and interruption events: https://spero.patientdesk.ai/blog/kuroko

    import soundfile as sf
    import kuroko

    model = kuroko.load(".")  # the directory that holds model.safetensors and config.json
    audio, rate = sf.read("conversation.wav", dtype="int16")  # (samples, 2): one speaker on each channel
    answers = kuroko.answers(model, kuroko.features(audio, rate))  # (2, 11, n): speaker, answer, 25 a second
    claims = answers[:, kuroko.ANSWERS.index("claim")]

This file is the model and what it needs to run, and nothing of how it was trained. It needs numpy, scipy, torch
and safetensors.
"""

from __future__ import annotations

import json
import math
from pathlib import Path

import numpy as np
import torch
from numpy.typing import NDArray
from scipy.signal import resample_poly
from torch import nn

RATE = 16_000  # the sample rate the model hears
HOP = 320  # 20 ms: one frame
WINDOW = 2 * HOP  # a frame holds the 40 ms that end with it
FRAME_HZ = RATE // HOP
ANSWER_HZ = FRAME_HZ // 2  # the model answers every other frame
MELS = 80
LOW_HZ, HIGH_HZ = 60.0, 7_600.0
WINDOW_S = 40.0  # the most the model looks back
AHEAD_S = ((0.0, 0.2), (0.2, 0.6), (0.6, 1.2), (1.2, 2.0))
ANSWERS = (
    "voice",  # the speaker is making a sound
    "claim",  # that sound claims the floor: the answer the events are read from
    "hold",  # the speaker holds the floor, pauses included
    *(f"{who} speaks {start:g} to {end:g} s ahead" for who in ("own", "partner") for start, end in AHEAD_S),
)
_BLOCK = 20_000  # frames turned into spectra at once


def mel_filters() -> NDArray[np.float32]:
    """Triangular filters (MELS, WINDOW // 2 + 1), evenly spaced on the mel scale from LOW_HZ to HIGH_HZ."""
    low, high = (2595.0 * math.log10(1.0 + hz / 700.0) for hz in (LOW_HZ, HIGH_HZ))
    edges = 700.0 * (10.0 ** (np.linspace(low, high, MELS + 2) / 2595.0) - 1.0)
    bins = np.fft.rfftfreq(WINDOW, 1.0 / RATE)
    rising = (bins[None] - edges[:-2, None]) / (edges[1:-1] - edges[:-2])[:, None]
    falling = (edges[2:, None] - bins[None]) / (edges[2:] - edges[1:-1])[:, None]
    filters: NDArray[np.float32] = np.clip(np.minimum(rising, falling), 0.0, None).astype(np.float32)
    return filters


def log_mel(audio: NDArray[np.float32]) -> NDArray[np.float32]:
    """Log-mel frames (frames, MELS) of 16 kHz audio: frame k holds the 40 ms that end at sample (k + 1) * HOP."""
    frames = len(audio) // HOP
    padded = np.concatenate([np.zeros(HOP, np.float32), audio[: frames * HOP]])
    windows = np.lib.stride_tricks.sliding_window_view(padded, WINDOW)[::HOP]
    taper = np.hanning(WINDOW).astype(np.float32)
    filters = mel_filters()
    out = np.empty((frames, MELS), np.float32)
    for start in range(0, frames, _BLOCK):
        power = np.abs(np.fft.rfft(windows[start : start + _BLOCK] * taper, axis=1)) ** 2
        out[start : start + _BLOCK] = np.log(power @ filters.T + 1e-10)
    return out


def features(audio: NDArray[np.int16] | NDArray[np.floating], rate: int) -> NDArray[np.float16]:
    """The frames of a conversation given as (samples, 2), one speaker on each channel, at `rate` Hz, as 16-bit
    samples or as floats between -1 and 1: (2, frames, MELS), in the half precision the model learned from.

    The model was trained on 16-bit recordings at 48 kHz, brought to 16 kHz this way. Read as floats, audio of
    more than 16 bits gives answers that differ in the second decimal."""
    scale = 32768.0 if audio.dtype == np.int16 else 1.0
    step = math.gcd(RATE, rate)
    sides = []
    for channel in (0, 1):
        heard = resample_poly(audio[:, channel].astype(np.float32) / scale, RATE // step, rate // step)
        sides.append(log_mel(heard.astype(np.float32)).astype(np.float16))
    return np.stack(sides)


def _turned(x: torch.Tensor, angle: torch.Tensor) -> torch.Tensor:
    """Rotary positions: each pair of `x`'s last dimension turned by its position's angle."""
    first, second = x[..., 0::2], x[..., 1::2]
    cos, sin = angle.cos(), angle.sin()
    return torch.stack([first * cos - second * sin, first * sin + second * cos], dim=-1).flatten(-2)


class Block(nn.Module):
    """One transformer layer whose attention sees the frames up to its own."""

    def __init__(self, width: int, heads: int) -> None:
        super().__init__()
        self.heads = heads
        self.before_attention = nn.LayerNorm(width)
        self.query_key_value = nn.Linear(width, 3 * width)
        self.attended = nn.Linear(width, width)
        self.before_mlp = nn.LayerNorm(width)
        self.mlp = nn.Sequential(nn.Linear(width, 4 * width), nn.GELU(), nn.Linear(4 * width, width))

    def forward(self, x: torch.Tensor, angle: torch.Tensor) -> torch.Tensor:
        batch, frames, width = x.shape
        parts = self.query_key_value(self.before_attention(x)).view(batch, frames, 3, self.heads, -1).transpose(1, 3)
        query, key, value = _turned(parts[:, :, 0], angle), _turned(parts[:, :, 1], angle), parts[:, :, 2]
        mixed = nn.functional.scaled_dot_product_attention(query, key, value, is_causal=True)
        x = x + self.attended(mixed.transpose(1, 2).reshape(batch, frames, width))
        out: torch.Tensor = x + self.mlp(self.before_mlp(x))
        return out


class Kuroko(nn.Module):
    """Frames of both speakers in, a logit for each of `ANSWERS` and each speaker out, at half the frame rate.
    Both speakers go through the same weights, each once as "own" and once as "partner"."""

    mean: torch.Tensor
    std: torch.Tensor
    turns: torch.Tensor  # how fast each pair of an attention head's dimensions turns with position

    def __init__(self, mels: int = MELS, width: int = 384, layers: int = 8, heads: int = 6) -> None:
        super().__init__()
        self.register_buffer("mean", torch.zeros(mels))
        self.register_buffer("std", torch.ones(mels))
        self.inputs = nn.Conv1d(mels, width, 5)
        self.halve = nn.Conv1d(width, width, 2, stride=2)
        self.join = nn.Linear(2 * width, width)
        self.blocks = nn.ModuleList(Block(width, heads) for _ in range(layers))
        self.after = nn.LayerNorm(width)
        self.answers = nn.Linear(width, len(ANSWERS))
        pairs = width // heads // 2
        self.register_buffer("turns", 1.0 / 10_000 ** (torch.arange(pairs) / pairs), persistent=False)

    def forward(self, mel: torch.Tensor) -> torch.Tensor:
        """`mel` (batch, 2, frames, mels), from where the model starts to listen, gives logits
        (batch, 2, frames // 2, len(ANSWERS)). An answer sees its own two input frames and what came before."""
        batch, speakers, frames, mels = mel.shape
        frames = frames // 2 * 2
        heard = ((mel[:, :, :frames] - self.mean) / self.std).reshape(batch * speakers, frames, mels).transpose(1, 2)
        heard = nn.functional.gelu(self.inputs(nn.functional.pad(heard, (4, 0))))  # zeros before the audio
        own = nn.functional.gelu(self.halve(heard)).transpose(1, 2).reshape(batch, speakers, frames // 2, -1)
        x = self.join(torch.cat([own, own.flip(1)], dim=-1)).reshape(batch * speakers, frames // 2, -1)
        angle = torch.arange(frames // 2, device=x.device)[:, None] * self.turns
        for block in self.blocks:
            x = block(x, angle)
        logits: torch.Tensor = self.answers(self.after(x)).reshape(batch, speakers, frames // 2, len(ANSWERS))
        return logits


def load(directory: str | Path = ".", device: str | torch.device = "cpu") -> Kuroko:
    """The model whose `config.json` and `model.safetensors` are in a directory, ready to answer."""
    from safetensors.torch import load_file

    config = json.loads((Path(directory) / "config.json").read_text())
    model = Kuroko(config["mels"], config["width"], config["layers"], config["heads"])
    model.load_state_dict(load_file(Path(directory) / "model.safetensors"))
    return model.to(device).eval()


@torch.no_grad()
def answers(model: Kuroko, mel: NDArray[np.floating]) -> NDArray[np.float32]:
    """The model's answers for a whole conversation given as (2, frames, MELS): probabilities
    (2, len(ANSWERS), frames // 2), the answer at index k having heard the audio up to (k + 1) / ANSWER_HZ seconds.

    The model listens to stretches of `WINDOW_S` that start half a stretch apart, and of each but the first only
    the answers of its second half are kept: every answer then comes from a model that has listened for at least
    half a stretch, or from the conversation's start.

    Each stretch goes through the model by itself and in full precision. Batched or in half precision, an answer
    shifts in its third decimal with what else is in the batch, and with it the frame at which a threshold is
    crossed: the answers before a moment must not change with the audio after it."""
    device = model.mean.device
    window = round(WINDOW_S * FRAME_HZ)
    half = window // 2
    frames = mel.shape[1] // 2 * 2
    out = np.zeros((2, len(ANSWERS), frames // 2), np.float32)
    for start in range(0, max(1, frames - half), half):
        stretch = np.zeros((1, 2, window, mel.shape[2]), np.float32)
        piece = mel[:, start : start + window]
        stretch[0, :, : piece.shape[1]] = piece
        scores = torch.sigmoid(model(torch.from_numpy(stretch).to(device))[0]).permute(0, 2, 1).cpu().numpy()
        keep = 0 if start == 0 else half // 2  # in answer frames
        end = min(frames, start + window) // 2
        out[:, :, start // 2 + keep : end] = scores[:, :, keep : end - start // 2]
    return out