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)

def LayerNorm(num_channels):
    """GroupNorm with 1 group (equivalent to LayerNorm for channels-first tensors)."""
    return torch.nn.GroupNorm(num_groups=1, num_channels=num_channels)

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 SAPEOverwrite(torch.nn.Module):
    """Overwrites the last `dim` channels with centered integer coordinates (in patch units) for each spatial dimension. Apply before
    each attention layer to inject fresh position information. No learnable parameters. See SAPE for the concatenation variant."""
    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[:, :-self.dim], 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):
    """ConvNeXt-style block: depthwise conv → norm → pointwise expand → act → pointwise shrink. Wrapped in Residual with drop path. Default norm is SoftsignCompanding(qat=False), default act is GELU."""
    mid = num_channels * mlp_ratio
    norm = LayerNorm(num_channels) if norm_layer is None else norm_layer()
    act = torch.nn.GELU(approximate='tanh') if act_layer is None else act_layer()
    return Residual(torch.nn.Sequential(
        ConvND(dim, num_channels, num_channels, kernel_size, padding=kernel_size//2, groups=num_channels),
        norm,
        ConvND(dim, num_channels, mid, 1),
        act,
        ConvND(dim, mid, num_channels, 1),
    ), 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 = LayerNorm(inner) if norm_layer is None else norm_layer()
        self.norm2 = LayerNorm(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