# -*- coding: utf-8 -*-
"""
Created on Tue Jun 21 10:53:51 2022
model
@author: fangjiasheng
"""
from keras.layers import Lambda, Dense, Reshape, Conv2D, BatchNormalization, LeakyReLU, Masking, MaxPool2D
from keras import layers, models, Sequential
import keras.backend as K


def direction_model(input_shape, output_shape):
    conv_num = 6

    # Input
    _input = layers.Input(shape=input_shape, dtype="float32", name="input")

    conv = Conv2D(16, (3, 3), padding='same')(_input)
    bn = BatchNormalization()(conv)
    relu = LeakyReLU(alpha=0.)(bn)
    max_pool = MaxPool2D()(relu)
    for i in range(conv_num):
        conv = Conv2D(16, (3, 3), padding='same')(max_pool)
        bn = BatchNormalization()(conv)
        relu = LeakyReLU(alpha=0.)(bn)
        max_pool = MaxPool2D()(relu)
    conv = Conv2D(16, (3, 3), padding='same')(max_pool)
    bn = BatchNormalization()(conv)
    relu = LeakyReLU(alpha=0.)(bn)
    max_pool = MaxPool2D((5, 5))(relu)

    dense = layers.Dense(output_shape, activation='softmax')(max_pool)
    squeeze = Lambda(lambda x: K.squeeze(x, axis=1))(dense)
    squeeze = Lambda(lambda x: K.squeeze(x, axis=1))(squeeze)

    model = models.Model(inputs=_input, outputs=squeeze)
    # model.summary(line_length=100)
    return model