"""Streaming phoneme recognizer: a small causal TCN with CTC output.

Design constraints (for later ESP32-S3 deployment via TFLite-Micro):
  - only Conv1d / BatchNorm / ReLU / residual add (all INT8-friendly and
    supported by esp-nn optimized kernels once folded/exported)
  - strictly causal (left padding only) so it can run on a live stream
  - ~350k parameters -> ~400 KB at INT8, fits the S3 easily

Input:  (B, T, 40) log-mel frames at 10 ms
Output: (B, T', NUM_CLASSES) logits at 20 ms (stem has stride 2)
"""

import torch
import torch.nn as nn

from .phones import NUM_CLASSES


class CausalConv1d(nn.Module):
    """Conv1d with left-only padding (streaming safe)."""

    def __init__(self, in_ch, out_ch, kernel, stride=1, dilation=1, groups=1):
        super().__init__()
        self.left_pad = dilation * (kernel - 1)
        self.conv = nn.Conv1d(in_ch, out_ch, kernel, stride=stride,
                              dilation=dilation, groups=groups)

    def forward(self, x):
        x = nn.functional.pad(x, (self.left_pad, 0))
        return self.conv(x)


class TCNBlock(nn.Module):
    """Depthwise-separable causal conv block with residual."""

    def __init__(self, channels, kernel, dilation):
        super().__init__()
        self.dw = CausalConv1d(channels, channels, kernel,
                               dilation=dilation, groups=channels)
        self.bn1 = nn.BatchNorm1d(channels)
        self.pw = nn.Conv1d(channels, channels, 1)
        self.bn2 = nn.BatchNorm1d(channels)
        self.act = nn.ReLU()

    def forward(self, x):
        y = self.act(self.bn1(self.dw(x)))
        y = self.bn2(self.pw(y))
        return self.act(x + y)


class PhonemeTCN(nn.Module):
    def __init__(self, n_mels=40, channels=192, kernel=5,
                 dilations=(1, 2, 4, 8, 1, 2, 4, 8)):
        super().__init__()
        self.stem = CausalConv1d(n_mels, channels, kernel, stride=2)
        self.stem_bn = nn.BatchNorm1d(channels)
        self.act = nn.ReLU()
        self.blocks = nn.Sequential(
            *[TCNBlock(channels, kernel, d) for d in dilations])
        self.head = nn.Conv1d(channels, NUM_CLASSES, 1)

    def forward(self, feats):
        """feats (B, T, n_mels) -> logits (B, T//2, NUM_CLASSES)."""
        x = feats.transpose(1, 2)              # (B, n_mels, T)
        x = self.act(self.stem_bn(self.stem(x)))
        x = self.blocks(x)
        return self.head(x).transpose(1, 2)    # (B, T', classes)

    @staticmethod
    def out_lengths(in_lengths):
        """Output frame count for input frame counts (stride-2 stem)."""
        return torch.div(in_lengths - 1, 2, rounding_mode="floor") + 1