from google.colab import drive
drive.mount('/content/drive')

!pip install seaborn openvino-dev
!pip install tensorflow_addons wandb git+https://github.com/leondgarse/keras_cv_attention_models.git
!wandb login

# Import Data Science Libraries
import numpy as np
import pandas as pd
import tensorflow as tf
from sklearn.model_selection import train_test_split
import itertools
import random
import math
# Import visualization libraries
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import cv2
import seaborn as sns

# Tensorflow Libraries
from tensorflow import keras
from tensorflow.keras import layers,models
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.layers import Dense, Dropout
from tensorflow.keras.callbacks import Callback, EarlyStopping,ModelCheckpoint
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras import Model
from tensorflow.keras.layers.experimental import preprocessing
from tensorflow.data.experimental import AUTOTUNE
from tensorflow.keras.callbacks import ReduceLROnPlateau
import tensorflow_addons as tfa
import tensorflow.keras.backend as K

# System libraries
from pathlib import Path
import os.path

# Metrics
from sklearn.metrics import classification_report, confusion_matrix
from wandb.keras import WandbCallback,WandbMetricsLogger

import wandb

SEED= 42
model_name = 'efficeintnetv2L-480' # @param {type:"string"}
USER_NAME = 'saidinesh' # @param {type:"string"}
wandb.init(entity=USER_NAME, project='mosquito-classification', job_type='train', name=model_name , sync_tensorboard=True)

sns.set(style='darkgrid')

# Detect TPU, return appropriate distribution strategy

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)

# Seed Everything to reproduce results for future use cases
def seed_everything(seed=42):
    # Seed value for TensorFlow
    tf.random.set_seed(seed)

    # Seed value for NumPy
    np.random.seed(seed)

    # Seed value for Python's random library
    random.seed(seed)

    # Force TensorFlow to use single thread
    # Multiple threads are a potential source of non-reproducible results.
    session_conf = tf.compat.v1.ConfigProto(
        intra_op_parallelism_threads=1,
        inter_op_parallelism_threads=1
    )

    # Make sure that TensorFlow uses a deterministic operation wherever possible
    tf.compat.v1.set_random_seed(seed)

    sess = tf.compat.v1.Session(graph=tf.compat.v1.get_default_graph(), config=session_conf)
    tf.compat.v1.keras.backend.set_session(sess)

seed_everything(SEED)

!wget https://raw.githubusercontent.com/mrdbourke/tensorflow-deep-learning/main/extras/helper_functions.py

# Import series of helper functions for our notebook
from helper_functions import create_tensorboard_callback, plot_loss_curves, unzip_data, compare_historys, walk_through_dir, pred_and_plot

BATCH_SIZE = 8 # @param {type:"number"}
INPUT_SIZE = "480" #@param ["384","480","512"]
TARGET_SIZE = (int(INPUT_SIZE),int(INPUT_SIZE)) #(224, 224)
BUFFER_SIZE = 60000
BATCH_SIZE = 8 * strategy.num_replicas_in_sync #64
HEIGHT=TARGET_SIZE[0]
WIDTH = TARGET_SIZE[1]
IMAGE_SIZE = TARGET_SIZE
INIT_LR = 0.00001 * (BATCH_SIZE/64)
EPOCHS = 40
DISABLE_CUTMIX = False # @param {type:"boolean"}
# ImageNet mean and std
MEAN = [0.485, 0.456, 0.406]
STD = [0.229, 0.224, 0.225]

gcs_path = 'gs://kds-9a7c3559b6f8b4ef686fe4a6c2224bed3900561891c6ec018ef0fcc2'
# walk_through_dir(gcs_path)
training_filepath = tf.io.gfile.glob(gcs_path + '/train_crop.tfrecords')
validation_filepath = tf.io.gfile.glob(gcs_path + '/val_crop.tfrecords')
test_filepath = validation_filepath

# Load TFRecord file from the folder as bytes
raw_training_dataset = tf.data.TFRecordDataset(training_filepath)
raw_validation_dataset = tf.data.TFRecordDataset(validation_filepath)
raw_test_dataset = tf.data.TFRecordDataset(test_filepath)

# Create a dictionary describing the features
labeled_feature_description = {
    'label': tf.io.FixedLenFeature([], tf.int64),
    'image': tf.io.FixedLenFeature([], tf.string)
}
#Class names of Mosquitoes
CLASSES = ['aegypti', 'albopictus', 'anopheles', 'culex', 'culiseta', 'japonicus-koreicus']

# Create a function to read and extract images from dataset
def _parse_labeled_image_function(example_proto):
    example = tf.io.parse_single_example(example_proto, labeled_feature_description)
    image = tf.io.decode_jpeg(example['image'])
    image = tf.cast(image, tf.float32) / 255.
    image = tf.image.resize(image, TARGET_SIZE)
    #label = tf.cast(example['label'], tf.int32)
    label = tf.one_hot(example['label'], depth=6)
    return image, label

# Parse labeled images, shuffle and batch
training_dataset = (
    raw_training_dataset
    .map(_parse_labeled_image_function)
    .shuffle(BUFFER_SIZE)
    .batch(BATCH_SIZE)
    .prefetch(AUTOTUNE)
)

# Parse unlabeled images and batch
validation_dataset = (
    raw_validation_dataset
    .map(_parse_labeled_image_function)
    .shuffle(BUFFER_SIZE)
    .batch(BATCH_SIZE)
    .prefetch(AUTOTUNE)
)

# Parse unlabeled images and batch
test_dataset = (
    raw_test_dataset
    .map(_parse_labeled_image_function)
    .batch(BATCH_SIZE)
    .prefetch(AUTOTUNE)
)

# Display images in a 5x5 grid

image_batch, label_batch = next(iter(training_dataset))

def display_images(image_batch, label_batch):
    plt.figure(figsize = [20,12])
    for i in range(25):
        plt.subplot(5,5,i+1)
        plt.imshow(image_batch[i])
        plt.title(CLASSES[np.argmax(label_batch[i].numpy())])
        plt.axis('off')
    plt.show()

display_images(image_batch, label_batch)