"""Same bounded live state and transport with explicit phase residual synthesis."""
import torch

from tinyasr.codec_live_v1 import LiveDecoder
from tinyasr.codec_live_bitstream_v1 import LiveBitReader
from tinyasr.codec_phase_residual_v1 import PhaseResidualDecoder
from tinyasr.codec_streaming_v1 import StreamingDecoder


class PhaseStreamingDecoder(StreamingDecoder):
    def __init__(self, codec):
        if not isinstance(codec.decoder, PhaseResidualDecoder):
            raise ValueError('Expected explicit phase residual decoder')
        super().__init__(codec)

    def _frame(self, features):
        context = torch.cat((self.history, features), -1)
        decoder = self.codec.decoder
        spectrum = decoder.complex_spectrum(decoder.blocks(decoder.input(context)), last=True)
        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


class PhaseLiveDecoder(LiveDecoder):
    def __init__(self, codec, model_sha256):
        self.audio = PhaseStreamingDecoder(codec)
        self.transport = LiveBitReader(model_sha256)