"""

CNN model of Digit Recognizer using Pytoch

Hoon Kim (id : KeyboardWarrior6969)

13/1/2020

"""
"""
**1. Introduction**

**2. Data Preparation**

**3. CNN Model**

- 3.1 Annealer
- 3.2 Model Definition
- 3.3 Optimizer

**4. Model Performance**

- 4.1 Validation and Training curves
"""
"""
### 1. Introduction

The aim of this notebook is to showcase a simple CNN build using the keras framework. 10 epochs were run utilising the Kaggle GPU architecture.
"""
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np 
from PIL import Image
import seaborn as sns
import cv2
import torch
from torchvision import models
from PIL import Image

import torch
import torch.nn as nn
from torch.autograd import Variable
from torch.utils.data import Dataset, DataLoader, TensorDataset
from sklearn.metrics import confusion_matrix
from albumentations import (
    HorizontalFlip, VerticalFlip, IAAPerspective, ShiftScaleRotate, CLAHE, RandomRotate90,
    Transpose, ShiftScaleRotate, Blur, OpticalDistortion, GridDistortion, HueSaturationValue,
    IAAAdditiveGaussianNoise, GaussNoise, MotionBlur, MedianBlur, IAAPiecewiseAffine, RandomResizedCrop,
    IAASharpen, IAAEmboss, RandomBrightnessContrast, Flip, OneOf, Compose, Normalize, Cutout, CoarseDropout, ShiftScaleRotate, CenterCrop, Resize
)
from albumentations.pytorch import ToTensorV2
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
torch.manual_seed(1234)
if device =='cuda':
    torch.cuda.manual_seed_all(1234)
"""
### 2. Data Preparation
"""
PATH = '../input/cassava-leaf-disease-classification'
train_img_path = PATH+ '/train_images'
test_img_ath = PATH + '/test_images'
# Load dataset
df_train = pd.read_csv(PATH + '/train.csv')

def path_extracter(df):
    x = [train_img_path + '/' + str(k) for k in df.image_id]
    y = [s for s in df.label]
    
    return x,y
x_train,y_train = path_extracter(df_train)
del df_train
from sklearn.model_selection import train_test_split
train_idx, val_idx = train_test_split(list(range(len(x_train))), test_size = 0.2)
x_val = [x_train[k] for k in val_idx]
y_val = [y_train[k] for k in val_idx]

x_train = [x_train[k] for k in train_idx]
y_train = [y_train[k] for k in train_idx]
"""
### 3. CNN Model
"""
"""
#### 3.1 Annealer
"""
class dataset(torch.utils.data.Dataset):
     
        def __init__(self, file_list, labels, transform = None):
            self.file_list = file_list
            self.transforms = transform
            self.label_list = labels
            
            
        def __len__(self):
            self.filelength = len(self.file_list)
            return self.filelength
        
        def __getitem__(self,idx):
            img_path = self.file_list[idx]
            x = cv2.imread(img_path)
            if self.transforms:
                x = self.transforms(image=x)['image']
            
            if self.label_list[idx] == 1:
                label = 1
            else:
                label = 0
                
            return x, label
img_transforms = Compose([
            RandomResizedCrop(224,224),
            Transpose(p=0.5),
            HorizontalFlip(p=0.5),
            VerticalFlip(p=0.5),
            ShiftScaleRotate(p=0.5),
            HueSaturationValue(hue_shift_limit=0.2, sat_shift_limit=0.2, val_shift_limit=0.2, p=0.5),
            RandomBrightnessContrast(brightness_limit=(-0.1,0.1), contrast_limit=(-0.1, 0.1), p=0.5),
            Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], max_pixel_value=255.0, p=1.0),
            CoarseDropout(p=0.5),
            Cutout(p=0.5),
            ToTensorV2(p=1.0),
        ], p=1.)
train_data = dataset(x_train, y_train, transform = img_transforms)
val_data  = dataset(x_val, y_val, transform = img_transforms)

train_loader = torch.utils.data.DataLoader(dataset = train_data, batch_size=16, shuffle=True )
val_loader = torch.utils.data.DataLoader(dataset = val_data, batch_size=16, shuffle=True )

"""
#### 3.2 Model Definition
"""
class Net(nn.Module):
    
    def __init__(self):
        super(Net, self).__init__()
        
        self.layer1 = nn.Sequential(
            nn.Conv2d(3,16,kernel_size=3, padding=0,stride=2),
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        
        self.layer2 = nn.Sequential(
            nn.Conv2d(16,32, kernel_size=3, padding=0, stride=2),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(2)
            )
        
        self.layer3 = nn.Sequential(
            nn.Conv2d(32,64, kernel_size=3, padding=0, stride=2),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        
        
        self.fc1 = nn.Linear(3*3*64,10)
        self.dropout = nn.Dropout(0.5)
        self.fc2 = nn.Linear(10,5)
        self.relu = nn.ReLU()
        
        
    def forward(self,x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = self.layer3(out)
        out = out.view(out.size(0),-1)
        out = self.relu(self.fc1(out))
        out = self.fc2(out)
        return out
"""
#### 3.3 Optimizer
"""
# ----------------------------------------
# Initialize model
model = Net().to(device)
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-3)
model.train()
# ----------------------------------------
history = {'training_acc' : [],'val_loss' : [],'val_acc' : [],'train_loss' : []}
def train(epoch):
    epoch_loss = 0
    epoch_accuracy = 0
    
    for data, label in train_loader:
        data = data.to(device)
        label = label.to(device)
        
        output = model(data)
        loss = criterion(output, label)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        acc = ((output.argmax(dim=1) == label).float().mean())
        epoch_accuracy += acc/len(train_loader)
        epoch_loss += loss/len(train_loader)
        
    print('Epoch : {}, train accuracy : {}, train loss : {}'.format(epoch+1, epoch_accuracy,epoch_loss))
    
    
    with torch.no_grad():
        epoch_val_accuracy=0
        epoch_val_loss =0
        for data, label in val_loader:
            data = data.to(device)
            label = label.to(device)
            
            val_output = model(data)
            val_loss = criterion(val_output,label)
            
            
            acc = ((val_output.argmax(dim=1) == label).float().mean())
            epoch_val_accuracy += acc/ len(val_loader)
            epoch_val_loss += val_loss/ len(val_loader)
            
        print('Epoch : {}, val_accuracy : {}, val_loss : {}'.format(epoch+1, epoch_val_accuracy,epoch_val_loss))
    history['training_acc'].append(epoch_accuracy)
    history['val_acc'].append(epoch_val_accuracy)
    history['train_loss'].append(epoch_loss)
    history['val_loss'].append(epoch_val_loss)


for epoch in range(0, 1):
    train(epoch)


"""
### 4. Model Performance
"""
"""
#### 4.1 Training and Validation Curves
"""
df = pd.DataFrame.from_dict(history, orient='columns')
# Plot the loss and accuracy curves for training and validation 
fig, ax = plt.subplots(2,1)
ax[0].plot(df.train_loss, color='b', label="Training loss")
ax[0].plot(df.val_loss, color='r', label="validation loss",axes =ax[0])
legend = ax[0].legend(loc='best', shadow=True)

ax[1].plot(df.training_acc, color='b', label="Training accuracy")
ax[1].plot(df.val_acc, color='r',label="Validation accuracy")
legend = ax[1].legend(loc='best', shadow=True)