# -*- coding: utf-8 -*-
# Author: allen
# Copyright (c) 2022 RuiCare All rights reserved.
from keras import Input, Model
from tensorflow.python.keras.layers import Activation
from tensorflow.keras.mixed_precision import experimental as mixed_precision
from keras.layers import Conv3D, UpSampling3D, concatenate, Dropout, MaxPooling3D, BatchNormalization


def down_sample(tensor, filters, post_active=True, pool=True, kernel_size=(3, 3, 3), strides=(1, 1, 1),
                dropout=0.3):
    """down sample conv unit"""
    tensor_new = Conv3D(filters, kernel_size, strides=strides, padding='same')(tensor)
    # tensor_new = Dropout(rate=dropout)(tensor_new)
    tensor_new = BatchNormalization(axis=-1)(tensor_new)
    tensor_new = Activation('relu')(tensor_new)
    tensor_new = Conv3D(filters, kernel_size, strides=strides, padding='same')(tensor_new)
    # tensor_new = Dropout(rate=dropout)(tensor_new)
    tensor_new = BatchNormalization(axis=-1)(tensor_new)
    tensor_before = Activation('relu')(tensor_new) if post_active else tensor_new
    tensor_after = MaxPooling3D(pool_size=(2, 2, 2))(tensor_before) if pool else tensor_before
    return tensor_before, tensor_after


def up_sample(tensor_down, tensor_up, filters, post_active=True, kernel_size=(3, 3, 3), strides=(1, 1, 1),
              dropout=0.3):
    tensor_new = UpSampling3D(size=(2, 2, 2))(tensor_down)
    tensor_new = concatenate([tensor_new, tensor_up], axis=-1)
    tensor_new = Conv3D(filters, kernel_size, strides=strides, padding='same')(tensor_new)
    # tensor_new = Dropout(rate=dropout)(tensor_new)
    tensor_new = BatchNormalization(axis=-1)(tensor_new)
    tensor_new = Activation('relu')(tensor_new)
    tensor_new = Conv3D(filters, kernel_size, strides=strides, padding='same')(tensor_new)
    # tensor_new = Dropout(rate=dropout)(tensor_new)
    tensor_new = BatchNormalization(axis=-1)(tensor_new)
    tensor_new = Activation('relu')(tensor_new) if post_active else tensor_new
    return tensor_new


def load_unet3d(input_shape=(16, 256, 256, 4), classes=1, activation='linear'):
    policy = mixed_precision.Policy('mixed_float16')
    mixed_precision.set_policy(policy)

    tensor_input = Input(shape=input_shape)  # 32,256,256,2

    # left1_before, left1_after = down_sample(tensor_input, 32)  # 16,256,256,32
    left2_before, left2_after = down_sample(tensor_input, 64)  # 8,128,128,64
    left3_before, left3_after = down_sample(left2_after, 128)  # 4,64,64,128
    left4_before, left4_after = down_sample(left3_after, 256)  # 2,32,32,256

    mid, _ = down_sample(left4_after, 512, pool=False)  # 2,32,32,512

    right1 = up_sample(mid, left4_before, 256)  # 4,64,64,256
    right2 = up_sample(right1, left3_before, 128)  # 8,128,128,128
    right3 = up_sample(right2, left2_before, 64)  # 16,256,256,64
    # right4 = up_sample(right3, left1_before, 32)  # 32,512,512,32

    tensor_output = Conv3D(classes, (1, 1, 1), activation=activation)(right3)
    model = Model(inputs=tensor_input, outputs=tensor_output)
    return model


if __name__ == '__main__':
    net = load_unet3d(input_shape=(16, 256, 256, 4), classes=1)
    print(net)
