import tensorflow as tf
print(tf.__version__)





import tensorflow as tf
import tensorflow_addons as tfa
import numpy as np
from tensorflow import keras
from tensorflow.keras import layers
from imutils import paths
import os
from sklearn.preprocessing import LabelBinarizer
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report
import cv2
from tensorflow.keras.models import Sequential
from tensorflow.keras.callbacks import ModelCheckpoint
from tqdm.notebook import tqdm
from imutils import paths
import tensorflow_datasets as tfds
import matplotlib.pyplot as plt
from tensorflow.keras.layers import Dense, GlobalAveragePooling2D,Dropout,Conv2D,BatchNormalization,Activation








#num of classes 
num_classes = 2

#input dim
dim = 480

#input shape
input_shape = (dim, dim, 3)

# image paths
train_images = list(paths.list_images("dataset/train1"))

filenames = []







unwanted_classes = ['.ipynb_checkpoints']

def prepare_images(image_paths):
    images = []
    labels = []

    for image in tqdm(image_paths):

        img_path = image.split("/")
        img_folder_name = img_path[-2]
        img_file_name = img_path[-1]

        if img_folder_name in unwanted_classes or img_file_name in unwanted_classes :
            continue
        
        try :
            if img_file_name.endswith('.png') :
                image_pixels = cv2.imread(image)
                #image_pixels = cv2.cvtColor(image_pixels, cv2.COLOR_BGR2RGB)
                image_pixels = cv2.cvtColor(image_pixels, cv2.COLOR_BGR2GRAY)
                image_pixels = cv2.resize(image_pixels, (dim,dim))
                image_pixels = image_pixels.reshape(dim, dim, 1)
                image_pixels = image_pixels/255.


                label = image.split("/")[-2]

                if len(image_pixels.shape) == 3 :
                    images.append(image_pixels)
                    labels.append(label)
                    filenames.append(image)
            else :
                image_pixels = plt.imread(image)
                image_pixels = cv2.cvtColor(image_pixels, cv2.COLOR_RGB2GRAY)
                image_pixels = cv2.resize(image_pixels, (dim,dim))
                image_pixels = image_pixels.reshape(dim, dim, 1)
                image_pixels = image_pixels/255.


                label = image.split("/")[-2]

                if len(image_pixels.shape) == 3 :
                    images.append(image_pixels)
                    labels.append(label)
                    filenames.append(image)
                    
        
            
        except Exception as ex:
            print("Exception :", ex)
            pass
        


    images = np.array(images)
    labels = np.array(labels)

    print(images.shape, labels.shape)

    return images, labels







#call the pre-moddelling method
X, y = prepare_images(train_images)

print(len(filenames))






#train and test split
(x_train,x_test, y_train, y_test)  = train_test_split(X, y,test_size=0.15, random_state=42)
print("train count :",len(x_train))
print("test count :",len(x_test))







#create a embedding visualize param
y_vis = y_train

# one hot encodeer
from sklearn import preprocessing
le = preprocessing.LabelEncoder()

y_train = le.fit_transform(y_train)
y_test = le.transform(y_test)







print(np.unique(y_vis))






from tensorflow.keras.layers import Input, AveragePooling2D


    
    
    
    
    
    
    
    

from tensorflow.keras.models import Model
from tensorflow.keras.layers import Dense, Dropout, Activation
from tensorflow.keras.layers import Conv2D
from tensorflow.keras.layers import AveragePooling2D
from tensorflow.keras.layers import GlobalAveragePooling2D
from tensorflow.keras.layers import Input, Concatenate
from tensorflow.keras.layers import BatchNormalization
from tensorflow.keras.regularizers import l2
import tensorflow.keras.backend as K


def conv_factory(x, concat_axis, nb_filter, dropout_rate=None, weight_decay=1E-4):
    """Apply BatchNorm, Relu 3x3Conv2D, optional dropout
    :parameter x: Input keras network
    :parameter concat_axis: int -- index of contatenate axis
    :parameter nb_filter: int -- number of filters
    :parameter dropout_rate: int -- dropout rate
    :parameter weight_decay: int -- weight decay factor
    :returns: keras network with b_norm, relu and Conv2D added
    :return type: keras network
    """
    x = BatchNormalization(axis=concat_axis, gamma_regularizer=l2(weight_decay), beta_regularizer=l2(weight_decay))(x)
    x = Activation('relu')(x)
    x = Conv2D(nb_filter, (1, 1), kernel_initializer="he_uniform", padding="same", use_bias=False, kernel_regularizer=l2(weight_decay))(x)
    x = BatchNormalization(axis=concat_axis, gamma_regularizer=l2(weight_decay), beta_regularizer=l2(weight_decay))(x)
    x = Activation('relu')(x)
    x = Conv2D(nb_filter, (3, 3), kernel_initializer="he_uniform", padding="same", use_bias=False, kernel_regularizer=l2(weight_decay))(x)
    if dropout_rate:
        x = Dropout(dropout_rate)(x)

    return x


def transition(x, concat_axis, nb_filter, dropout_rate=None, weight_decay=1E-4):
    """Apply BatchNorm, Relu 1x1Conv2D, optional dropout and Maxpooling2D
    :parameter x: keras model
    :parameter concat_axis: int -- index of contatenate axis
    :parameter nb_filter: int -- number of filters
    :parameter dropout_rate: int -- dropout rate
    :parameter weight_decay: int -- weight decay factor
    :returns: model
    :return type: keras model, after applying batch_norm, relu-conv, dropout, maxpool
    """

    x = BatchNormalization(axis=concat_axis, gamma_regularizer=l2(weight_decay), beta_regularizer=l2(weight_decay))(x)
    x = Activation('relu')(x)
    x = Conv2D(nb_filter, (1, 1), kernel_initializer="he_uniform", padding="same", use_bias=False, kernel_regularizer=l2(weight_decay))(x)
    if dropout_rate:
        x = Dropout(dropout_rate)(x)
    x = AveragePooling2D((2, 2), strides=(2, 2))(x)

    return x


def denseblock(x, concat_axis, nb_layers, nb_filter, growth_rate, dropout_rate=None, weight_decay=1E-4):
    """
    Build a denseblock where the output of each conv_factory is fed to subsequent ones
    :parameter x: keras model
    :parameter concat_axis: int -- index of contatenate axis
    :parameter nb_layers: int -- the number of layers of conv_factory to append to the model.
    :parameter nb_filter: int -- number of filters
    :parameter dropout_rate: int -- dropout rate
    :parameter weight_decay: int -- weight decay factor
    :returns: keras model with nb_layers of conv_factory appended
    :return type: keras model
    """

    list_feat = [x]

    for i in range(nb_layers):
        x = conv_factory(x, concat_axis, growth_rate,
                         dropout_rate, weight_decay)
        list_feat.append(x)
        x = Concatenate(axis=concat_axis)(list_feat)
        nb_filter += growth_rate
        #print (nb_filter)

    return x, nb_filter


def denseblock_altern(x, concat_axis, nb_layers, nb_filter, growth_rate, dropout_rate=None, weight_decay=1E-4):
    """Build a denseblock where the output of each conv_factory is fed to subsequent ones. (Alternative of denseblock)
    :parameter x: keras model
    :parameter concat_axis: int -- index of contatenate axis
    :parameter nb_layers: int -- the number of layers of conv_factory to append to the model.
    :parameter nb_filter: int -- number of filters
    :parameter dropout_rate: int -- dropout rate
    :parameter weight_decay: int -- weight decay factor
    :returns: keras model with nb_layers of conv_factory appended
    :return type: keras model
    * The main difference between this implementation and the implementation
    above is that the one above
    """

    for i in range(nb_layers):
        merge_tensor = conv_factory(x, concat_axis, growth_rate,  dropout_rate, weight_decay)
        x = Concatenate(axis=concat_axis)([merge_tensor, x])
        nb_filter += growth_rate

    return x, nb_filter


def DenseNet(nb_classes, img_dim, depth, nb_dense_block, growth_rate, nb_filter, dropout_rate=None, weight_decay=1E-4):
    """ 
    Build the DenseNet model
    :parameter nb_classes: int -- number of classes
    :parameter img_dim: tuple -- (channels, rows, columns)
    :parameter depth: int -- how many layers
    :parameter nb_dense_block: int -- number of dense blocks to add to end
    :parameter growth_rate: int -- number of filters to add
    :parameter nb_filter: int -- number of filters
    :parameter dropout_rate: float -- dropout rate
    :parameter weight_decay: float -- weight decay
    :returns: keras model with nb_layers of conv_factory appended
    :return type: keras model
    """
    
    if K.image_data_format() == "channels_first":
        concat_axis = 1
    elif K.image_data_format() == "channels_last":
        concat_axis = -1

    model_input = Input(shape=img_dim)

    assert (depth - 4) % 3 == 0, "Depth must be 3 N + 4"

    # layers in each dense block
    nb_layers = int((depth - 4) / 3)

    # Initial convolution
    x = Conv2D(nb_filter, (7, 7), strides=(2, 2), kernel_initializer="he_uniform", padding="same", name="initial_conv2D", use_bias=False, kernel_regularizer=l2(weight_decay))(model_input)

    # Add dense blocks
    nb_layers1 = [6,12,32,32,48,32,48,64,32]  #3*3 convolutional layer of each denseblock ，
    for block_idx in range(nb_dense_block - 1):
        x, nb_filter = denseblock(x, concat_axis, nb_layers1[block_idx], nb_filter, growth_rate, dropout_rate=dropout_rate, weight_decay=weight_decay)
        # add transition
        x = transition(x, concat_axis,nb_filter, dropout_rate=dropout_rate, weight_decay=weight_decay)

    # The last denseblock does not have a transition
    x, nb_filter = denseblock(x, concat_axis, nb_layers1[nb_dense_block-1], nb_filter, growth_rate, dropout_rate=dropout_rate, weight_decay=weight_decay)

    x = BatchNormalization(axis=concat_axis, gamma_regularizer=l2(weight_decay), beta_regularizer=l2(weight_decay))(x)
    x = Activation('relu')(x)
    x = GlobalAveragePooling2D(data_format=K.image_data_format())(x)
    x = Dense(nb_classes, activation='softmax', kernel_regularizer=l2(weight_decay), bias_regularizer=l2(weight_decay))(x)

    densenet = Model(inputs=[model_input], outputs=[x], name="DenseNet")

    return densenet







nb_classes = 2
img_dim = (dim,dim,1)
depth = 6*3+4
nb_dense_block = 4
growth_rate = 16
nb_filter = 12
dropout_rate = 0.2
weight_decay = 1E-4
learning_rate = 1E-3

model = DenseNet(nb_classes, img_dim, depth, nb_dense_block, growth_rate, nb_filter, dropout_rate=dropout_rate, weight_decay=weight_decay)

# Model output
print("\n\n", model.summary())

# Build optimizer
opt = keras.optimizers.Adam(lr=learning_rate, beta_1=0.9, beta_2=0.999, epsilon=1e-08)
model.compile(loss="binary_crossentropy", optimizer=opt,
    metrics=["accuracy"])






nb_classes = 2
img_dim = (dim,dim,1)
depth = 6*3+4
nb_dense_block = 4
growth_rate = 16
nb_filter = 12
dropout_rate = 0.2
weight_decay = 1E-4
learning_rate = 1E-3

model = DenseNet(nb_classes, img_dim, depth, nb_dense_block, growth_rate, nb_filter, dropout_rate=dropout_rate, weight_decay=weight_decay)
# Model output
#model.summary()

# Build optimizer
opt = keras.optimizers.Adam(lr=learning_rate, beta_1=0.9, beta_2=0.999, epsilon=1e-08)
model.compile(loss="binary_crossentropy", optimizer=opt,
    metrics=["accuracy"])






num_epochs = 150
batch_size = 8








## Define callback
from tensorflow.keras.callbacks import CSVLogger
import os

os.makedirs('saved_models/try2/', exist_ok=True)

csv_logger = CSVLogger('saved_models/try2/10_Nov_2_class_log.csv', append=True, separator=',')
checkpoint = ModelCheckpoint('saved_models/try2/model.{epoch:02d}.h5', monitor='val_loss', verbose=2, save_best_only=False, mode='min', period=1)

history = model.fit(x=x_train, y=y_train, batch_size=batch_size, epochs=num_epochs, callbacks = [checkpoint, csv_logger])

accuracy = model.evaluate(x_test, y_test)[1]
print(f"Test accuracy: {round(accuracy * 100, 2)}%")





print("Done..!!!")
