import os
import tensorflow as tf
from config import *
from tensorflow.keras import Model,layers

os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
# custom layer for reshaping last layer

class yolo_dynamic_reshape(tf.keras.layers.Layer):
  def __init__(self, **kwargs):
    super(yolo_dynamic_reshape,self).__init__()
    self.scaler=32
    self.last_item=(5+len(class_names))
  def call(self,input_layer,output_layer):
    # print(input_layer.shape)
    if input_layer.shape[1] is not None:
      return layers.Reshape(((input_layer.shape[1]//self.scaler),(input_layer.shape[2]//self.scaler),num_anchors,self.last_item))(output_layer)
    return layers.Reshape((0,0,num_anchors,self.last_item))(output_layer)

def create_model():

    def space_to_depth_x2(x):
        return tf.nn.space_to_depth(x,block_size=2)

    x_input=layers.Input(shape=(416,416,3))
    x=layers.Lambda(lambda x:x/255.)(x_input)
    x=layers.Conv2D(32,(3,3),strides=(1,1),padding='same',name='conv_1',use_bias=False)(x)
    x=layers.BatchNormalization(name='norm_1')(x)
    x=layers.LeakyReLU(alpha=0.1)(x)
    x=layers.MaxPooling2D(pool_size=(2,2))(x)

    x=layers.Conv2D(64,(3,3),strides=(1,1),padding='same',name='conv_2',use_bias=False)(x)
    x=layers.BatchNormalization(name='norm_2')(x)
    x=layers.LeakyReLU(alpha=0.1)(x)
    x=layers.MaxPooling2D(pool_size=(2,2))(x)

    x=layers.Conv2D(128,(3,3),strides=(1,1),padding='same',name='conv_3',use_bias=False)(x)
    x=layers.BatchNormalization(name='norm_3')(x)
    x=layers.LeakyReLU(alpha=0.1)(x)

    x=layers.Conv2D(64,(1,1),strides=(1,1),padding='same',name='conv_4',use_bias=False)(x)
    x=layers.BatchNormalization(name='norm_4')(x)
    x=layers.LeakyReLU(alpha=0.1)(x)

    x=layers.Conv2D(128,(3,3),strides=(1,1),padding='same',name='conv_5',use_bias=False)(x)
    x=layers.BatchNormalization(name='norm_5')(x)
    x=layers.LeakyReLU(alpha=0.1)(x)
    x=layers.MaxPooling2D(pool_size=(2,2))(x)

    x=layers.Conv2D(256,(3,3),strides=(1,1),padding='same',name='conv_6',use_bias=False)(x)
    x=layers.BatchNormalization(name='norm_6')(x)
    x=layers.LeakyReLU(alpha=0.1)(x)

    x=layers.Conv2D(128,(1,1),strides=(1,1),padding='same',name='conv_7',use_bias=False)(x)
    x=layers.BatchNormalization(name='norm_7')(x)
    x=layers.LeakyReLU(alpha=0.1)(x)

    x=layers.Conv2D(256,(3,3),strides=(1,1),padding='same',name='conv_8',use_bias=False)(x)
    x=layers.BatchNormalization(name='norm_8')(x)
    x=layers.LeakyReLU(alpha=0.1)(x)
    x=layers.MaxPooling2D(pool_size=(2,2))(x)

    x=layers.Conv2D(512,(3,3),strides=(1,1),padding='same',name='conv_9',use_bias=False)(x)
    x=layers.BatchNormalization(name='norm_9')(x)
    x=layers.LeakyReLU(alpha=0.1)(x)

    x=layers.Conv2D(256,(1,1),strides=(1,1),padding='same',name='conv_10',use_bias=False)(x)
    x=layers.BatchNormalization(name='norm_10')(x)
    x=layers.LeakyReLU(alpha=0.1)(x)

    x=layers.Conv2D(512,(3,3),strides=(1,1),padding='same',name='conv_11',use_bias=False)(x)
    x=layers.BatchNormalization(name='norm_11')(x)
    x=layers.LeakyReLU(alpha=0.1)(x)

    x=layers.Conv2D(256,(1,1),strides=(1,1),padding='same',name='conv_12',use_bias=False)(x)
    x=layers.BatchNormalization(name='norm_12')(x)
    x=layers.LeakyReLU(alpha=0.1)(x)

    x=layers.Conv2D(512,(3,3),strides=(1,1),padding='same',name='conv_13',use_bias=False)(x)
    x=layers.BatchNormalization(name='norm_13')(x)
    x=layers.LeakyReLU(alpha=0.1)(x)

    skip_connection = x

    x=layers.MaxPooling2D(pool_size=(2,2))(x)

    x=layers.Conv2D(1024,(3,3),strides=(1,1),padding='same',name='conv_14',use_bias=False)(x)
    x=layers.BatchNormalization(name='norm_14')(x)
    x=layers.LeakyReLU(alpha=0.1)(x)

    x=layers.Conv2D(512,(1,1),strides=(1,1),padding='same',name='conv_15',use_bias=False)(x)
    x=layers.BatchNormalization(name='norm_15')(x)
    x=layers.LeakyReLU(alpha=0.1)(x)

    x=layers.Conv2D(1024,(3,3),strides=(1,1),padding='same',name='conv_16',use_bias=False)(x)
    x=layers.BatchNormalization(name='norm_16')(x)
    x=layers.LeakyReLU(alpha=0.1)(x)

    x=layers.Conv2D(512,(1,1),strides=(1,1),padding='same',name='conv_17',use_bias=False)(x)
    x=layers.BatchNormalization(name='norm_17')(x)
    x=layers.LeakyReLU(alpha=0.1)(x)

    x=layers.Conv2D(1024,(3,3),strides=(1,1),padding='same',name='conv_18',use_bias=False)(x)
    x=layers.BatchNormalization(name='norm_18')(x)
    x=layers.LeakyReLU(alpha=0.1)(x)

    x=layers.Conv2D(1024,(3,3),strides=(1,1),padding='same',name='conv_19',use_bias=False)(x)
    x=layers.BatchNormalization(name='norm_19')(x)
    x=layers.LeakyReLU(alpha=0.1)(x)

    x=layers.Conv2D(1024,(3,3),strides=(1,1),padding='same',name='conv_20',use_bias=False)(x)
    x=layers.BatchNormalization(name='norm_20')(x)
    x=layers.LeakyReLU(alpha=0.1)(x)

    skip_connection=layers.Conv2D(64,(1,1),strides=(1,1),padding='same',name='conv_21',use_bias=False)(skip_connection)
    skip_connection=layers.BatchNormalization(name='norm_21')(skip_connection)
    skip_connection=layers.LeakyReLU(alpha=0.1)(skip_connection)
    skip_connection=layers.Lambda(space_to_depth_x2)(skip_connection) # halfs the resolution and add more depth

    x=layers.concatenate([skip_connection,x])
    x=layers.Conv2D(1024,(3,3),strides=(1,1),padding='same',name='conv_22',use_bias=False)(x)
    x=layers.BatchNormalization(name='norm_22')(x)
    x=layers.LeakyReLU(alpha=0.1)(x)

    x=layers.Conv2D((num_anchors*(5+len(class_names))),(1,1),strides=(1,1),padding='same',name='conv_23')(x)
    out=yolo_dynamic_reshape()(x_input,x)

    model=Model(x_input,out,name='yolo_v2_model')
    # model.summary()
    return model

def load_model(path):
    # model=tf.keras.models.load_model(path,custom_objects={'yolo_dynamic_reshape':yolo_dynamic_reshape},compile=False)
    model=create_model()
    model.load_weights(path)
    # model=tf.keras.models.load_model(path)
    return model