"""
Vanilla U-Net for RockGlacierNet — baseline model.

Standard Ronneberger et al. (2015) U-Net: encoder/decoder with plain skip
connections (NO attention gates, NO transformer blocks). Kept architecturally
identical to the Attention U-Net (same filters/depth) so the ONLY difference is
the attention mechanism — a fair baseline comparison.

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

Encoder filters : 64 -> 128 -> 256 -> 512
Bottleneck      : 512
Decoder filters : 256 -> 128 -> 64 -> 32
"""

import tensorflow as tf
from tensorflow.keras.layers import (
    Input, Conv2D, Conv2DTranspose, MaxPooling2D,
    BatchNormalization, Activation, concatenate
)
from tensorflow.keras.models import Model


def _double_conv(x, filters, dropout=0.0, name=""):
    """Two Conv3x3 + BN + ReLU blocks (standard U-Net block)."""
    x = Conv2D(filters, 3, padding='same', name=f'{name}_c1')(x)
    x = BatchNormalization(name=f'{name}_bn1')(x)
    x = Activation('relu', name=f'{name}_r1')(x)
    if dropout > 0:
        x = tf.keras.layers.Dropout(dropout, name=f'{name}_dp')(x)
    x = Conv2D(filters, 3, padding='same', name=f'{name}_c2')(x)
    x = BatchNormalization(name=f'{name}_bn2')(x)
    x = Activation('relu', name=f'{name}_r2')(x)
    return x


def build_unet(input_shape=(128, 128, 12), num_classes=1):
    """Plain U-Net (no attention). Signature matches build_attention_unet."""
    inputs = Input(input_shape, name='input')

    # Encoder
    e1 = _double_conv(inputs, 64,  dropout=0.1, name='enc1')   # 128^2,  64
    p1 = MaxPooling2D(2, name='pool1')(e1)
    e2 = _double_conv(p1,     128, dropout=0.1, name='enc2')   # 64^2,  128
    p2 = MaxPooling2D(2, name='pool2')(e2)
    e3 = _double_conv(p2,     256, dropout=0.2, name='enc3')   # 32^2,  256
    p3 = MaxPooling2D(2, name='pool3')(e3)
    e4 = _double_conv(p3,     512, dropout=0.2, name='enc4')   # 16^2,  512
    p4 = MaxPooling2D(2, name='pool4')(e4)

    # Bottleneck
    b  = _double_conv(p4,     512, dropout=0.3, name='btn')    # 8^2,   512

    # Decoder (plain skip concatenation, no attention gates)
    u4 = Conv2DTranspose(256, 2, strides=2, padding='same', name='up4')(b)
    d4 = _double_conv(concatenate([e4, u4], name='cat4'), 256, dropout=0.2, name='dec4')
    u3 = Conv2DTranspose(128, 2, strides=2, padding='same', name='up3')(d4)
    d3 = _double_conv(concatenate([e3, u3], name='cat3'), 128, dropout=0.2, name='dec3')
    u2 = Conv2DTranspose(64, 2, strides=2, padding='same', name='up2')(d3)
    d2 = _double_conv(concatenate([e2, u2], name='cat2'), 64,  dropout=0.1, name='dec2')
    u1 = Conv2DTranspose(32, 2, strides=2, padding='same', name='up1')(d2)
    d1 = _double_conv(concatenate([e1, u1], name='cat1'), 32,  dropout=0.1, name='dec1')

    outputs = Conv2D(num_classes, 1, activation='sigmoid', name='output')(d1)
    return Model(inputs, outputs, name='UNet')