"""
<center><img src="https://raw.githubusercontent.com/dimitreOliveira/MachineLearning/master/Kaggle/Cassava%20Leaf%20Disease%20Classification/banner.png" width="1000"></center>
<br>
<center><h1>Cassava Leaf Disease - Supervised Contrastive Learning - Inference</h1></center>
<br>

- This is the inference part of the work, the training notebook can be found here [Cassava Leaf - Supervised Contrastive Learning](https://www.kaggle.com/dimitreoliveira/cassava-leaf-supervised-contrastive-learning)
- keras-applications GitHub repository can be found [here](https://www.kaggle.com/dimitreoliveira/kerasapplications)
- efficientnet GitHub repository can be found [here](https://www.kaggle.com/dimitreoliveira/efficientnet-git)
- Dataset source `center cropped` [512x512](https://www.kaggle.com/dimitreoliveira/cassava-leaf-disease-50-tfrecords-center-512x512)
- Dataset source `external data` `center cropped` [512x512](https://www.kaggle.com/dimitreoliveira/cassava-leaf-disease-50-tfrecords-external-512x512)
- Dataset source [discussion thread](https://www.kaggle.com/c/cassava-leaf-disease-classification/discussion/198744)
- Dataset [creation source](https://www.kaggle.com/dimitreoliveira/cassava-leaf-disease-stratified-tfrecords-256x256)
"""
"""
## Dependencies
"""
!pip install --quiet /kaggle/input/kerasapplications
!pip install --quiet /kaggle/input/efficientnet-git
import math, os, re, warnings, random, glob
import numpy as np
import pandas as pd
import tensorflow as tf
import tensorflow.keras.layers as L
import tensorflow.keras.backend as K
from tensorflow.keras import Sequential, Model
import efficientnet.tfkeras as efn

def seed_everything(seed=0):
    random.seed(seed)
    np.random.seed(seed)
    tf.random.set_seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    os.environ['TF_DETERMINISTIC_OPS'] = '1'

seed = 0
seed_everything(seed)
warnings.filterwarnings('ignore')
"""
### Hardware configuration
"""
# TPU or GPU detection
# Detect hardware, return appropriate distribution strategy
try:
    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()
    print(f'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()

AUTO = tf.data.experimental.AUTOTUNE
REPLICAS = strategy.num_replicas_in_sync
print(f'REPLICAS: {REPLICAS}')
"""
# Model parameters
"""
BATCH_SIZE = 16 * REPLICAS
HEIGHT = 512
WIDTH = 512 
CHANNELS = 3
N_CLASSES = 5
TTA_STEPS = 0 # Do TTA if > 0 
USE_REGULAR = False
USE_SCL = True
"""
# Augmentation
"""
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) # rotate 270º
    elif p_rotate > .5:
        image = tf.image.rot90(image, k=2) # rotate 180º
    elif p_rotate > .25:
        image = tf.image.rot90(image, k=1) # rotate 90º
        
    # Pixel-level transforms
    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)
        
    # Crops
    if p_crop > .6:
        if p_crop > .9:
            image = tf.image.central_crop(image, central_fraction=.5)
        elif p_crop > .8:
            image = tf.image.central_crop(image, central_fraction=.6)
        elif p_crop > .7:
            image = tf.image.central_crop(image, central_fraction=.7)
        else:
            image = tf.image.central_crop(image, central_fraction=.8)
    elif p_crop > .3:
        crop_size = tf.random.uniform([], int(HEIGHT*.6), HEIGHT, dtype=tf.int32)
        image = tf.image.random_crop(image, size=[crop_size, crop_size, CHANNELS])
            
    image = tf.image.resize(image, size=[HEIGHT, WIDTH])
        
    return image, label
"""
## Auxiliary functions
"""
# Datasets utility functions
def get_name(file_path):
    parts = tf.strings.split(file_path, os.path.sep)
    name = parts[-1]
    return name

def decode_image(image_data):
    image = tf.image.decode_jpeg(image_data, channels=3)
    image = tf.cast(image, tf.float32) / 255.0
    
#     image = center_crop(image)
    return image

def center_crop(image):
    image = tf.reshape(image, [600, 800, CHANNELS]) # Original shape
    
    h, w = image.shape[0], image.shape[1]
    if h > w:
        image = tf.image.crop_to_bounding_box(image, (h - w) // 2, 0, w, w)
    else:
        image = tf.image.crop_to_bounding_box(image, 0, (w - h) // 2, h, h)
        
    image = tf.image.resize(image, [HEIGHT, WIDTH]) # Expected shape
    return image

def resize_image(image, label):
    image = tf.image.resize(image, [HEIGHT, WIDTH])
    image = tf.reshape(image, [HEIGHT, WIDTH, CHANNELS])
    return image, label

def process_path(file_path):
    name = get_name(file_path)
    img = tf.io.read_file(file_path)
    img = decode_image(img)
    return img, name

def get_dataset(files_path, shuffled=False, tta=False, extension='jpg'):
    dataset = tf.data.Dataset.list_files(f'{files_path}*{extension}', shuffle=shuffled)
    dataset = dataset.map(process_path, num_parallel_calls=AUTO)
    if tta:
        dataset = dataset.map(data_augment, num_parallel_calls=AUTO)
    dataset = dataset.map(resize_image, num_parallel_calls=AUTO)
    dataset = dataset.batch(BATCH_SIZE)
    dataset = dataset.prefetch(AUTO)
    return dataset

def count_data_items(filenames):
    n = [int(re.compile(r"-([0-9]*)\.").search(filename).group(1)) for filename in filenames]
    return np.sum(n)
"""
# Load data
"""
database_base_path = '/kaggle/input/cassava-leaf-disease-classification/'
submission = pd.read_csv(f'{database_base_path}sample_submission.csv')
display(submission.head())

TEST_FILENAMES = tf.io.gfile.glob(f'{database_base_path}test_tfrecords/ld_test*.tfrec')
NUM_TEST_IMAGES = count_data_items(TEST_FILENAMES)
print(f'GCS: test: {NUM_TEST_IMAGES}')
model_path_list = glob.glob('/kaggle/input/cassava-leaf-supervised-contrastive-learning/model_reg*.h5')
model_path_list.sort()

print('Models to predict:')
print(*model_path_list, sep='\n')

model_path_list_scl = glob.glob('/kaggle/input/cassava-leaf-supervised-contrastive-learning/model_scl*.h5')
model_path_list_scl.sort()

print('\nModels to predict:')
print(*model_path_list_scl, sep='\n')
"""
# Model
"""
def encoder_fn(input_shape):
    inputs = L.Input(shape=input_shape, name='inputs')
    base_model = efn.EfficientNetB3(input_tensor=inputs, 
                                    include_top=False, 
                                    weights=None, 
                                    pooling='avg')
    
    model = Model(inputs=inputs, outputs=base_model.outputs)

    return model


def classifier_fn(input_shape, N_CLASSES, encoder, trainable=True):
    for layer in encoder.layers:
        layer.trainable = trainable
        
    inputs = L.Input(shape=input_shape, name='inputs')
    
    features = encoder(inputs)
    features = L.Dropout(.5)(features)
    features = L.Dense(1000, activation='relu')(features)
    features = L.Dropout(.5)(features)
    outputs = L.Dense(N_CLASSES, activation='softmax', name='outputs', dtype='float32')(features)

    model = Model(inputs=inputs, outputs=outputs)

    return model
"""
# Test set predictions
"""
files_path = f'{database_base_path}test_images/'
test_size = len(os.listdir(files_path))
test_preds = np.zeros((test_size, N_CLASSES))


if USE_REGULAR:
    print('Inference for regular trainining models')
    with strategy.scope():
        encoder = encoder_fn((None, None, CHANNELS))
        model = classifier_fn((None, None, CHANNELS), N_CLASSES, encoder)

    for model_path in model_path_list:
        print(model_path)
        K.clear_session()
        model.load_weights(model_path)

        if TTA_STEPS > 0:
            test_ds = get_dataset(files_path, tta=True).repeat()
            ct_steps = TTA_STEPS * ((test_size/BATCH_SIZE) + 1)
            preds = model.predict(test_ds, steps=ct_steps, verbose=1)[:(test_size * TTA_STEPS)]
            preds = np.mean(preds.reshape(test_size, TTA_STEPS, N_CLASSES, order='F'), axis=1)
            test_preds += preds / len(model_path_list)
        else:
            test_ds = get_dataset(files_path, tta=False)
            x_test = test_ds.map(lambda image, image_name: image)
            test_preds += model.predict(x_test) / len(model_path_list)
        
        
if USE_SCL:
    print('\nInference for "Supervised Contrastive Learning" models')
    with strategy.scope():
        encoder = encoder_fn((None, None, CHANNELS))
        model = classifier_fn((None, None, CHANNELS), N_CLASSES, encoder, trainable=False)

    for model_path in model_path_list_scl:
        print(model_path)
        K.clear_session()
        model.load_weights(model_path)

        if TTA_STEPS > 0:
            test_ds = get_dataset(files_path, tta=True).repeat()
            ct_steps = TTA_STEPS * ((test_size/BATCH_SIZE) + 1)
            preds = model.predict(test_ds, steps=ct_steps, verbose=1)[:(test_size * TTA_STEPS)]
            preds = np.mean(preds.reshape(test_size, TTA_STEPS, N_CLASSES, order='F'), axis=1)
            test_preds += preds / len(model_path_list_scl)
        else:
            test_ds = get_dataset(files_path, tta=False)
            x_test = test_ds.map(lambda image, image_name: image)
            test_preds += model.predict(x_test) / len(model_path_list_scl)
    
test_preds = np.argmax(test_preds, axis=-1)
test_names_ds = get_dataset(files_path)
image_names = [img_name.numpy().decode('utf-8') for img, img_name in iter(test_names_ds.unbatch())]
submission = pd.DataFrame({'image_id': image_names, 'label': test_preds})
submission.to_csv('submission.csv', index=False)
display(submission.head())