"""
# Reduce training time by using npy files
"""
"""
This notebook is an attempt at reducing training times by using npy image files instead of raw jpg images. The notebook itself is mostly adapted from 
pestipeti wonderful [starter notebook](https://www.kaggle.com/pestipeti/cassava-pytorch-starter-train).


## Training time comparison

### cv2 2-fold training 
- fold0 - train_set(4:03 mins), val_set(2:01 mins)
- fold1 - train_set(3:47 mins), val_set(1:51 mins)

### npy 2-fold training
- fold0 - train_set(1:28 mins), val_set(1:55 mins)
- fold1 - train_set(1:05 mins), val_set(1:09 mins)


Any feedback on the method and dataset would be most welcome.`
"""
"""
## References

1. [plot_confusion_matrix](https://deeplizard.com/learn/video/0LhiS6yu2qQ)
2. [sklearn metrics example](https://towardsdatascience.com/confusion-matrix-for-your-multi-class-machine-learning-model-ff9aa3bf7826)
3. [multi_class_classification](https://towardsdatascience.com/multi-class-classification-extracting-performance-metrics-from-the-confusion-matrix-b379b427a872)
"""
"""
## Library imports
"""
# basic imports
import os
import numpy as np
import pandas as pd
import random
import itertools
from tqdm.notebook import tqdm
import math

# augumentations library
from albumentations.pytorch import ToTensorV2
import albumentations as A
import cv2

# DL library imports
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader

# metrics calculation
from sklearn.metrics import accuracy_score, confusion_matrix
from sklearn.model_selection import KFold, StratifiedKFold

# basic plotting library
import matplotlib.pyplot as plt

# interactive plots
import plotly.express as px
import plotly.graph_objects as go
from plotly.subplots import make_subplots

import warnings  
warnings.filterwarnings('ignore')
"""
temp = torch.Tensor([4.0, 30.0, 23.0])
device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device('cpu')
temp = temp.to(device)
np.save('temp.npy', temp.cpu().data.numpy())
a = np.load('temp.npy')
"""
#!rm -rf /output/kaggle/working/*
"""
## Config files
"""
pipeline = {'train' : True, 'lr_find' : False, 'test' : False}
model_cfg = {'model_architecture': 'resnet50', 'model_name': 'cv2_npy_test_v1',
             'init_lr': 3e-4, 'weight_path': '../input/cassava-pytorch-r50-baseline'}

train_cfg = {'batch_size': 16, 'shuffle': False, 'num_workers': 4, 'checkpt_every' : 1 }
valid_cfg = {'batch_size': 16, 'shuffle': False, 'num_workers': 4, 'validate_every' : 1 }
test_cfg  = {'batch_size': 16, 'shuffle': False, 'num_workers': 4}

DIR_INPUT = '../input/cassava-leaf-disease-classification'
NPY_FOLDER = '../input/cassava-npy-train-images/train_npy_images'
SEED = 42
N_FOLDS = 2 #if pipeline['DEBUG'] else 5
N_EPOCHS = 1 #if pipeline['DEBUG'] else 10
BATCH_SIZE = 32
SIZE = 256
index_label_map = {
                0: "Cassava Bacterial Blight (CBB)", 
                1: "Cassava Brown Streak Disease (CBSD)",
                2: "Cassava Green Mottle (CGM)", 
                3: "Cassava Mosaic Disease (CMD)", 
                4: "Healthy"
                }

class_names = [value for key,value in index_label_map.items()]
"""
## Helper functions
"""
def find_no_of_trainable_params(model):
    total_trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    #print(total_trainable_params)
    return total_trainable_params
def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    os.environ["PYTHONHASHSEED"] = str(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    
#RANDOM_STATE = 42
set_seed(SEED)
def plot_confusion_matrix(cm, classes, normalize=False, title='Confusion matrix', cmap=plt.cm.Blues):
    if normalize:
        cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
        print("Normalized confusion matrix")
    else:
        print('Confusion matrix, without normalization')

    #print(cm)
    plt.imshow(cm, interpolation='nearest', cmap=cmap)
    plt.title(title)
    #plt.colorbar()
    tick_marks = np.arange(len(classes))
    plt.xticks(tick_marks, classes, rotation=45)
    plt.yticks(tick_marks, classes)

    fmt = '.2f' if normalize else 'd'
    thresh = cm.max() / 2.
    for i, j in itertools.product(range(cm.shape[0]), range(cm.shape[1])):
        plt.text(j, i, format(cm[i, j], fmt), horizontalalignment="center", color="white" if cm[i, j] > thresh else "black")

    plt.tight_layout()
    plt.ylabel('True label')
    plt.xlabel('Predicted label')
"""
## Dataset 
"""
train_df = pd.read_csv(f'{DIR_INPUT}/train.csv')
train_df[['cls0', 'cls1', 'cls2', 'cls3', 'cls4']] = train_labels = pd.get_dummies(train_df.iloc[:, 1])
train_df['npy_image_id'] = train_df['image_id'].str.replace('jpg', 'npy')
train_labels = train_df.iloc[:, 1].values
print(train_df.shape)
train_df.head()
class CassavaDataset(Dataset):
    
    def __init__(self, df, dataset='train', transforms=None):
        self.df = df
        self.transforms=transforms
        self.dataset=dataset
        
    def __len__(self):
        return self.df.shape[0]
    
    def __getitem__(self, idx):        
        #image_src = f'{DIR_INPUT}/{self.dataset}_images/{self.df.loc[idx, "image_id"]}'
        #image = cv2.imread(image_src, cv2.IMREAD_COLOR)
        #image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        image = np.load(f'{NPY_FOLDER}/{self.df.loc[idx, "npy_image_id"]}')

        if self.dataset == 'train':
            labels = self.df.loc[idx, ['cls0', 'cls1', 'cls2', 'cls3', 'cls4']].values
            labels = torch.from_numpy(labels.astype(np.int8))
            labels = labels.unsqueeze(-1)
        
        else:
            labels = torch.Tensor(1)
        
        if self.transforms:
            transformed = self.transforms(image=image)
            image = transformed['image']

        return image, labels
"""
## Transforms for Augumentations
"""
transforms_train = A.Compose([
    A.RandomResizedCrop(height=SIZE, width=SIZE, p=1.0),
    #A.Flip(),
    #A.ShiftScaleRotate(rotate_limit=1.0, p=0.8),
    A.Normalize([0.4303133, 0.49675637, 0.3135656], 
                         [0.2379062, 0.24065569, 0.22874062], p=1.0),
    ToTensorV2(p=1.0),
])

transforms_valid = A.Compose([
    A.Resize(height=SIZE, width=SIZE, p=1.0),
    A.Normalize([0.4303133, 0.49675637, 0.3135656], 
                         [0.2379062, 0.24065569, 0.22874062], p=1.0),
    ToTensorV2(p=1.0),
])

transforms_test = A.Compose([
    A.Resize(height=SIZE, width=SIZE, p=1.0),
    A.Normalize([0.4303133, 0.49675637, 0.3135656], 
                         [0.2379062, 0.24065569, 0.22874062], p=1.0),
    ToTensorV2(p=1.0),
])
"""
## CV strategy
"""
folds = StratifiedKFold(n_splits=N_FOLDS, shuffle=True, random_state=SEED)
oof_preds = np.zeros((train_df.shape[0],))
"""
## Model class
"""
class CassavaModel(nn.Module):
    
    def __init__(self, num_classes=5, use_pretrained_weights=False):
        super().__init__()
        self.backbone = torchvision.models.resnet18(pretrained=use_pretrained_weights)
        in_features = self.backbone.fc.in_features
        self.logit = nn.Linear(in_features, num_classes)
        
    def forward(self, x):
        batch_size, C, H, W = x.shape
        
        x = self.backbone.conv1(x)
        x = self.backbone.bn1(x)
        x = self.backbone.relu(x)
        x = self.backbone.maxpool(x)

        x = self.backbone.layer1(x)
        x = self.backbone.layer2(x)
        x = self.backbone.layer3(x)
        x = self.backbone.layer4(x)
        
        x = F.adaptive_avg_pool2d(x,1).reshape(batch_size,-1)
        x = F.dropout(x, 0.25, self.training)

        x = self.logit(x)

        return x
"""
## Loss function
"""
class DenseCrossEntropy(nn.Module):

    def __init__(self):
        super(DenseCrossEntropy, self).__init__()
        
        
    def forward(self, logits, labels):
        logits = logits.float()
        labels = labels.float()
        
        logprobs = F.log_softmax(logits, dim=-1)
        
        loss = -labels * logprobs
        loss = loss.sum(-1)

        return loss.mean()

# creating loss function instance
criterion = DenseCrossEntropy()
#criterion = nn.CrossEntropyLoss()
## Device as cpu or tpu
device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device('cpu')
print(device)
"""
## Lr_find
"""
def plot_lr_finder_results(lr_finder): 
    # Create subplot grid
    fig = make_subplots(rows=1, cols=2)
    # layout ={'title': 'Lr_finder_result'}
    
    # Create a line (trace) for the lr vs loss, gradient of loss
    trace0 = go.Scatter(x=lr_finder['log_lr'], y=lr_finder['smooth_loss'],name='log_lr vs smooth_loss')
    trace1 = go.Scatter(x=lr_finder['log_lr'], y=lr_finder['grad_loss'],name='log_lr vs loss gradient')

    # Add subplot trace & assign to each grid
    fig.add_trace(trace0, row=1, col=1);
    fig.add_trace(trace1, row=1, col=2);
    #iplot(fig, show_link=False)
    fig.write_html(model_cfg['model_name'] + '_lr_find.html');
def find_lr(model, data_loader, optimizer, init_value = 1e-8, final_value=100.0, beta = 0.98, num_batches = 200):
    assert(num_batches > 0)
    mult = (final_value / init_value) ** (1/num_batches)
    lr = init_value
    optimizer.param_groups[0]['lr'] = lr
    batch_num = 0
    avg_loss = 0.0
    best_loss = 0.0
    smooth_losses = []
    raw_losses = []
    log_lrs = []
    dataloader_it = iter(data_loader)
    progress_bar = tqdm(range(num_batches))
        
    for idx in progress_bar:
        batch_num += 1
        try:
            images, labels = next(dataloader_it)
            #print(images.shape)
        except:
            dataloader_it = iter(data_loader)
            images, labels = next(dataloader_it)

        # Move input and label tensors to the default device
        images = images.to(device, dtype=torch.float)
        labels = labels.to(device, dtype=torch.float)
        
        # handle exception in criterion
        try:
            # Forward pass
            log_ps = model(images)
            loss = criterion(log_ps, labels.squeeze(-1))
        except:
            if len(smooth_losses) > 1:
                grad_loss = np.gradient(smooth_losses)
            else:
                grad_loss = 0.0
            lr_finder_results = {'log_lr':log_lrs, 'raw_loss':raw_losses, 
                                 'smooth_loss':smooth_losses, 'grad_loss': grad_loss}
            return lr_finder_results 
                    
        #Compute the smoothed loss
        avg_loss = beta * avg_loss + (1-beta) *loss.item()
        smoothed_loss = avg_loss / (1 - beta**batch_num)
        
        #Stop if the loss is exploding
        if batch_num > 1 and smoothed_loss > 50 * best_loss:
            if len(smooth_losses) > 1:
                grad_loss = np.gradient(smooth_losses)
            else:
                grad_loss = 0.0
            lr_finder_results = {'log_lr':log_lrs, 'raw_loss':raw_losses, 
                                 'smooth_loss':smooth_losses, 'grad_loss': grad_loss}
            return lr_finder_results
        
        #Record the best loss
        if smoothed_loss < best_loss or batch_num==1:
            best_loss = smoothed_loss
        
        #Store the values
        raw_losses.append(loss.item())
        smooth_losses.append(smoothed_loss)
        log_lrs.append(math.log10(lr))
        
        # Backward pass
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        # print info
        progress_bar.set_description(f"loss: {loss.item()},smoothed_loss: {smoothed_loss},lr : {lr}")

        #Update the lr for the next step
        lr *= mult
        optimizer.param_groups[0]['lr'] = lr
    
    grad_loss = np.gradient(smooth_losses)
    lr_finder_results = {'log_lr':log_lrs, 'raw_loss':raw_losses, 
                         'smooth_loss':smooth_losses, 'grad_loss': grad_loss}
    return lr_finder_results
if pipeline['lr_find'] == True:
    # create Dataset
    temp_train_dataset = CassavaDataset(df=train_df, dataset='train', transforms=transforms_train)
    temp_train_dataloader = DataLoader(temp_train_dataset, batch_size=BATCH_SIZE, num_workers=4, shuffle=True)
    
    # create model instance
    model = CassavaModel(num_classes=5, use_pretrained_weights=True)
    model.to(device)
    optimizer = optim.Adam(model.parameters(), lr=model_cfg['init_lr'])
    
    lr_finder_results = find_lr(model, temp_train_dataloader, optimizer)
    plot_lr_finder_results(lr_finder_results)
"""
## One fold train and validation function
"""
def train_one_fold(i_fold, model, optimizer, dataloader_train, dataloader_valid):
    
    train_fold_results = []

    for epoch in range(N_EPOCHS):
        print('  Epoch {}/{}'.format(epoch + 1, N_EPOCHS))
        model.train()
        tr_loss = 0
        lr_list = []
        
        # training iterator
        tr_iterator = iter(dataloader_train)
        train_progress_bar = tqdm(range(len(dataloader_train)))
    
        for idx in train_progress_bar:
            try:
                images, labels = next(tr_iterator)
                #print(images.shape)
            except StopIteration:
                tr_iterator = iter(dataloader_train)
                images, labels = next(tr_iterator)

            images = images.to(device, dtype=torch.float)
            labels = labels.to(device, dtype=torch.float)
            
            # Forward pass
            outputs = model(images)

            # Backward pass
            loss = criterion(outputs, labels.squeeze(-1))                
            tr_loss += loss.item()
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
            
            lr_list.append(optimizer.state_dict()["param_groups"][0]['lr'])

            # print to console
            train_progress_bar.set_description(f"Train_loss: {tr_loss} loss(avg): {tr_loss/(idx+1)}")
        
        lr_list = np.array(lr_list)
        np.save(model_cfg['model_name'] + '_' + str(i_fold) + 'fold_lr_list.npy', lr_list)

        # Validate
        model.eval()
        val_loss = 0.0
        val_preds = None
        val_labels = None
        
        valid_iterator = iter(dataloader_valid)
        valid_progress_bar = tqdm(range(len(dataloader_valid)))

        for idx in valid_progress_bar:
            try:
                images, labels = next(valid_iterator)
            except StopIteration:
                tr_iterator = iter(dataloader_valid)
                images, labels = next(valid_iterator)

            if val_labels is None:
                val_labels = labels.clone().squeeze(-1)
            else:
                val_labels = torch.cat((val_labels, labels.squeeze(-1)), dim=0)

            images = images.to(device, dtype=torch.float)
            labels = labels.to(device, dtype=torch.float)

            with torch.no_grad():
                outputs = model(images)

                loss = criterion(outputs, labels.squeeze(-1))
                val_loss += loss.item()

                preds = torch.softmax(outputs, dim=1).data.cpu()

                if val_preds is None:
                    val_preds = preds
                else:
                    val_preds = torch.cat((val_preds, preds), dim=0)
            
            # print to console
            valid_progress_bar.set_description(f"val_loss: {val_loss} loss(avg): {val_loss/(idx+1)}")

        
        val_preds = torch.argmax(val_preds, dim=1)
        val_labels = torch.argmax(val_labels, dim=1)
        
        np.save(model_cfg['model_name'] + '_val_preds_' + str(i_fold) + '.npy', val_preds.cpu().data.numpy())
        np.save(model_cfg['model_name'] + '_val_labels_' + str(i_fold) + '.npy', val_labels.cpu().data.numpy())

        train_fold_results.append({
            'fold': i_fold,
            'epoch': epoch,
            'train_loss': tr_loss / len(dataloader_train),
            'valid_loss': val_loss / len(dataloader_valid),
            'valid_score': accuracy_score(val_labels, val_preds)
        })

    return val_preds, train_fold_results
"""
## Training and validation function calls
"""
if pipeline['train'] == True:
    submissions = None
    train_results = []

    for i_fold, (train_idx, valid_idx) in enumerate(folds.split(train_df, train_labels)):
        print("Fold {}/{}".format(i_fold + 1, N_FOLDS))

        valid = train_df.iloc[valid_idx]
        valid.reset_index(drop=True, inplace=True)

        train = train_df.iloc[train_idx]
        train.reset_index(drop=True, inplace=True)    

        dataset_train = CassavaDataset(df=train, dataset='train', transforms=transforms_train)
        dataset_valid = CassavaDataset(df=valid, dataset='train', transforms=transforms_valid)

        dataloader_train = DataLoader(dataset_train, batch_size=BATCH_SIZE, num_workers=4, shuffle=True)
        dataloader_valid = DataLoader(dataset_valid, batch_size=BATCH_SIZE, num_workers=4, shuffle=False)

        model = CassavaModel(num_classes=5, use_pretrained_weights=True)
        model.to(device)

        plist = [{'params': model.parameters(), 'lr': 1e-4}]
        optimizer = optim.Adam(plist, lr=model_cfg['init_lr'])

        val_preds, train_fold_results = train_one_fold(i_fold, model, optimizer, 
                                                       dataloader_train, dataloader_valid)
        oof_preds[valid_idx] = val_preds.numpy()
        train_results = train_results + train_fold_results

        torch.save({
            'fold': i_fold,
            'lr': optimizer.state_dict()["param_groups"][0]['lr'],
            'model_state_dict': model.cpu().state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            # 'scheduler_state_dict'
            # 'scaler_state_dict'
        }, f"{model_cfg['model_name']}_fold_{i_fold}.pth")

    print("{}-Folds CV score: {:.4f}".format(N_FOLDS, accuracy_score(train_labels, oof_preds)))
"""
## Plot training results
"""
def plot_training_results():
    fig = make_subplots(rows=2, cols=1)

    colors = [
        ('#d32f2f', '#ef5350'),
        ('#303f9f', '#5c6bc0'),
        ('#00796b', '#26a69a'),
        ('#fbc02d', '#ffeb3b'),
        ('#5d4037', '#8d6e63'),
    ]

    for i in range(N_FOLDS):
        data = train_results[train_results['fold'] == i]

        fig.add_trace(go.Scatter(x=data['epoch'].values,
                                 y=data['train_loss'].values,
                                 mode='lines',
                                 visible='legendonly' if i > 0 else True,
                                 line=dict(color=colors[i][0], width=2),
                                 name='Train loss - Fold #{}'.format(i)),
                     row=1, col=1)

        fig.add_trace(go.Scatter(x=data['epoch'],
                                 y=data['valid_loss'].values,
                                 mode='lines+markers',
                                 visible='legendonly' if i > 0 else True,
                                 line=dict(color=colors[i][1], width=2),
                                 name='Valid loss - Fold #{}'.format(i)),
                     row=1, col=1)

        fig.add_trace(go.Scatter(x=data['epoch'].values,
                                 y=data['valid_score'].values,
                                 mode='lines+markers',
                                 line=dict(color=colors[i][0], width=2),
                                 name='Valid score - Fold #{}'.format(i),
                                 showlegend=False),
                     row=2, col=1)

    fig.update_layout({
      "annotations": [
        {
          "x": 0.225, 
          "y": 1.0, 
          "font": {"size": 16}, 
          "text": "Train / valid losses", 
          "xref": "paper", 
          "yref": "paper", 
          "xanchor": "center", 
          "yanchor": "bottom", 
          "showarrow": False
        }, 
        {
          "x": 0.775, 
          "y": 1.0, 
          "font": {"size": 16}, 
          "text": "Validation scores", 
          "xref": "paper", 
          "yref": "paper", 
          "xanchor": "center", 
          "yanchor": "bottom", 
          "showarrow": False
        }, 
      ]
    })

    fig.show()
"""
val_preds_0 = np.load('./R18_imagenet_v2_val_preds_0.npy')
val_labels_0 = np.load('./R18_imagenet_v2_val_labels_0.npy')

cm = confusion_matrix(val_labels_0, val_preds_0)
print(cm)
plt.figure(figsize=(8,8))
plot_confusion_matrix(cm, classes=class_names, normalize=True)
"""
if pipeline['train'] == True:
    train_results = pd.DataFrame(train_results)
    print(train_results.head(10))
    
    final_results = train_results[train_results['epoch']==train_results['epoch'].max()]
    print(final_results)
    print(final_results['valid_score'].mean(), final_results['valid_score'].std())
    plot_training_results()
"""
## Testing function
"""
if pipeline["test"] == True:
    # read submission file
    submission_df = pd.read_csv(DIR_INPUT + '/sample_submission.csv')
    submission_df.iloc[:, 1] = 4
    #print(submission_df.head())
    submission_df.to_csv('submission.csv', index=False)
"""
if pipeline["test"] == True:
    # read submission file
    submission_df = pd.read_csv(DIR_INPUT + '/sample_submission.csv')
    submission_df.iloc[:, 1] = 0
    #print(submission_df.head())


    # just for debugging purporse, adding 1 more row
    if submission_df.shape[0] == 1:
        submission_df = pd.DataFrame([{'image_id': '2216849948.jpg', 'label': 0},{'image_id': '2216849948.jpg', 'label': 0}])
        submission_df.reset_index(drop=True, inplace=True)
    #print(submission_df.head())


    # Creating test dataset and dataloaders
    dataset_test = CassavaDataset(df=submission_df, dataset='test', transforms=transforms_test)
    dataloader_test = DataLoader(dataset_test, batch_size=BATCH_SIZE, num_workers=4, shuffle=False)
    
    
    # placeholder for final submission csv
    submissions = None

    """
    1. Iterate and store predictions (one-hot encoded format) of N-folds of model 
    2. Average the predictions of all folds
    3. argmax of mean one-hot encoded prediction is output
    """
    for i_fold in range(N_FOLDS):
        print(f'Inference for {i_fold}th fold')
        model = CassavaModel(num_classes=5, use_pretrained_weights=False)
        model.to(device)

        checkpoint = torch.load(f"{model_cfg['weight_path']}/{model_cfg['model_name']}_fold_{i_fold}.pth", map_location=device)
        model.load_state_dict(checkpoint['model_state_dict'], strict=True)
        model.eval()
        test_preds = None

        for step, (images, _) in enumerate(dataloader_test):
            images = images.to(device, dtype=torch.float)
            with torch.no_grad():
                outputs = model(images)
                preds = torch.softmax(outputs, dim=1).data.cpu()
                if test_preds is None:
                    test_preds = preds
                else:
                    test_preds = torch.cat((test_preds, preds), dim=0)

        # submission_df[['label']] = test_preds.argmax(test_preds, dim=1)
        # submission_df.to_csv('submission_fold_{}.csv'.format(i_fold), index=False)

        # logits avg
        if submissions is None:
            submissions = test_preds / N_FOLDS
        else:
            submissions += test_preds / N_FOLDS
            
        
    #print(submissions[:10])
    # argmax of predictions and write to csv
    submission_df['label'] = torch.argmax(submissions, dim=1)
    submission_df.to_csv('submission.csv', index=False)
    #print(submission_df.head())
"""