import numpy as np
import tensorflow as tf
from tensorflow.keras import layers, models, optimizers, losses, metrics, Model
import os

class CountNet(Model):
	def __init__(self, K, input_shape):
		super(CountNet, self).__init__()
		self.K = K
		self.cnn_1         = layers.Conv2D(filters=64, kernel_size=3, strides=1, padding='same', input_shape=input_shape)
		self.batch_norm_1  = layers.BatchNormalization()
		# self.cnn_2         = layers.Conv2D(filters=512, kernel_size=3, strides=1, padding='same')
		# self.batch_norm_2  = layers.BatchNormalization()
		self.cnn_3         = layers.Conv2D(filters=128, kernel_size=3, strides=1, padding='same')
		self.pool_3        = layers.MaxPooling2D(pool_size=(2, 2))
		self.batch_norm_3  = layers.BatchNormalization()
		self.cnn_4         = layers.Conv2D(filters=128, kernel_size=3, strides=1, padding='same')
		self.batch_norm_4  = layers.BatchNormalization()
		self.cnn_5         = layers.Conv2D(filters=256, kernel_size=3, strides=1, padding='same')
		self.pool_5        = layers.MaxPooling2D(pool_size=(2, 2))
		self.batch_norm_5  = layers.BatchNormalization()
		# self.cnn_6         = layers.Conv2D(filters=64, kernel_size=3, strides=1, padding='same')
		# self.batch_norm_6  = layers.BatchNormalization()
		self.cnn_7         = layers.Conv2D(filters=256, kernel_size=3, strides=1, padding='same')
		self.pool_7        = layers.MaxPooling2D(pool_size=(2, 2))
		self.batch_norm_7  = layers.BatchNormalization()
		self.cnn_8         = layers.Conv2D(filters=512, kernel_size=3, strides=1, padding='same')
		self.batch_norm_8  = layers.BatchNormalization()
		self.cnn_9         = layers.Conv2D(filters=512, kernel_size=3, strides=1, padding='same')
		self.pool_9        = layers.MaxPooling2D(pool_size=(2, 2))
		self.batch_norm_9  = layers.BatchNormalization()
		self.cnn_10        = layers.Conv2D(filters=1024, kernel_size=3, strides=1, padding='same')
		self.pool_10       = layers.MaxPooling2D(pool_size=(2, 2))
		self.batch_norm_10 = layers.BatchNormalization()
		self.cnn_11        = layers.Conv2D(filters=1024, kernel_size=3, strides=1, padding='same')
		self.pool_11       = layers.MaxPooling2D(pool_size=(2, 2))
		self.batch_norm_11 = layers.BatchNormalization()
		self.cnn_12        = layers.Conv2D(filters=2048, kernel_size=3, strides=1, padding='same')
		self.pool_12       = layers.MaxPooling2D(pool_size=(4, 4))
		self.cnn_13        = layers.Conv2D(filters=1024 * self.K, kernel_size=3, strides=1, padding='same', activation='sigmoid')
		self.cnn_14       =  layers.Conv2D(filters=self.K, kernel_size=3, strides=1, padding='same', activation='sigmoid')
		self.activation    = layers.Activation('relu')

	def call(self, x):
		x = self.cnn_1(x)
		x = self.batch_norm_1(x)
		x = self.activation(x)
		# x = self.cnn_2(x)
		# x = self.batch_norm_2(x)
		# x = self.activation(x)
		x = self.cnn_3(x) 
		x = self.pool_3(x)
		x = self.batch_norm_3(x)
		x = self.activation(x)
		x = self.cnn_4(x)  
		x = self.batch_norm_4(x)
		x = self.activation(x)
		x = self.cnn_5(x)
		x = self.pool_5(x)
		x = self.batch_norm_5(x)
		x = self.activation(x)
		# x = self.cnn_6(x)
		# x = self.batch_norm_6(x)
		# x = self.activation(x)
		x = self.cnn_7(x)
		x = self.pool_7(x)
		x = self.batch_norm_7(x)
		x = self.activation(x)
		x = self.cnn_8(x)
		x = self.batch_norm_8(x)
		x = self.activation(x)
		x = self.cnn_9(x)
		x = self.pool_9(x)
		x = self.batch_norm_9(x)
		x = self.activation(x)
		x = self.cnn_10(x)
		x = self.pool_10(x)
		x = self.batch_norm_10(x)
		x = self.activation(x)
		x = self.cnn_11(x)
		x = self.pool_11(x)
		x = self.batch_norm_11(x)
		x = self.activation(x)
		x = self.cnn_12(x)
		x = self.pool_12(x)
		x_out = self.cnn_13(x)
		prob_out = self.cnn_14(x)
		x_out = tf.reshape(x_out, (x_out.shape[0], self.K, -1))
		prob_out = tf.reshape(prob_out, (prob_out.shape[0], self.K))
		# print(x_out.shape, prob_out.shape)
		return x_out, prob_out