"""Controlled encoder-scale ablation; fixed frame-wise LayerNorm, no extra bits.

The failed baseline's continuous features grew to hundreds in RMS, saturating
the tanh quantizer. Normalize before both continuous decoding and quantization
so the warmup and transition have a consistent feature-scale constraint.
This changes the encoder architecture and must be identified in checkpoints.
"""
import torch
import torch.nn.functional as F

from tinyasr.tinycall import TinyCallCodec, TinyCallEncoder


ARCHITECTURE = 'unit_layernorm_encoder_v1'


class UnitScaleEncoder(TinyCallEncoder):
    def forward(self, audio):
        features = super().forward(audio)
        return F.layer_norm(features.transpose(1, 2), (self.dimension,), eps=1e-5).transpose(1, 2)


class UnitScaleCodec(TinyCallCodec):
    def __init__(self):
        super().__init__()
        # Preserve the baseline's exact initial weights and RNG consumption.
        # Fixed LayerNorm has neither learned parameters nor running statistics.
        with torch.random.fork_rng(devices=[]):
            encoder = UnitScaleEncoder()
        encoder.load_state_dict(self.encoder.state_dict(), strict=True)
        self.encoder = encoder