import torch, math
from timm.layers import use_fused_attn

def ConvND(dim, *args, **kwargs):
    conv_cls = [torch.nn.Conv1d, torch.nn.Conv2d, torch.nn.Conv3d][dim-1]
    return conv_cls(*args, **kwargs)

def ConvTransposeND(dim, *args, **kwargs):
    conv_cls = [torch.nn.ConvTranspose1d, torch.nn.ConvTranspose2d, torch.nn.ConvTranspose3d][dim-1]
    return conv_cls(*args, **kwargs)

class LayerNormND(torch.nn.Module):
    """Per-position channel normalization for channels-first tensors of any spatial dimensionality.
    Permutes to channels-last, applies F.layer_norm, permutes back. Matches timm's LayerNorm2d."""
    def __init__(self, num_channels, eps=1e-6):
        super().__init__()
        self.weight = torch.nn.Parameter(torch.ones(num_channels))
        self.bias = torch.nn.Parameter(torch.zeros(num_channels))
        self.eps = eps
        self.num_channels = num_channels

    def forward(self, x):
        return torch.nn.functional.layer_norm(
            x.movedim(1, -1), (self.num_channels,), self.weight, self.bias, self.eps
        ).movedim(-1, 1)

class _ChannelsLast(torch.nn.Module):
    def forward(self, x):
        return x.movedim(1, -1)

class _ChannelsFirst(torch.nn.Module):
    def forward(self, x):
        return x.movedim(-1, 1)

class Residual(torch.nn.Module):
    """During training, the residual branch is dropped with probability `drop_prob` and scaled by 1/(1-drop_prob) when kept, preserving the expected value of the output (same scaling convention as timm's DropPath)."""
    def __init__(self, main, drop_prob=0.0):
        super().__init__()
        self.main = main
        self.drop_prob = drop_prob

    def forward(self, x):
        if self.drop_prob and self.training:
            if torch.rand(()).item() < self.drop_prob:
                return x
            return x + self.main(x) / (1.0 - self.drop_prob)
        return x + self.main(x)

class Softsign(torch.nn.Module):
    """Learnable per-channel companding: r * x / (σ + |x|). The output is guaranteed to lie in [-(2^(bits-1)-1), 2^(bits-1)-1], i.e. the representable range of a signed `bits`-bit integer (default [-127, 127] for 8-bit). Requires 3 MACs per element: one abs, one add, one div (the multiply by r is a fixed scalar and can be fused)."""
    def __init__(self, dim, num_channels, bits=8):
        super().__init__()
        self.shape = [1, -1] + [1]*dim
        self.r = 2**(bits - 1) - 1
        self._σ = torch.nn.Parameter(torch.full((num_channels,), float(self.r - 1)))

    def forward(self, x):
        σ = (self._σ.abs() + 1e-6).view(self.shape)
        return self.r * x / (σ + x.abs())

class UniformTrainingNoise(torch.nn.Module):
    """Adds uniform noise in [-delta, delta] during training only, where delta = k + 0.5. With k=0 (delta=0.5). With k >= 1, it simulates the wider reconstruction tolerance of near-lossless codecs (e.g. JPEG-LS allows reconstructed values to differ by +/-k from the original)."""
    def __init__(self, k=0):
        super().__init__()
        self.delta = k + 0.5
    def forward(self, x):
        if self.training:
            x = x + torch.rand_like(x)*(2*self.delta) - self.delta
        return x

class ChannelAffine(torch.nn.Module):
    """Per-channel affine transform: x * γ (+ β if bias=True). Learnable scale γ is initialized to 1 and optional bias β to 0."""
    def __init__(self, dim, num_channels, bias=False):
        super().__init__()
        self.shape = [1, -1] + [1] * dim
        self.γ = torch.nn.Parameter(torch.ones(num_channels))
        self.β = torch.nn.Parameter(torch.zeros(num_channels)) if bias else None

    def forward(self, x):
        x = x * self.γ.view(self.shape)
        if self.β is not None:
            x = x + self.β.view(self.shape)
        return x

def SoftsignCompanding(dim, num_channels, qat=False, bits=8, k=0, affine=True, bias=False):
    """(Quantization-aware) learnable softsign companding. Somewhat similar to to DyTanh (arXiv:2503.10622) but uses softsign for lower MAC cost and optionally uses the limited range companded output as an opportunity to quantize activations."""
    return torch.nn.Sequential(
        Softsign(dim, num_channels, bits),
        UniformTrainingNoise(k) if qat else torch.nn.Identity(),
        ChannelAffine(dim, num_channels, bias=bias) if affine else torch.nn.Identity()
    )

class SAPE(torch.nn.Module):
    """Stacked Absolute Position Embedding. Concatenates `dim` new channels encoding centered integer coordinates (in patch units) for each spatial dimension. No learnable parameters. See NaViT (arXiv:2307.06304) for an ablation on a similar position encoding strategy."""
    def __init__(self, dim):
        super().__init__()
        self.dim = dim

    def forward(self, x):
        spatial = x.shape[2:]
        coords = []
        for i, size in enumerate(spatial):
            pos = torch.arange(size, device=x.device, dtype=x.dtype) - (size - 1) / 2
            shape = [1, 1] + [1] * self.dim
            shape[2 + i] = size
            coords.append(pos.view(shape).expand(1, 1, *spatial))
        pos = torch.cat(coords, dim=1).expand(x.shape[0], -1, *spatial)
        return torch.cat([x, pos], dim=1)


class AttentionND(torch.nn.Module):
    """N-dimensional multi-head self-attention. Flattens spatial dims into the sequence and uses torch.nn.functional.scaled_dot_product_attention. QKV and output projections are 1×1 convolutions. The internal attention dimension is num_heads × head_dim, which may differ from num_channels."""
    def __init__(self, dim, num_channels, num_heads, head_dim=32):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = head_dim
        attn_dim = num_heads * head_dim
        self.qkv = ConvND(dim, num_channels, 3 * attn_dim, 1, bias=False)
        self.proj = ConvND(dim, attn_dim, num_channels, 1)

    def forward(self, x):
        B, C, *spatial = x.shape
        qkv = self.qkv(x).flatten(2)                               # (B, 3*attn_dim, N)
        qkv = qkv.unflatten(1, (3, self.num_heads, self.head_dim))  # (B, 3, H, Dh, N)
        q, k, v = qkv.permute(1, 0, 2, 4, 3).unbind(0)             # 3× (B, H, N, Dh)
        out = torch.nn.functional.scaled_dot_product_attention(q, k, v)  # (B, H, N, Dh)
        return self.proj(out.transpose(2, 3).reshape(B, self.num_heads * self.head_dim, *spatial))

def ConvBlockND(dim, num_channels, kernel_size=5, mlp_ratio=4, drop_prob=0.0,
                norm_layer=None, act_layer=None, conv_mlp=True):
    """ConvNeXt-style block: depthwise conv → norm → pointwise expand → act → pointwise shrink. Wrapped in Residual
    with drop path. When conv_mlp=False, the pointwise layers use nn.Linear (channels-last) instead of 1×1 conv."""
    mid = num_channels * mlp_ratio
    norm = LayerNormND(num_channels) if norm_layer is None else norm_layer()
    act = torch.nn.GELU(approximate='tanh') if act_layer is None else act_layer()
    layers = [ConvND(dim, num_channels, num_channels, kernel_size, padding=kernel_size//2, groups=num_channels), norm]
    if conv_mlp:
        layers += [ConvND(dim, num_channels, mid, 1), act, ConvND(dim, mid, num_channels, 1)]
    else:
        layers += [_ChannelsLast(), torch.nn.Linear(num_channels, mid), act,
                   torch.nn.Linear(mid, num_channels), _ChannelsFirst()]
    return Residual(torch.nn.Sequential(*layers), drop_prob)

class SAPETransformerBlockND(torch.nn.Module):
    """Pre-norm transformer block with per-layer SAPE position injection. Takes (B, num_channels, *spatial),
    internally operates at num_channels + dim channels (with dim position channels via SAPE), outputs
    (B, num_channels, *spatial). Interleaves cleanly with ConvBlockND at the same num_channels."""
    def __init__(self, dim, num_channels, num_heads, head_dim=32, mlp_ratio=4, drop_prob=0.0,
                 norm_layer=None, act_layer=None):
        super().__init__()
        inner = num_channels + dim
        self.num_channels = num_channels
        self.sape = SAPE(dim)
        self.norm1 = LayerNormND(inner) if norm_layer is None else norm_layer()
        self.norm2 = LayerNormND(inner) if norm_layer is None else norm_layer()
        self.attn = AttentionND(dim, inner, num_heads, head_dim)
        act = torch.nn.GELU(approximate='tanh') if act_layer is None else act_layer()
        mid = inner * mlp_ratio
        self.mlp = torch.nn.Sequential(
            ConvND(dim, inner, mid, 1),
            ConvND(dim, mid, mid, 3, padding=1, groups=mid),
            act,
            ConvND(dim, mid, inner, 1),
        )
        self.drop_prob = drop_prob

    def forward(self, x):
        x_pos = self.sape(x)                              # (B, C+dim, *spatial)
        if self.drop_prob and self.training:
            if torch.rand(()).item() >= self.drop_prob:
                x_pos = x_pos + self.attn(self.norm1(x_pos)) / (1.0 - self.drop_prob)
            if torch.rand(()).item() >= self.drop_prob:
                x_pos = x_pos + self.mlp(self.norm2(x_pos)) / (1.0 - self.drop_prob)
        else:
            x_pos = x_pos + self.attn(self.norm1(x_pos))
            x_pos = x_pos + self.mlp(self.norm2(x_pos))
        return x_pos[:, :self.num_channels]                # strip position