# This Python 3 environment comes with many helpful analytics libraries installed
# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python
# For example, here's several helpful packages to load

import numpy as np # linear algebra
import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)'
from keras.preprocessing.image import ImageDataGenerator, load_img, img_to_array
import tensorflow as tf
import albumentations as A
import os
from keras.models import Sequential
from keras.layers import Conv2D, MaxPooling2D
from keras.layers import Activation, Dropout, Flatten, Dense
import shutil
from tensorflow.compat.v1 import ConfigProto
from tensorflow.compat.v1 import InteractiveSession
from tensorflow.python.client import device_lib


### TODO -- sort images out into working folder, with subdirectory for each class.  Allows Generator to automatically determine labels.  

def main():
    config = ConfigProto()
    config.gpu_options.allow_growth = True
    session = InteractiveSession(config=config)
    
    print(device_lib.list_local_devices())
    # Input data files are available in the read-only "../input/" directory
    # For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory

    """ import os
    for dirname, _, filenames in os.walk('/kaggle/input'):
        for filename in filenames:
            print(os.path.join(dirname, filename)) """

    # You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using "Save & Run All" 
    # You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session
    
    sort_images()
    
    batch_size = 128
    
    ## Data Augmentation
    train_datagen = ImageDataGenerator(
            rescale=1./255,
            rotation_range=20,
            width_shift_range=0.2,
            height_shift_range=0.2,
            zoom_range=0.4,
            horizontal_flip=True,
            vertical_flip=True,
            channel_shift_range=0.05,
            brightness_range=[0.1, 0.2],
            preprocessing_function=compress,
            validation_split=0.2,
            fill_mode='reflect')
    
    ## When testing, we don't want to augment the images, only rescale to similar size
    test_datagen = ImageDataGenerator(rescale=1./255,)
    
    # this is a generator that will read pictures found in
    # subfolers of 'kaggle/working/train', and indefinitely generate
    # batches of augmented image data
    train_generator = train_datagen.flow_from_directory(
        'train',  # this is the target directory
        target_size=(150, 150),  # all images will be resized to 150x150
        batch_size=batch_size,
        subset="training",
        color_mode="rgb",
        class_mode='categorical')  # since we use categorical_crossentropy loss, we need categorical labels

    # this is a similar generator, for validation data
    validation_generator = train_datagen.flow_from_directory(
        'train',
        target_size=(150, 150),
        batch_size=batch_size,
        subset="validation",
        color_mode="rgb",
        class_mode='categorical')
 

    ## Build Model
    model = Sequential()
    model.add(Conv2D(32, (3, 3), input_shape=(150, 150, 3)))
    model.add(Activation('relu'))
    model.add(MaxPooling2D(pool_size=(2, 2), padding='same'))

    model.add(Conv2D(32, (3, 3)))
    model.add(Activation('relu'))
    model.add(MaxPooling2D(pool_size=(2, 2), padding='same'))
    
    model.add(Conv2D(64, (3, 3)))
    model.add(Activation('relu'))
    model.add(MaxPooling2D(pool_size=(2, 2), padding='same'))

    model.add(Conv2D(64, (3, 3)))
    model.add(Activation('relu'))
    model.add(MaxPooling2D(pool_size=(2, 2), padding='same'))   # the model so far outputs 3D feature maps (height, width, features)
    
    model.add(Flatten())  # this converts our 3D feature maps to 1D feature vectors
    model.add(Dense(64))
    model.add(Activation('relu'))
    model.add(Dropout(0.5))
    model.add(Dense(5))   # 5 output nodes, one for each classification
    model.add(Activation('softmax'))
    
    model.compile(loss='categorical_crossentropy',
              optimizer='adam',
              metrics=['accuracy'])
    
    model.fit_generator(
        train_generator,
        steps_per_epoch=16000 // batch_size,
        epochs=10,
        validation_data=validation_generator,
        validation_steps=800 // batch_size)
    model.save_weights('first_try.h5')  # always save your weights after training or during training

def compress(image):
    transform = get_train_transforms()
    return transform(image=image)["image"]
    
    
def get_train_transforms(aug=None):
    return A.Compose(
        [   A.OneOf([
                A.OneOf([   
                A.MotionBlur(p=0.9), 
                #A.MultiplicativeNoise(per_channel=False, p=0.3),
                #A.JpegCompression(quality_lower=98, quality_upper=100, p=0.3),
                ], p=0.7),
                A.OneOf([
                #A.HueSaturationValue(hue_shift_limit=2, p=0.1),
                #A.RandomBrightness(p=0.4),
                #A.RandomContrast(p=0.4),
                A.IAASharpen(p=0.5)
                ], p=0.1)
            ])
        ],
        p=1.0,
    )
    
def sort_images():
    os.mkdir("/kaggle/working/train/")
    os.mkdir("/kaggle/working/train/CBB")
    os.mkdir("/kaggle/working/train/CBSD")
    os.mkdir("/kaggle/working/train/CGM")
    os.mkdir("/kaggle/working/train/CMD")
    os.mkdir("/kaggle/working/train/Healthy")
    
    switch = {
        0 : "CBB",
        1 : "CBSD",
        2 : "CGM",
        3 : "CMD",
        4 : "Healthy" 
    }
    
    data = pd.read_csv('../input/cassava-leaf-disease-classification/train.csv')
    images = data["image_id"]
    lis = []
    for image_name in images:
        cat = data[data["image_id"] == image_name].label.values[0]
        shutil.copy("../input/cassava-leaf-disease-classification/train_images/%s" % (image_name), "../working/train/%s/%s" % (switch[cat], image_name))
    print("Images sorted")   
    
if __name__ == '__main__':
    main()