"""Explicit padding-safe native Matcha, with unchanged learned state layout.

``create(config)`` accepts the unchanged capacity-192 parent configuration.
Callers must record ARCHITECTURE_ID separately in branch/checkpoint metadata:
identical state keys do not identify the changed forward contract. In particular,
load the CONTROL129000 no-adapter parent strictly, not a rich-text adapter state.
The returned MatchaTTS retains its native ``forward`` and ``synthesise`` methods.
No trainer, frontend, optimizer, dropout, vocoder, or global forward hooks change.

Only PyTorch is required to import the numerical helpers. ``native_types`` and
``create`` additionally require the installed native Matcha dependency stack.
"""
from functools import lru_cache, partial
from pathlib import Path
import sys

import torch
from torch import nn


ARCHITECTURE_ID = "inflect-v3mini-matcha-padding-safe-v1"
ACOUSTIC_PARAMETERS = 10_780_425


def boolean_sequence_mask(mask, x, *, prefix=False):
    """Validate Bx1xT binary validity; True means an allowed attention key.

    Additive masks (including 0/-inf) are deliberately rejected, not cast to bool.
    Every example must contain at least one valid frame.
    """
    if x.ndim != 3 or x.shape[0] == 0 or x.shape[-1] == 0:
        raise ValueError("Expected a nonempty BxCxT tensor")
    if mask.shape != (x.shape[0], 1, x.shape[-1]) or mask.device != x.device:
        raise ValueError("Mask must have shape Bx1xT on the input device")
    if not torch.all((mask == 0) | (mask == 1)):
        raise ValueError("Mask must be finite binary validity, not additive attention bias")
    valid = mask.bool()
    if not valid.any(dim=-1).all():
        raise ValueError("Every example must contain a nonempty valid sequence")
    if prefix and torch.any(valid[..., 1:] & ~valid[..., :-1]):
        raise ValueError("Decoder masks must be right-padded contiguous prefixes")
    return valid


def _zero_padding(x, valid):
    # Multiplication by zero does not remove NaN/Inf padding.
    return x.masked_fill(~valid, 0)


def _group_norm(x, valid, norm):
    b, c, length = x.shape
    if c != norm.num_channels:
        raise ValueError("GroupNorm channel count differs from input")
    # Match native accumulation precision under autocast without lowering FP64.
    work = _zero_padding(x, valid)
    if work.dtype in (torch.float16, torch.bfloat16):
        work = work.float()
    grouped = work.reshape(b, norm.num_groups, c // norm.num_groups, length)
    weights = valid.unsqueeze(1)
    count = weights.sum((2, 3), keepdim=True).to(work.dtype) * (c // norm.num_groups)
    mean = grouped.sum((2, 3), keepdim=True) / count
    centered = (grouped - mean).masked_fill(~weights, 0)
    variance = centered.square().sum((2, 3), keepdim=True) / count
    output = (centered * torch.rsqrt(variance + norm.eps)).reshape(b, c, length)
    if norm.affine:
        output = output * norm.weight.to(work.dtype)[None, :, None]
        output = output + norm.bias.to(work.dtype)[None, :, None]
    return _zero_padding(output.to(x.dtype), valid)


def valid_frame_group_norm(x, mask, norm):
    """Biased GroupNorm statistics over valid time and channels within each group.

    Uses the supplied native nn.GroupNorm's eps and learned affine tensors;
    registers no parameters/buffers. Invalid outputs and input gradients are zero.
    """
    return _group_norm(x, boolean_sequence_mask(mask, x), norm)


def masked_flow_loss(prediction, target, mask):
    """CFM squared error divided by valid frames times mel channels, not padding."""
    if prediction.shape != target.shape:
        raise ValueError("Flow prediction and target shapes must match")
    valid = boolean_sequence_mask(mask, prediction)
    prediction, target = _zero_padding(prediction, valid), _zero_padding(target, valid)
    if prediction.dtype in (torch.float16, torch.bfloat16):
        prediction, target = prediction.float(), target.float()
    error = prediction - target
    return error.square().sum() / (valid.sum() * target.shape[1])


@lru_cache(maxsize=1)
def native_decoder_types():
    """Return (PaddingSafeBlock1D, PaddingSafeDecoder) lazily.

    Constructors run the vendor constructors first. Replacement initialization is
    RNG-neutral and strictly copies state, retaining registration order and names.
    No vendor classes, existing models, or attention processors are patched.
    """
    vendor = str(Path(__file__).resolve().parents[1] / "external/Matcha-TTS-primary")
    if vendor not in sys.path:
        sys.path.insert(0, vendor)
    from matcha.models.components.decoder import Block1D, Decoder

    class PaddingSafeBlock1D(Block1D):
        def forward(self, x, mask):
            return self._forward_validated(x, boolean_sequence_mask(mask, x))

        def _forward_validated(self, x, valid):
            h = self.block[0](_zero_padding(x, valid))
            return _zero_padding(self.block[2](_group_norm(h, valid, self.block[1])), valid)

    class PaddingSafeDecoder(Decoder):
        def __init__(self, *args, **kwargs):
            super().__init__(*args, **kwargs)
            from matcha.models.components.transformer import BasicTransformerBlock

            for stage in (self.down_blocks, self.mid_blocks, self.up_blocks):
                for parts in stage:
                    if any(not isinstance(block, BasicTransformerBlock) for block in parts[1]):
                        raise ValueError("Padding-safe variant supports native transformer blocks only")
            for module in list(self.modules()):
                for name, child in list(module.named_children()):
                    if isinstance(child, Block1D):
                        conv, norm = child.block[:2]
                        with torch.random.fork_rng(devices=[]):
                            replacement = PaddingSafeBlock1D(
                                conv.in_channels, conv.out_channels, groups=norm.num_groups)
                        replacement.load_state_dict(child.state_dict(), strict=True)
                        setattr(module, name, replacement)

        @staticmethod
        def _resnet(resnet, x, valid, t):
            # Same native ResnetBlock1D operations, but the decoder has already
            # validated this mask. Avoid per-Block1D device-to-host reductions.
            x = _zero_padding(x, valid)
            h = resnet.block1._forward_validated(x, valid)
            h += resnet.mlp(t).unsqueeze(-1)
            h = resnet.block2._forward_validated(h, valid)
            return h + resnet.res_conv(x)

        @staticmethod
        def _attention(x, blocks, valid, t):
            x = _zero_padding(x, valid).transpose(1, 2)
            # Explicit singleton query axis: Diffusers expands heads and SDPA
            # broadcasts queries. A BxT mask alone can misbroadcast for B > 1.
            for block in blocks:
                x = block(hidden_states=x, attention_mask=valid, timestep=t)
            return _zero_padding(x.transpose(1, 2), valid)

        def forward(self, x, mask, mu, t, spks=None, cond=None):
            valid = boolean_sequence_mask(mask, x, prefix=True)
            if mu.shape != x.shape:
                raise ValueError("Flow input and text mean shapes must match")
            if cond is not None:
                raise ValueError("Native padding-safe decoder has no adapter conditioning")
            x, mu = _zero_padding(x, valid), _zero_padding(mu, valid)
            if not torch.isfinite(x).all() or not torch.isfinite(mu).all():
                raise ValueError("Valid flow input and text mean must be finite")
            t = self.time_mlp(self.time_embeddings(t))
            x = torch.cat((x, mu), dim=1)
            if spks is not None:
                x = torch.cat((x, spks.unsqueeze(-1).expand(-1, -1, x.shape[-1])), dim=1)

            hiddens, masks = [], []
            current = valid
            for index, (resnet, blocks, downsample) in enumerate(self.down_blocks):
                x = self._resnet(resnet, x, current, t)
                x = self._attention(x, blocks, current, t)
                hiddens.append(x)
                masks.append(current)
                x = downsample(x)
                # Native's final down block is stride one, not a downsample.
                if index < len(self.down_blocks) - 1:
                    current = current[..., ::2]
                x = _zero_padding(x, current)

            for resnet, blocks in self.mid_blocks:
                x = self._resnet(resnet, x, current, t)
                x = self._attention(x, blocks, current, t)

            for resnet, blocks, upsample in self.up_blocks:
                current, skip = masks.pop(), hiddens.pop()
                # Stride-two transposed convolutions produce one extra right
                # frame for odd skip lengths. Crop before the next convolution.
                x = x[..., :skip.shape[-1]]
                x = self._resnet(resnet, torch.cat((x, skip), dim=1), current, t)
                x = self._attention(x, blocks, current, t)
                x = upsample(x)

            x = self.final_block._forward_validated(x, valid)
            return _zero_padding(self.final_proj(x), valid)

    return PaddingSafeBlock1D, PaddingSafeDecoder


@lru_cache(maxsize=1)
def native_types():
    """Return decoder types plus PaddingSafeCFM with native training dependencies."""
    PaddingSafeBlock1D, PaddingSafeDecoder = native_decoder_types()
    from matcha.models.components.flow_matching import CFM

    class PaddingSafeCFM(CFM):
        def __init__(self, in_channels, out_channel, cfm_params, decoder_params,
                     n_spks=1, spk_emb_dim=64):
            super().__init__(in_channels, out_channel, cfm_params, decoder_params, n_spks, spk_emb_dim)
            with torch.random.fork_rng(devices=[]):
                estimator = PaddingSafeDecoder(
                    in_channels=in_channels + (spk_emb_dim if n_spks > 1 else 0),
                    out_channels=out_channel, **decoder_params)
            estimator.load_state_dict(self.estimator.state_dict(), strict=True)
            self.estimator = estimator

        def compute_loss(self, x1, mask, mu, spks=None, cond=None):
            valid = boolean_sequence_mask(mask, x1, prefix=True)
            if mu.shape != x1.shape:
                raise ValueError("Flow target and text mean shapes must match")
            x1 = _zero_padding(x1, valid)
            if not torch.isfinite(x1).all():
                raise ValueError("Valid flow target must be finite")
            t = torch.rand([mu.shape[0], 1, 1], device=mu.device, dtype=mu.dtype)
            z = torch.randn_like(x1)
            y = (1 - (1 - self.sigma_min) * t) * z + t * x1
            u = x1 - (1 - self.sigma_min) * z
            prediction = self.estimator(y, valid, mu, t.squeeze(), spks, cond)
            return masked_flow_loss(prediction, u, valid), y

    return PaddingSafeBlock1D, PaddingSafeDecoder, PaddingSafeCFM


def create(config):
    """Build FP32 native-compatible acoustic model; strict state_dict loading only.

    Config is the unchanged capacity-192/full-context config, including its parent
    architecture marker. Persist ARCHITECTURE_ID outside it to distinguish this
    experimental forward contract. This function does not load a checkpoint or
    create/restore an optimizer; native optimizer configuration is preserved.
    """
    from scripts import train_v3mini_unet_capacity_192 as native

    values = native.native_config(config)
    _, _, PaddingSafeCFM = native_types()
    from omegaconf import OmegaConf
    from matcha.models.matcha_tts import MatchaTTS

    for key in ("encoder", "decoder", "cfm"):
        values[key] = OmegaConf.create(values[key])
    model = MatchaTTS(**values, optimizer=partial(torch.optim.Adam, lr=1e-4))
    original = model.decoder
    with torch.random.fork_rng(devices=[]):
        replacement = PaddingSafeCFM(
            2 * values["encoder"].encoder_params.n_feats,
            values["encoder"].encoder_params.n_feats,
            values["cfm"], values["decoder"], values["n_spks"], values["spk_emb_dim"])
    replacement.load_state_dict(original.state_dict(), strict=True)
    model.decoder = replacement
    model.architecture_id = ARCHITECTURE_ID
    model.float()
    if sum(p.numel() for p in model.parameters()) != ACOUSTIC_PARAMETERS:
        raise ValueError("Padding-safe acoustic parameter count must be exactly 10,780,425")
    native.parameter_budget(model)
    return model