text: """
  构建CycleGAN需要的各个模块
  """
  import tensorflow as tf
  import os 
  import sys
  base = os.path.dirname(os.path.abspath(__file__))
  sys.path.append(os.path.join(base, '../../'))
  # import basemodels.BaseLayers_UnMixed as BaseLayers
  import basemodels.BaseLayers_Mixed as BaseLayers
  ###########################################################Generator blocks###########################################################
  class c7s1_k_2D(tf.keras.Model):
      def __init__(self,filters,sn=False,use_bias=False,activation="relu",**kwargs):
          super(c7s1_k_2D,self).__init__()
          self.filters_zp = filters
          self.l1_conv2d = BaseLayers.Conv2D(filters=filters,kernel_size=[7,7],strides=[1,1],padding="REFLECT",use_bias=use_bias,sn=sn,**kwargs)
          self.l2_norm = BaseLayers.InstanceNorm2D()
          if activation.lower()=="relu":
              self.l3_activation = BaseLayers.ReLU()
          elif activation.lower()=="sigmoid":
              self.l3_activation = tf.keras.layers.Activation("sigmoid")
          elif activation.lower()=="tanh":
              self.l3_activation = tf.keras.layers.Activation("tanh")
          else:
              raise ValueError("Un supported activation "+activation)
          
      def build(self,input_shape):
          self.l1_conv2d.build(input_shape=input_shape)
          input_shape[-1]=int(self.filters_zp)
          self.l2_norm.build(input_shape=input_shape)
          self.built = True
      def call(self,x,training=True):
          x = self.l1_conv2d(x,training=training)
          x = self.l2_norm(x,training=training)
          y = self.l3_activation(x,training=training)
          return y
  class dk_2D(tf.keras.Model):
      def __init__(self,filters,sn=False,use_bias=False,**kwargs):
          super(dk_2D,self).__init__()
          self.filters_zp = filters
          self.l1_conv2d = BaseLayers.Conv2D(filters=filters,kernel_size=[3,3],strides=[2,2],padding="REFLECT",use_bias=use_bias,sn=sn,**kwargs)
          self.l2_norm = BaseLayers.InstanceNorm2D()
          self.l3_activation = BaseLayers.ReLU()
      def build(self,input_shape):
          self.l1_conv2d.build(input_shape=input_shape)
          for i in range(2):
              input_shape[i+1] = int(tf.math.ceil(input_shape[i+1]/2))
          input_shape[-1]=int(self.filters_zp)
          self.l2_norm.build(input_shape=input_shape)
          self.built = True
      def call(self,x,training=True):
          x = self.l1_conv2d(x,training=training)
          x = self.l2_norm(x,training=training)
          y = self.l3_activation(x,training=training)
          return y
  class rk_2D(tf.keras.Model):
      def __init__(self,filters,sn=False,use_bias=False,**kwargs):
          super(rk_2D,self).__init__()
          self.l1_conv2d= BaseLayers.Conv2D(filters=filters,kernel_size=[3,3],strides=[1,1],padding="REFLECT",use_bias=use_bias,sn=sn,**kwargs)
          self.l2_norm = BaseLayers.InstanceNorm2D()
          self.l3_activation = BaseLayers.ReLU()
          self.l4_conv2d= BaseLayers.Conv2D(filters=filters,kernel_size=[3,3],strides=[1,1],padding="REFLECT",use_bias=use_bias,sn=sn,**kwargs)
          self.l5_norm = BaseLayers.InstanceNorm2D()
      def build(self,input_shape):
          self.l1_conv2d.build(input_shape=input_shape)
          self.l2_norm.build(input_shape=input_shape)
          self.l4_conv2d.build(input_shape=input_shape)
          self.l5_norm.build(input_shape=input_shape)
          self.built = True
      def call(self,x,training=True):
          y = self.l1_conv2d(x,training=training)
          y = self.l2_norm(y,training=training)
          y = self.l3_activation(y,training=training)
          y = self.l4_conv2d(y,training=training)
          y = self.l5_norm(y,training=training)
          return y+x
  
  class n_res_blocks_2D(tf.keras.Model):
      def __init__(self,filters,sn=False,use_bias=False,n=6,**kwargs):
          super(n_res_blocks_2D,self).__init__()
          self.rs_list = []
          for _ in range(n):
              self.rs_list.append(rk_2D(filters=filters,sn=sn,use_bias=use_bias,**kwargs))
      def build(self,input_shape):
          for item in self.rs_list:
              item.build(input_shape=input_shape)
          self.built = True
      def call(self,x,training=True):
          in_put = x
          for item in self.rs_list:
              out_put = item(in_put,training=training)
              in_put = out_put
          return out_put
  
  class uk_2D(tf.keras.Model):
      def __init__(self,filters,sn=False,use_bias=False,**kwargs):
          super(uk_2D,self).__init__()
          self.filters_zp = filters
          self.l1_up= BaseLayers.UpSample2D(filters=filters,kernel_size=[3,3],strides=[2,2],padding="SAME",use_bias=use_bias,sn=sn,**kwargs)
          self.l2_norm = BaseLayers.InstanceNorm2D()
          self.l3_activation = BaseLayers.ReLU()
      def build(self,input_shape):
          self.l1_up.build(input_shape=input_shape)
          for i in range(3):
              input_shape[i+1] = input_shape[i+1]*2
          input_shape[-1]=int(self.filters_zp)    
          self.l2_norm.build(input_shape=input_shape)
          self.built = True
      def call(self,x,training=True):
          x = self.l1_up(x,training=training)
          x = self.l2_norm(x,training=training)
          y = self.l3_activation(x,training=training)
          return y
  ###########################################################Discriminator blocks###########################################################
  class ck_2D(tf.keras.Model):
      def __init__(self,filters,sn=False,norm=True,use_bias=False,**kwargs):
          super(ck_2D,self).__init__()
          self.filters_zp = filters
          self.l1_conv2d= BaseLayers.Conv2D(filters=filters,kernel_size=[4,4],strides=[2,2],padding='SAME',use_bias=use_bias,sn=sn,**kwargs)
          if norm:
              self.l2_norm = BaseLayers.InstanceNorm2D()
          else:
              self.l2_norm = tf.keras.layers.Activation("linear")
          self.l3_activation = BaseLayers.LeakyReLU(alpha=0.2)
      def build(self,input_shape):
          self.l1_conv2d.build(input_shape=input_shape)
          for i in range(2):
              input_shape[i+1] = int(tf.math.ceil(input_shape[i+1]/2))
          input_shape[-1]=int(self.filters_zp)
          self.l2_norm.build(input_shape=input_shape)
          self.built = True
      def call(self,x,training=True):
          x = self.l1_conv2d(x,training=training)
          x = self.l2_norm(x,training=training)
          y = self.l3_activation(x,training=training)
          return y
  class last_conv_2D(tf.keras.Model):
      def __init__(self,use_sigmoid=True,sn=False,**kwargs):
          super(last_conv_2D,self).__init__()
          self.l1_conv2d= BaseLayers.Conv2D(filters=1,kernel_size=[4,4],strides=[1,1],padding='SAME',use_bias=True,sn=sn,**kwargs)
          if use_sigmoid:
              self.l2_activation = tf.keras.layers.Activation("sigmoid")
          else:
              self.l2_activation = tf.keras.layers.Activation("linear")
      def build(self,input_shape):
          self.l1_conv2d.build(input_shape=input_shape)
          self.built = True
      def call(self,x,training=True):
          x = self.l1_conv2d(x,training=training)
          y = self.l2_activation(x,training=training)
          return y
  if __name__ == "__main__":
      physical_devices = tf.config.experimental.list_physical_devices(device_type='GPU')
      tf.config.experimental.set_memory_growth(physical_devices[0], True)
      from tensorflow.keras.mixed_precision import experimental as mixed_precision
      policy = mixed_precision.Policy('mixed_float16')
      mixed_precision.set_policy(policy)
      a = c7s1_k_2D(6,sn=True,use_bias=True)
      a.build(input_shape=[None,256,256,3])
      w = tf.random.normal([1,256,256,3])
      y=a(w)
      w = tf.random.normal([1,256,256,3])
      y=a(w)
      w = tf.random.normal([1,256,256,3])
      y=a(w)
      print(y.shape)
      print(y.dtype)
      print(len(a.trainable_variables))
      a = dk_2D(64)#(256+1*2-3)//2 +1=128
      a.build(input_shape=[None,256,256,3])
      w = tf.constant(tf.random.normal(([1,256,256,3])))
      print(a(w).shape)
      print(a(w).dtype)
      print(len(a.trainable_variables))
      a = rk_2D(32)#(256+1*2-3)//2 +1=128
      a.build(input_shape=[None,256,256,32])
      w = tf.constant(tf.random.normal(([1,256,256,32])))
      print(a(w).shape)
      print(a(w).dtype)
      print(len(a.trainable_variables))
      a = n_res_blocks_2D(32,n=9)
      a.build(input_shape=[None,256,256,32])
      w = tf.constant(tf.random.normal([1,256,256,32]))
      print(a(w).shape)
      print(a(w).dtype)
      print(len(a.trainable_variables))
      a = n_res_blocks_2D(32,n=9,use_bias=True,sn=False)
      a.build(input_shape=[None,256,256,32])
      w = tf.constant(tf.random.normal([1,256,256,32]))
      print(a(w).shape)
      print(a(w).dtype)
      print(len(a.trainable_variables))
      a = n_res_blocks_2D(32,n=9,use_bias=True,sn=True)
      a.build(input_shape=[None,256,256,32])
      w = tf.constant(tf.random.normal([1,256,256,32]))
      print(a(w).shape)
      print(a(w).dtype)
      print(len(a.trainable_variables))
      a = uk_2D(5)
      a.build(input_shape=[None,256,256,3])
      w = tf.constant(tf.zeros((1,256,256,3)))#生成随机值会溢出 所以用全0值代替
      print(a(w).shape)
      print(a(w).dtype)
      print(len(a.trainable_variables))
  
      a = ck_2D(64)
      a.build(input_shape=[None,256,256,3])
      w = tf.constant(tf.random.normal((1,256,256,3)))
      print(a(w).shape)
      print(a(w).dtype)
      print(len(a.trainable_variables))  
  
      a = ck_2D(64,norm=False)
      a.build(input_shape=[None,256,256,3])
      w = tf.constant(tf.random.normal((1,256,256,3)))
      print(a(w).shape)
      print(a(w).dtype)
      print(len(a.trainable_variables))  
      a = last_conv_2D(use_sigmoid=True)
      a.build(input_shape=[None,512,512,3])
      w = tf.constant(tf.random.normal((1,512,512,3)))
      print(a(w).shape)
      print(a(w).dtype)
      print(len(a.trainable_variables))
      # a = last_conv_2D(use_sigmoid=False)
      # a.build(input_shape=[None,512,512,3])
      # w = tf.constant(tf.random.normal((1,512,512,3)))
      # print(a(w).shape)
      # print(a(w).dtype)
      # print(len(a.trainable_variables))
      # a = last_conv_2D(use_sigmoid=False,sn=True)
      # a.build(input_shape=[None,512,512,3])
      # w = tf.constant(tf.random.normal((1,512,512,3)))
      # print(a(w).shape)
      # print(a(w).dtype)
      # print(len(a.trainable_variables))
  
      # a = last_conv_2D(use_sigmoid=False,sn=False)
      # a.build(input_shape=[None,512,512,3])
      # w = tf.constant(tf.random.normal((1,512,512,3)))
      # print(a(w).shape)
      # print(a(w).dtype)
      # print(len(a.trainable_variables))
  
  
  
  
  