"""Exact streaming for ONE causal, stride-one Conv1d. Not a Demucs adapter.

Run: python chunking.py
Requires torch==2.10.0. The tests compare streaming with a full causal pass.
"""
from __future__ import annotations
import json
from collections.abc import Iterable, Iterator
import torch
from torch import Tensor, nn
from torch.nn import functional as F


def causal_chunks(conv: nn.Conv1d, chunks: Iterable[Tensor]) -> Iterator[Tensor]:
    if conv.stride != (1,) or conv.padding != (0,):
        raise ValueError("Requires stride=1 and padding=0")
    context = (conv.kernel_size[0] - 1) * conv.dilation[0]
    state: Tensor | None = None
    for chunk in chunks:
        if chunk.ndim != 3 or chunk.shape[1] != conv.in_channels or chunk.shape[-1] == 0:
            raise ValueError("Expected a nonempty [batch, in_channels, time] chunk")
        if state is None:
            state = chunk.new_zeros(chunk.shape[0], chunk.shape[1], context)
        if state.shape[:2] != chunk.shape[:2] or state.dtype != chunk.dtype or state.device != chunk.device:
            raise ValueError("Chunk stream changed batch, channels, dtype or device")
        with torch.inference_mode():
            joined = torch.cat((state, chunk), dim=-1)
            output = conv(joined)
            # Clone the tail. A view would keep the entire joined allocation alive.
            state = joined[..., -context:].clone() if context else joined[..., :0].clone()
        yield output


def main() -> None:
    torch.set_num_threads(1)
    torch.manual_seed(20260924)
    tests = []
    for kernel, dilation, chunk_size in [(5, 1, 17), (5, 2, 3), (1, 1, 11)]:
        conv = nn.Conv1d(2, 4, kernel, dilation=dilation).eval()
        x = torch.randn(1, 2, 103)
        context = (kernel - 1) * dilation
        with torch.inference_mode():
            full = conv(F.pad(x, (context, 0)))
            # Test-only concatenation. A production sink consumes one output at a time.
            streamed = torch.cat(list(causal_chunks(conv, x.split(chunk_size, dim=-1))), dim=-1)
        torch.testing.assert_close(full, streamed, rtol=1e-5, atol=1e-6)
        tests.append({"kernel": kernel, "dilation": dilation, "chunk_size": chunk_size,
                      "max_abs_error": (full-streamed).abs().max().item(), "status": "passed"})
    print(json.dumps({"torch": torch.__version__, "tests": tests}, indent=2))

if __name__ == "__main__":
    main()