"""
SegFormer (MiT-B0 scale) for RockGlacierNet — transformer baseline.

Faithful compact implementation of Xie et al. (2021) "SegFormer":
  - Hierarchical Transformer encoder (4 stages, overlap patch embed)
  - Efficient self-attention (spatial-reduction) + Mix-FFN (depthwise conv)
  - Lightweight all-MLP decoder that fuses the 4 multi-scale features

Input : (128, 128, N_BANDS)
Output: (128, 128, 1) sigmoid — binary glacier mask

Encoder dims  : 32 -> 64 -> 160 -> 256   (MiT-B0)
Decoder dim   : 256
"""

import tensorflow as tf
from tensorflow.keras import layers, Model


def _overlap_patch_embed(x, out_ch, patch, stride, name):
    x = layers.Conv2D(out_ch, patch, strides=stride, padding='same', name=f'{name}_proj')(x)
    x = layers.LayerNormalization(epsilon=1e-6, name=f'{name}_ln')(x)
    return x


def _efficient_attention(x, dim, num_heads, sr_ratio, name):
    """Spatial-reduction self-attention on a (B, H, W, C) tensor."""
    H, W = x.shape[1], x.shape[2]
    q = layers.Reshape((H * W, dim), name=f'{name}_qseq')(x)
    if sr_ratio > 1:
        kv = layers.Conv2D(dim, sr_ratio, strides=sr_ratio, padding='same',
                           name=f'{name}_sr')(x)
        kv = layers.LayerNormalization(epsilon=1e-6, name=f'{name}_srln')(kv)
        h2, w2 = kv.shape[1], kv.shape[2]
        kv = layers.Reshape((h2 * w2, dim), name=f'{name}_kvseq')(kv)
    else:
        kv = q
    attn = layers.MultiHeadAttention(num_heads=num_heads, key_dim=max(dim // num_heads, 1),
                                     name=f'{name}_mha')(q, kv)
    attn = layers.Reshape((H, W, dim), name=f'{name}_out')(attn)
    return attn


def _mix_ffn(x, dim, expansion, name):
    h = layers.Dense(dim * expansion, name=f'{name}_fc1')(x)
    h = layers.DepthwiseConv2D(3, padding='same', name=f'{name}_dw')(h)
    h = layers.Activation('gelu', name=f'{name}_gelu')(h)
    h = layers.Dense(dim, name=f'{name}_fc2')(h)
    return h


def _transformer_block(x, dim, num_heads, sr_ratio, name):
    a = layers.LayerNormalization(epsilon=1e-6, name=f'{name}_ln1')(x)
    a = _efficient_attention(a, dim, num_heads, sr_ratio, f'{name}_attn')
    x = layers.Add(name=f'{name}_add1')([x, a])
    f = layers.LayerNormalization(epsilon=1e-6, name=f'{name}_ln2')(x)
    f = _mix_ffn(f, dim, 4, f'{name}_ffn')
    x = layers.Add(name=f'{name}_add2')([x, f])
    return x


def build_segformer(input_shape=(128, 128, 12), num_classes=1):
    """SegFormer MiT-B0. Signature matches the other build_* functions."""
    dims       = [32, 64, 160, 256]
    heads      = [1, 2, 5, 8]
    sr_ratios  = [8, 4, 2, 1]
    depths     = [2, 2, 2, 2]
    patch_cfg  = [(7, 4), (3, 2), (3, 2), (3, 2)]   # (kernel, stride) per stage
    dec_dim    = 256

    inp = layers.Input(input_shape, name='input')
    x = inp
    feats = []
    for i in range(4):
        p, s = patch_cfg[i]
        x = _overlap_patch_embed(x, dims[i], p, s, f'stage{i+1}_embed')
        for d in range(depths[i]):
            x = _transformer_block(x, dims[i], heads[i], sr_ratios[i], f'stage{i+1}_blk{d}')
        x = layers.LayerNormalization(epsilon=1e-6, name=f'stage{i+1}_norm')(x)
        feats.append(x)

    # All-MLP decoder
    target_hw = feats[0].shape[1]          # 1/4 resolution (32 for 128 input)
    ups = []
    for i, f in enumerate(feats):
        u = layers.Dense(dec_dim, name=f'dec_proj{i}')(f)
        scale = target_hw // f.shape[1]
        if scale > 1:
            u = layers.UpSampling2D(size=scale, interpolation='bilinear',
                                    name=f'dec_up{i}')(u)
        ups.append(u)
    x = layers.Concatenate(name='dec_concat')(ups)
    x = layers.Conv2D(dec_dim, 1, name='dec_fuse')(x)
    x = layers.BatchNormalization(name='dec_bn')(x)
    x = layers.Activation('relu', name='dec_relu')(x)
    x = layers.Dropout(0.1, name='dec_drop')(x)

    x = layers.UpSampling2D(size=input_shape[0] // target_hw, interpolation='bilinear',
                            name='final_up')(x)
    out = layers.Conv2D(num_classes, 1, activation='sigmoid', name='output')(x)
    return Model(inp, out, name='SegFormer')