"""
<center style = "font-family: 'Lucida Console', 'Courier New', monospace;">
    <img src = "https://www.bargainmarket.com.au/images/banner-pets-dog-cat-boarding.png" width=600 height = 400>
    <h1 style = "background: rgb(238,174,202);
background: radial-gradient(circle, rgba(238,174,202,1) 0%, rgba(81,154,204,1) 0%);">PetFinder.my - Pawpularity Contest 🐶🐱</h1>
</center>

<br>
<h2 style = "font-family: Consolas">What is PetFinder.my❓</h2>
<p style = "font-family : Lucida Sans Typewriter">
<a href = "https://petfinder.my/">PetFinder.my</a> is Malaysia’s leading animal welfare platform, featuring over 180,000 animals with 54,000 happily adopted. PetFinder collaborates closely with animal lovers, media, corporations, and global organizations to improve animal welfare.
</p>

<br>
<h2 style = "font-family: Consolas">What should we do📝</h2>
<p style = "font-family : Lucida Sans Typewriter">
In this competition, you’ll analyze raw images and metadata to predict the “Pawpularity” of pet photos. You'll train and test your model on PetFinder.my's thousands of pet profiles.</p>

<br>
<h2 style = "font-family: Consolas">Data Information💽</h2>
<p style = "font-family : Lucida Sans Typewriter">Check <a href = "https://www.kaggle.com/c/petfinder-pawpularity-score/data">competition page</a> for details</p>

<h2 style = "font-family : Comic Sans MS">Let's dive in ⬇️</h2>
"""
"""
<center style = "background: rgb(238,174,202);
background: radial-gradient(circle, rgba(238,174,202,1) 0%, rgba(81,154,204,1) 0%);"><h1 style = "font-family:'Bodoni MT'">Libraries📚</h1></center>
"""
!/opt/conda/bin/python3.7 -m pip install --upgrade pip
! pip install -q efficientnet
##--------------------------
# Import necessary libraries
##--------------------------
import os

import pandas as pd
from pandas_profiling import ProfileReport

import matplotlib.pyplot as plt
%matplotlib inline
import seaborn as sns
sns.set()

import numpy as np

from PIL import Image

import tensorflow as tf
from tensorflow.keras.layers import *
import efficientnet.tfkeras as efn
from sklearn.model_selection import KFold,StratifiedKFold

from kaggle_datasets import KaggleDatasets

import warnings
warnings.filterwarnings('ignore')
"""
<center style = "background: rgb(238,174,202);
background: radial-gradient(circle, rgba(238,174,202,1) 0%, rgba(81,154,204,1) 0%);"><h1 style = "font-family:'Bodoni MT'">Checking TPU access💻</h1></center>
"""
AUTO = tf.data.experimental.AUTOTUNE

try:
    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()  
    print('Running on TPU ', tpu.master())
except ValueError:
    tpu = None

if tpu:
    tf.config.experimental_connect_to_cluster(tpu)
    tf.tpu.experimental.initialize_tpu_system(tpu)
    strategy = tf.distribute.experimental.TPUStrategy(tpu)
else:
    strategy = tf.distribute.get_strategy()

print("REPLICAS: ", strategy.num_replicas_in_sync)
"""
<center style = "background: rgb(238,174,202);
background: radial-gradient(circle, rgba(238,174,202,1) 0%, rgba(81,154,204,1) 0%);"><h1 style = "font-family:'Bodoni MT'">Initializations🎬</h1></center>
"""
GCS_DS_PATH = KaggleDatasets().get_gcs_path("petfinder-pawpularity-score")
print(GCS_DS_PATH)

TRAIN_PATH = GCS_DS_PATH + "/train/"
TEST_PATH = GCS_DS_PATH + "/TEST/"

TRAIN_P = "../input/petfinder-pawpularity-score/train"
TEST_P = "../input/petfinder-pawpularity-score/test"

train_df = pd.read_csv("../input/petfinder-pawpularity-score/train.csv")
test_df = pd.read_csv("../input/petfinder-pawpularity-score/test.csv")

TRAIN_NS = os.listdir(TRAIN_P) 
TEST_NS = os.listdir(TEST_P)
TRAIN_PS = [os.path.join(TRAIN_P,f) for f in TRAIN_NS]
TEST_PS = [os.path.join(TEST_P,f) for f in TEST_NS]

HEIGHT, WIDTH = 512, 512
CHANNELS = 3
SEED = 2021
NUM_CLASSES = 10
"""
<center style = "background: rgb(238,174,202);
background: radial-gradient(circle, rgba(238,174,202,1) 0%, rgba(81,154,204,1) 0%);"><h1 style = "font-family:'Bodoni MT'">Data Visualization📊</h1></center>
"""
##------------------------------
# Display Random Training Images
##------------------------------

rand_ls = np.random.randint(9000, size=(9))
fig,ax = plt.subplots(nrows=3,ncols=3,figsize=(18,10))
for count,i in enumerate(rand_ls):
    r,c = count//3, count%3
    ax[r,c].imshow(np.asarray(Image.open(TRAIN_PS[i])))
    ax[r,c].axis("off")
    label = train_df[train_df["Id"] == TRAIN_NS[i].split(".")[0]]["Pawpularity"].values[0]
    ax[r,c].set_title(f"Pawpularity Score : {label}",fontsize=15, fontfamily="monospace", fontweight="bold")
##--------------------------------
# Height, Width of Training Images
##--------------------------------
train_h, train_w = [], []

for path in TRAIN_PS:
    im = Image.open(path)
    w,h = im.size
    train_h.append(h)
    train_w.append(w)

'''plt.figure(figsize=(10,10))
sns.jointplot(x=train_w, y=train_h,kind="kde")
plt.xlabel("Width of Image")
plt.ylabel("Height of Image")'''
print(f"Mean Height : {sum(train_h)/len(train_h)} | Mean Width : {sum(train_w)/len(train_w)}")
##---------------------------
# Metadata of Training images
# Pandas Profiling
##---------------------------
train_df["Width"] = train_w
train_df["Height"] = train_h

train_report = ProfileReport(train_df,title="Metadata of Training images")
#train_report.to_file("./train_metadata.html")
train_report
"""
<center style = "background: rgb(238,174,202);
background: radial-gradient(circle, rgba(238,174,202,1) 0%, rgba(81,154,204,1) 0%);"><h1 style = "font-family:'Bodoni MT'">Data Preprocessing🧰</h1></center>
"""
##----------------------
#Preprocessing Functions
##----------------------

def process_img(filepath,label):
    image = tf.io.read_file(filepath)
    image = tf.image.decode_jpeg(image, channels=CHANNELS)
    image = tf.image.convert_image_dtype(image, tf.float32) 
    image = tf.image.resize(image, [HEIGHT,WIDTH])
    return image,label


def data_augment(image, label):
    p_spatial = tf.random.uniform([], 0, 1.0, dtype=tf.float32)
    p_rotate = tf.random.uniform([], 0, 1.0, dtype=tf.float32)
    p_pixel_1 = tf.random.uniform([], 0, 1.0, dtype=tf.float32)
    p_pixel_2 = tf.random.uniform([], 0, 1.0, dtype=tf.float32)
    p_pixel_3 = tf.random.uniform([], 0, 1.0, dtype=tf.float32)
    p_crop = tf.random.uniform([], 0, 1.0, dtype=tf.float32)
            
    # Flips
    image = tf.image.random_flip_left_right(image)
    image = tf.image.random_flip_up_down(image)
    if p_spatial > .75:
        image = tf.image.transpose(image)
        
    # Rotates
    if p_rotate > .75:
        image = tf.image.rot90(image, k=3) 
    elif p_rotate > .5:
        image = tf.image.rot90(image, k=2) 
    elif p_rotate > .25:
        image = tf.image.rot90(image, k=1) 
        
    
    if p_pixel_1 >= .4:
        image = tf.image.random_saturation(image, lower=.7, upper=1.3)
    if p_pixel_2 >= .4:
        image = tf.image.random_contrast(image, lower=.8, upper=1.2)
    if p_pixel_3 >= .4:
        image = tf.image.random_brightness(image, max_delta=.1)
        
    
    if p_crop > .7:
        if p_crop > .9:
            image = tf.image.central_crop(image, central_fraction=.7)
        elif p_crop > .8:
            image = tf.image.central_crop(image, central_fraction=.8)
        else:
            image = tf.image.central_crop(image, central_fraction=.9)
    elif p_crop > .4:
        crop_size = tf.random.uniform([], int(HEIGHT*.8), HEIGHT, dtype=tf.int32)
        image = tf.image.random_crop(image, size=[crop_size, crop_size, CHANNELS])
    
    image = tf.image.resize(image, [HEIGHT,WIDTH])
    return image,label

def get_dataset(filenames,labels, training=True):
    dataset = tf.data.Dataset.from_tensor_slices((filenames,labels))
    dataset = dataset.map(process_img,num_parallel_calls=AUTO)
    dataset = dataset.map(data_augment,num_parallel_calls=AUTO)
    dataset = dataset.cache()
    dataset = dataset.repeat()
    if training:
        dataset = dataset.shuffle(2048)
    dataset = dataset.batch(BATCH_SIZE)
    dataset = dataset.prefetch(AUTO)
    return dataset
files_ls = tf.io.gfile.glob(TRAIN_PATH + "*.jpg" )
files_df = pd.DataFrame(files_ls, columns = ["filepath"])

bin_ranges = [0, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100]
bin_names = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

train_df["label"] = pd.cut(np.array(train_df["Pawpularity"]),bins=bin_ranges,labels=bin_names)

real_labels = np.array(train_df["Pawpularity"])
paw_values = []
for i in bin_names:
    val = list(train_df[train_df["label"] == i]["Pawpularity"])
    paw_values.append(sum(val)/len(val))

pseudo_labels = np.array(train_df["label"])
                       
'''onehot_labels = np.zeros((len(labels), NUM_CLASSES))
for i in range(len(pseudo_labels)):
    onehot_labels[i][pseudo_labels[i]] = 1'''
    
train_df.head()
"""
<center style = "background: rgb(238,174,202);
background: radial-gradient(circle, rgba(238,174,202,1) 0%, rgba(81,154,204,1) 0%);"><h1 style = "font-family:'Bodoni MT'">Setting Up Training Process🍳</h1></center>
"""
def create_model():
    
    pretrained = efn.EfficientNetB7(include_top=False, weights='noisy-student',input_shape=[HEIGHT,WIDTH, CHANNELS])
    '''pretrained = tf.keras.applications.DenseNet121(weights= "imagenet",
                                                   include_top=False,
                                                   input_shape=(HEIGHT,WIDTH,CHANNELS),pooling=None)'''        
    x = pretrained.output
    x = tf.keras.layers.GlobalAveragePooling2D() (x)
    x = tf.keras.layers.Dense(512, activation = "relu") (x)
    x = tf.keras.layers.Dropout(0.3) (x)
    x = tf.keras.layers.Dense(256, activation="relu") (x)
    x = tf.keras.layers.Dropout(0.3) (x)
    #outputs = tf.keras.layers.Dense(NUM_CLASSES,activation="softmax", dtype='float32')(x)
    outputs = tf.keras.layers.Dense(1, dtype='float32')(x)
        
    model = tf.keras.Model(pretrained.input, outputs)
    return model
import tensorflow_addons as tfa

def compile_model(model, lr=0.001):
    
    optimizer = tf.keras.optimizers.Adam(lr=lr)
    
    #loss = tf.keras.losses.CategoricalCrossentropy()
    loss = tf.keras.losses.MeanSquaredError()
   
    metrics = [tf.keras.metrics.RootMeanSquaredError(name='rmse')]
    '''
    tfa.metrics.F1Score(num_classes = NUM_CLASSES,average = "macro", name = "f1_score"),
    tf.keras.metrics.CategoricalAccuracy(name='acc')
    '''
    model.compile(optimizer=optimizer, loss=loss, metrics=metrics)

    return model
METRIC = "val_rmse"

def create_callbacks(kfold,metric = METRIC):
    
    cpk_path = f'./best_model_{kfold}.h5'
    
    checkpoint = tf.keras.callbacks.ModelCheckpoint(
        filepath=cpk_path,
        monitor= metric,
        mode='min',
        save_best_only=True,
        verbose=1,
    )

    reducelr = tf.keras.callbacks.ReduceLROnPlateau(
        monitor= metric,
        mode='min',
        factor=0.1,
        patience=3,
        verbose=0
    )

    earlystop = tf.keras.callbacks.EarlyStopping(
        monitor= metric,
        mode='min',
        patience=10, 
        verbose=1
    )
    
    callbacks = [checkpoint, reducelr, earlystop]         
    
    return callbacks
"""
<center style = "background: rgb(238,174,202);
background: radial-gradient(circle, rgba(238,174,202,1) 0%, rgba(81,154,204,1) 0%);"><h1 style = "font-family:'Bodoni MT'">Training💡</h1></center>
"""
EPOCHS = 10
VERBOSE = 1
N_SPLITS = 10
BATCH_SIZE = 32

kfold = StratifiedKFold(n_splits=N_SPLITS, shuffle=True, random_state=SEED)
history = {}


for fold,(tID,vID) in enumerate(kfold.split(files_ls,pseudo_labels)):
    tFiles, tLabels = list(files_df.iloc[tID]["filepath"]),real_labels[tID]  #, onehot_labels[tID]
    vFiles, vLabels = list(files_df.iloc[vID]["filepath"]), real_labels[vID]  #, onehot_labels[vID]
    print("Number of Training Images: ",len(tID))
    print("Number of Validation Images: ",len(vID))
    
    STEPS_PER_EPOCH  = len(tID)//BATCH_SIZE
    VALID_STEPS = len(vID)//BATCH_SIZE
    
    tf.keras.backend.clear_session()
    
    train_ds = get_dataset(tFiles,tLabels, training = True)
    val_ds = get_dataset(vFiles, vLabels, training = False)
    
    with strategy.scope():
        model = create_model()
        model = compile_model(model, lr=0.0001)
        callbacks = create_callbacks(kfold = fold)
    
        print("------------------Fold - ",fold+1," --------------------------")
        history[fold] = model.fit(
                            train_ds,
                            epochs=EPOCHS,
                            callbacks=callbacks,
                            validation_data = val_ds,
                            verbose=VERBOSE,
                            steps_per_epoch = STEPS_PER_EPOCH,
                            validation_steps=VALID_STEPS
                           )
"""
<center style = "background: rgb(238,174,202);
background: radial-gradient(circle, rgba(238,174,202,1) 0%, rgba(81,154,204,1) 0%);"><h1 style = "font-family:'Bodoni MT'">How did the Training go📜</h1></center>
"""
plt.figure(figsize=(8*N_SPLITS,24))

for i in range(N_SPLITS):
    acc = history[i].history['rmse']
    val_acc = history[i].history['val_rmse']
    loss = history[i].history['loss']
    val_loss = history[i].history['val_loss']
    epochs_range = range(len(history[i].history['val_loss'])) 
    
    plt.subplot(N_SPLITS, 2,i*2+1)
    plt.plot(epochs_range, acc, label='Training RMSE')
    plt.plot(epochs_range, val_acc, label='Validation  RMSE')
    plt.legend(loc='lower right')
    plt.title(f'FOLD:{str(i)} Training and Validation  RMSE')
    
    plt.subplot(N_SPLITS, 2, i*2+2)
    plt.plot(epochs_range, loss, label='Training Loss')
    plt.plot(epochs_range, val_loss, label='Validation Loss')
    plt.legend(loc='upper right')
    plt.title(f'FOLD:{str(i)} Training and Validation Loss')

plt.show()