"""Bounded inference for the frozen TinyCall-inspired codec.

Finite context is recomputed on each 20 ms hop. This favors a small, auditable
state over speed; model parameters and caller-owned input/output are excluded
from the persistent state accounting. Only batch-one CPU inference is supported.
"""
import math

import torch


class StreamingEncoder:
    hop = 320
    # Encoder receptive field is 1438 samples. Four complete past hops are
    # sufficient, and preserving hop alignment preserves every strided phase.
    history_samples = 1280

    def __init__(self, codec):
        if next(codec.parameters()).device.type != 'cpu':
            raise ValueError('Streaming v1 supports CPU inference only')
        self.codec = codec
        self.history = torch.empty(0)
        self.pending = torch.empty(0)
        self.samples = 0
        self.frames = 0
        self.finished = False

    @property
    def state_bytes(self):
        return sum(t.untyped_storage().nbytes() for t in [self.history, self.pending])

    def _frame(self, hop):
        context = torch.cat((self.history, hop))
        features = self.codec.encoder(context[None])[..., -1:]
        _, codes = self.codec.quantizer(features)
        self.history = context[-self.history_samples:].clone()
        self.frames += 1
        return codes

    @torch.inference_mode()
    def push(self, audio):
        if self.finished:
            raise ValueError('Encoder stream already finished')
        if audio.ndim != 1 or audio.device.type != 'cpu' or audio.dtype != torch.float32:
            raise ValueError('Expected one-dimensional CPU float32 waveform')
        if not torch.isfinite(audio).all():
            raise ValueError('Nonfinite waveform')
        self.samples += audio.numel()
        # The temporary concatenation is input-sized, but no view of it survives
        # this call; persistent storage is independent of the input duration.
        data = torch.cat((self.pending, audio))
        complete = data.numel() // self.hop * self.hop
        frames = [self._frame(data[i:i + self.hop]) for i in range(0, complete, self.hop)]
        self.pending = data[complete:].clone()
        return torch.cat(frames, -1) if frames else torch.empty((1, 4, 6, 0), dtype=torch.long)

    @torch.inference_mode()
    def finish(self):
        if self.finished or not self.samples:
            raise ValueError('Empty or already finished encoder stream')
        self.finished = True
        if self.pending.numel():
            # Let each strided layer pad its own final partial activation, as
            # in whole-waveform inference. Padding the raw waveform instead
            # would generate biased hidden activations beyond its true end.
            result = self._frame(self.pending)
        else:
            result = torch.empty((1, 4, 6, 0), dtype=torch.long)
        self.pending = torch.empty(0)
        return result


class StreamingDecoder:
    hop = 320
    history_frames = 30

    def __init__(self, codec):
        if next(codec.parameters()).device.type != 'cpu':
            raise ValueError('Streaming v1 supports CPU inference only')
        self.codec = codec
        self.history = torch.empty((1, 64, 0))
        self.previous_right = torch.empty(0)
        window = codec.decoder.window
        self.denominator = window[:self.hop].square() + window[self.hop:].square()
        self.frames = 0
        self.emitted = 0
        self.finished = False

    @property
    def state_bytes(self):
        return sum(t.untyped_storage().nbytes() for t in [self.history, self.previous_right, self.denominator])

    def _frame(self, features):
        context = torch.cat((self.history, features), -1)
        decoder = self.codec.decoder
        values = decoder.spectrum(decoder.blocks(decoder.input(context)))[..., -1]
        log_magnitude, phase = values.chunk(2, dim=1)
        magnitude = log_magnitude.clamp(-10, 5).exp()
        phase = math.pi * phase.tanh()
        spectrum = torch.complex(magnitude * phase.cos(), magnitude * phase.sin())
        frame = torch.fft.irfft(spectrum, n=decoder.n_fft)[0] * decoder.window
        self.history = context[..., -self.history_frames:].clone()
        output = ((self.previous_right + frame[:self.hop]) / self.denominator
                  if self.previous_right.numel() else torch.empty(0))
        self.previous_right = frame[self.hop:].clone()
        self.emitted += output.numel()
        return output

    @torch.inference_mode()
    def push(self, codes):
        if self.finished:
            raise ValueError('Decoder stream already finished')
        if codes.ndim != 4 or codes.shape[:3] != (1, 4, 6) or codes.dtype != torch.long or codes.device.type != 'cpu':
            raise ValueError('Expected CPU int64 codes [1,4,6,frames]')
        frames = []
        for i in range(codes.shape[-1]):
            features = self.codec.quantizer.decode(codes[..., i:i + 1])
            frames.append(self._frame(features))
            self.frames += 1
        return torch.cat(frames) if frames else torch.empty(0)

    @torch.inference_mode()
    def finish(self, samples):
        if self.finished or not self.frames:
            raise ValueError('Empty or already finished decoder stream')
        if not isinstance(samples, int) or not (self.frames - 1) * self.hop < samples <= self.frames * self.hop:
            raise ValueError('Sample count inconsistent with received code frames')
        self.finished = True
        # Match the offline decoder's single terminal zero latent. It is not a
        # quantized code and does not carry an undeclared transmitted feature.
        remaining = samples - self.emitted
        output = self._frame(torch.zeros((1, 64, 1)))[:remaining]
        self.emitted = samples
        return output