# DeepLabV3+ Model
import tensorflow as tf
from keras.layers import Input
from keras.layers import Conv2D
from keras.layers import Concatenate
from keras.layers import UpSampling2D
from keras.models import Model
from tensorflow.keras.applications import ResNet50

from .config import IMAGE_SIZE, LEARNING_RATE, LLF_FILTERS, RESNET_WEIGHTS
from .layers import ConvBlock, AtrousSpatialPyramidPooling


def create_model(image_size=IMAGE_SIZE, learning_rate=LEARNING_RATE,
                 compile_model=True, in_channels=4):
    """
    DeepLabV3+ 모델 생성

    Args:
        image_size: 입력 이미지 크기
        learning_rate: 학습률
        compile_model: True면 모델 컴파일 포함
        in_channels: 입력 채널 수 (4 = R,G,B,MNDWI)

    Returns:
        DeepLabV3+ Keras Model
    """
    # Input Layer
    InputL = Input(shape=(image_size, image_size, in_channels), name="InputLayer")

    # 4채널 → 3채널 변환 (ResNet50은 3채널만 지원)
    if in_channels != 3:
        x = Conv2D(3, kernel_size=1, padding='same', name="ChannelReduce")(InputL)
    else:
        x = InputL

    # Backbone: ResNet50 (pretrained on ImageNet)
    # 별도 input_shape로 생성 후 중간 레이어 출력을 추출
    resnet50_base = ResNet50(
        include_top=False,
        weights=RESNET_WEIGHTS,
        input_shape=(image_size, image_size, 3)
    )
    # ResNet50의 중간 출력을 가져오기 위한 멀티출력 모델 생성
    resnet50 = Model(
        inputs=resnet50_base.input,
        outputs={
            'deep': resnet50_base.get_layer('conv4_block6_2_relu').output,
            'low': resnet50_base.get_layer('conv2_block3_2_relu').output,
        },
        name="ResNet50-Backbone"
    )
    features_dict = resnet50(x)

    # ASPP Phase - Deep CNN features
    DCNN = features_dict['deep']
    ASPP = AtrousSpatialPyramidPooling(DCNN)
    ASPP = UpSampling2D(
        size=(image_size // 4 // ASPP.shape[1], image_size // 4 // ASPP.shape[2]),
        name="AtrousSpatial"
    )(ASPP)

    # Low-Level Features (LLF) Phase
    LLF = features_dict['low']
    LLF = ConvBlock(filters=LLF_FILTERS, kernel_size=1, name="LLF-ConvBlock")(LLF)

    # Decoder: Combine ASPP and LLF
    combined = Concatenate(axis=-1, name="Combine-LLF-ASPP")([ASPP, LLF])
    features = ConvBlock(name="Top-ConvBlock-1")(combined)
    features = ConvBlock(name="Top-ConvBlock-2")(features)

    # Upsample to original size
    upsample = UpSampling2D(
        size=(image_size // features.shape[1], image_size // features.shape[1]),
        interpolation='bilinear',
        name="Top-UpSample"
    )(features)

    # Output Mask (1 channel for binary segmentation)
    PredMask = Conv2D(
        1,
        kernel_size=3,
        strides=1,
        padding='same',
        activation='sigmoid',
        use_bias=False,
        name="OutputMask"
    )(upsample)

    # Build Model
    model = Model(InputL, PredMask, name="DeepLabV3-Plus")

    # Compile if requested
    if compile_model:
        optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)
        model.compile(loss='binary_crossentropy', optimizer=optimizer)

    return model


def load_model(model_path):
    """
    저장된 모델 로드

    Args:
        model_path: 모델 파일 경로 (.h5)

    Returns:
        로드된 Keras Model
    """
    from keras.models import load_model as keras_load_model
    from .layers import ConvBlock

    custom_objects = {
        'ConvBlock': ConvBlock
    }

    return keras_load_model(model_path, custom_objects=custom_objects)