# IMPORTS

# Imports from the standard library
from collections import OrderedDict
import json, sys, time, argparse, copy

# Imports from basic additional libraries
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image

# Pytorch imports
import torch
from torch import nn, optim
import torch.nn.functional as F
from torchvision import datasets, transforms, models

# GLOBAL CONSTANTS

# Constants regarding the dataset
train_data_dir = 'flowers/train'
validation_data_dir = 'flowers/valid'
testing_data_dir = 'flowers/test'
no_of_categories = 102
available_base_models = ['alexnet', 'densenet121', 'vgg11', 'vgg11_bn', 'vgg19']

# Constants for torchvision dataset normalization
std_norm_mean = [0.485, 0.456, 0.406]
std_norm_std = [0.229, 0.224, 0.225]

# Constants for batches
_std_batch_size = 32
_std_shuffle = False

# Constants for model building
_std_base_model = 'alexnet'
_std_hidden_1_out = 256
_std_hidden_2_out = 256
_std_final_out = no_of_categories
_std_dropout_rate = 0.25

# Constants for optimizer
_std_learning_rate = 1e-3

# Constants for training session
_std_epochs = 5
_std_monitor = True

# Constants for infrastructure
_std_device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# PROGRAM

def create_model_from_checkpoint(checkpoint_path='best_model.pth', 
                                 cat_to_name_path='cat_to_name.json', 
                                 device=_std_device):
    '''
    Description:
    ============
    
    Keyword arguments:
    - checkpoint_path:
      Path to checkpoint as str (default: best_model.pth)
    - cat_to_name_path:
      Path to cat_to_name_file as str (default: cat_to_name.json)
    - device:
      device for the model (default: 'cuda:0' if available, else 'cpu')
    Returns:
    - model:
      torch model
    - criterion
      loss function for torch model
    - optimizer
      optimizer for torch model
    - name_from_class
      dict for retrieves names from classes
    '''
    
    # Load the checkpoint
    if torch.cuda.is_available() and device=='cuda:0':
        map_location=lambda storage, loc: storage.cuda()
    else:
        map_location='cpu'
    
    model_and_optim_specs = torch.load(checkpoint_path, map_location=map_location)
    
    # Create model incl. criterion and optimizer
    model = model_and_optim_specs['model']
    # model.to(device) # Push model to device
    criterion = model_and_optim_specs['criterion']
    optimizer = model_and_optim_specs['optimizer']
    
    # Create mapping to names of classes
    if cat_to_name_path:
        with open(cat_to_name_path, 'r') as f:
            cat_to_name = json.load(f)
        
        class_to_idx = model_and_optim_specs['class_mapping']
        name_from_class = {value: cat_to_name[key] for key, value in model_and_optim_specs['class_mapping'].items()}
    else:
        name_from_class = None
    
    return model, criterion, optimizer, name_from_class

def _process_image(image_path, cropping=True, device=_std_device):
    '''
    Description:
    ============
    Helper function.
    Retrieve an image as a torch tensor from an directory path. Allows cropping
    to 224x224 to be turned off for showing the image.
    
    Arguments:
    - image_path:
      Path to image as string
    Keyword Arguments:
    - cropping:
      Defines if cropping to 224x224 is being applies which
      is necessary for prediction, but not for showing the image (default: True)
    - device:
      device for the model (default: 'cuda:0' if available, else 'cpu')
      
    Returns:
    - image:
      Torch tensor of the image
    '''
    
    # Open the image
    image = Image.open(image_path)
    
    # Define transform as requested (default)
    if cropping:
        transform = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(mean=std_norm_mean, std=std_norm_std)
        ])
    else:
        transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(mean=std_norm_mean, std=std_norm_std)
        ])
    
    # Transform the image
    image = transform(image)
    image = image.to(device)
    
    return image

def _imshow(image, 
            ax=None, 
            title=None):
    
    '''
    Description:
    ============
    Helper function.
    Produces an matplotlib imshow from an torch image.
    ax for subplot and title of image can be defined.
    
    Arguments:
    - image:
      image as torch tensor
    Keyword Arguments:
    - ax:
      Can be defined if image shall be plotted as subplot (default: None)
    - title:
      Can be defined if title shall be printed as string
    Returns:
    - ax:
      matplotlib axis
    '''

    if ax is None:
        fig, ax = plt.subplots()
    
    # PyTorch tensors assume the color channel is the first dimension
    # but matplotlib assumes is the third dimension
    image = image.numpy().transpose((1, 2, 0))
    
    # Undo preprocessing
    mean = np.array(std_norm_mean)
    std = np.array(std_norm_std)
    image = std * image + mean
    
    # Image needs to be clipped between 0 and 1 or it looks like noise when displayed
    image = np.clip(image, 0, 1)
    
    # Showing the image
    ax.imshow(image)
    ax.axis('off')
    ax.title.set_text(title)
    
    return ax

def predict(image_path, 
            checkpoint_path='best_model.pth', 
            cat_to_name_path='cat_to_name.json',
            k=5,
            device=_std_device):
    
    '''
    Description:
    ============
    Given a model, an image path and a dict to map from classes to flower names, 
    this function predicts the top k classes with names. Prediction can be process
    on device
    
    Arguments:
    - image_path:
      Path to image as string
    Keyword Arguments:
    - checkpoint_path:
      Path to checkpoint containing model, criterion, optimizer and class_to_idx as str
      (default: best_model.pth)
    - cat_to_name_path:
      Path to json containing cat to name dict as str (default: cat_to_name.json)
    - k:
      Number of classes
    - device:
      device for the model (default: 'cuda:0' if available, else 'cpu')
    Returns:
    - top_probabilities:
      list of top k probabilities as float
    - top_classes:
      list of top k classes as int
    - top_names:
      list of top k names as str
    '''
    
    model, criterion, _, name_from_class = create_model_from_checkpoint(checkpoint_path=checkpoint_path, 
                                                                                cat_to_name_path=cat_to_name_path, 
                                                                                device=device)
    
    # Inferring top 5 probabilities and classes for given image
    with torch.no_grad():
        model.eval()
        image = _process_image(image_path, device=device)
        image = image.unsqueeze_(0)
        output = model.forward(image)
        probabilities = torch.exp(output)
        top_probabilities, top_classes = probabilities.topk(k, dim=1)
        model.train()
    
    # Retrieving list of top 5 probabilities, classes and class names
    top_probabilities = top_probabilities.numpy().tolist()[0]
    top_classes = top_classes.numpy().tolist()[0]
    
    # Retrieving names from top 5 classes
    if name_from_class:
        top_names = [name_from_class[top_class] for top_class in top_classes]
    else:
        top_names = None
    
    return top_probabilities, top_classes, top_names

def show_prediction(image_path, 
                    checkpoint_path='best_model.pth',
                    cat_to_name_path='cat_to_name.json', 
                    device=_std_device): 
    '''
    Description:
    ============
    Shows a bar chart with the top five predictions and prints picture with it.
    
    Arguments:
    - image_path:
      Path to image as string
    Keyword Arguments:
    - checkpoint_path:
      Path to checkpoint containing model, criterion, optimizer and class_to_idx as str
      (default: best_model.pth)
    - cat_to_name_path:
      Path to json containing cat to name dict as str (default: cat_to_name.json)
    - device:
      device for the model (default: 'cuda:0' if available, else 'cpu')
    '''
    
    # Calculating top 5 probabilities with class names
    top_probabilities, _, top_names = predict(image_path, 
                                              checkpoint_path=checkpoint_path, 
                                              cat_to_name_path=cat_to_name_path, 
                                              device=_std_device)
    
    # Setting of the chart for top 5 classes for image with probabilities
    fig, ax = plt.subplots(1, 2, figsize=(15,5))
    
    # Showing bar chart of top 5 classes 
    ax[0].barh(list(reversed(top_names)), list(reversed(top_probabilities)))
    ax[0].title.set_text('Top 5 classes for image (with probabilities)')
    
    # Retrieving the image and plotting it with name
    image = _process_image(image_path, cropping=False)
    ax[1] = _imshow(image, title='It\'s probably a(n) ' + top_names[0] + '!', ax=ax[1])

# MAIN FUNCTION

def main(path_to_image, checkpoint, category_names='cat_to_name.json', k=5, device='cpu'):
    '''
    Description:
    ============
    Function runs when predict.py is being executed.
    It reads data from an image, loads a checkpoint and calculates the according
    k highest probabilities and categories giving the category names when a mapping
    file is being passed. Can run on gpu when defined and available.
    
    Arguments:
    - path_to_image:
      Path to the image to predict as str.
    - checkpoint:
      Path to checkpoint/model as str.
    Keyword Arguments:
    - category_names:
      Path to mapping file as str.
      (default: cat_to_name.json)
    - k:
      Number of k highest probabilities and categories to calculate 
      as int (default: 5)
    - device:
      device for the model (default: 'cpu' if available, else 'cpu')
    '''
    
    top_probabilities, top_classes, top_names = predict(path_to_image, 
                                                        checkpoint_path=checkpoint, 
                                                        cat_to_name_path=category_names, 
                                                        k=k,
                                                        device=device)
    
    if top_names:
        print('Top', k, 'classes with probabilities:', *zip(top_classes, top_names), top_probabilities)
    else:
        print('Top', k, 'classes with probabilities:', top_classes, top_probabilities)
        
if __name__ == '__main__':
    
    parser = argparse.ArgumentParser()
    parser.add_argument('path_to_image', action='store', type=str,
                        help='Path to image')
    parser.add_argument('checkpoint', action='store', type=str,
                        help='Path to checkpoint')
    parser.add_argument('--top_k', action='store', type=int, default=1,
                        help='Number of classes with top probabilities to show')
    parser.add_argument('--category_names', action='store', type=str, default=None,
                        help='Name of the file to map classes to names (default: cat_to_name.json)')
    parser.add_argument('--gpu', action='store_true', 
                        help='Device (default: cpu, gpu possible only if available)')
    args = parser.parse_args()
    
    path_to_image = args.path_to_image
    checkpoint = args.checkpoint
    k = args.top_k
    category_names = args.category_names
    device = torch.device("cuda:0" if (torch.cuda.is_available() and args.gpu) else "cpu")  
    
    main(path_to_image, checkpoint, k=k, category_names=category_names, device=device)
# IMPORTS

# Imports from the standard library
from collections import OrderedDict
import json, sys, time, argparse, copy

# Imports from basic additional libraries
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image

# Pytorch imports
import torch
from torch import nn, optim
import torch.nn.functional as F
from torchvision import datasets, transforms, models

# GLOBAL CONSTANTS

# Constants regarding the dataset
train_data_dir = 'flowers/train'
validation_data_dir = 'flowers/valid'
testing_data_dir = 'flowers/test'
no_of_categories = 102
available_base_models = ['alexnet', 'densenet121', 'vgg11', 'vgg11_bn', 'vgg19']

# Constants for torchvision dataset normalization
std_norm_mean = [0.485, 0.456, 0.406]
std_norm_std = [0.229, 0.224, 0.225]

# Constants for batches
_std_batch_size = 32
_std_shuffle = False

# Constants for model building
_std_base_model = 'alexnet'
_std_hidden_1_out = 256
_std_hidden_2_out = 256
_std_final_out = no_of_categories
_std_dropout_rate = 0.25

# Constants for optimizer
_std_learning_rate = 1e-3

# Constants for training session
_std_epochs = 5
_std_monitor = True

# Constants for infrastructure
_std_device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# PROGRAM

def create_model_from_checkpoint(checkpoint_path='best_model.pth', 
                                 cat_to_name_path='cat_to_name.json', 
                                 device=_std_device):
    '''
    Description:
    ============
    
    Keyword arguments:
    - checkpoint_path:
      Path to checkpoint as str (default: best_model.pth)
    - cat_to_name_path:
      Path to cat_to_name_file as str (default: cat_to_name.json)
    - device:
      device for the model (default: 'cuda:0' if available, else 'cpu')
    Returns:
    - model:
      torch model
    - criterion
      loss function for torch model
    - optimizer
      optimizer for torch model
    - name_from_class
      dict for retrieves names from classes
    '''
    
    # Load the checkpoint
    if torch.cuda.is_available() and device=='cuda:0':
        map_location=lambda storage, loc: storage.cuda()
    else:
        map_location='cpu'
    
    model_and_optim_specs = torch.load(checkpoint_path, map_location=map_location)
    
    # Create model incl. criterion and optimizer
    model = model_and_optim_specs['model']
    # model.to(device) # Push model to device
    criterion = model_and_optim_specs['criterion']
    optimizer = model_and_optim_specs['optimizer']
    
    # Create mapping to names of classes
    if cat_to_name_path:
        with open(cat_to_name_path, 'r') as f:
            cat_to_name = json.load(f)
        
        class_to_idx = model_and_optim_specs['class_mapping']
        name_from_class = {value: cat_to_name[key] for key, value in model_and_optim_specs['class_mapping'].items()}
    else:
        name_from_class = None
    
    return model, criterion, optimizer, name_from_class

def _process_image(image_path, cropping=True, device=_std_device):
    '''
    Description:
    ============
    Helper function.
    Retrieve an image as a torch tensor from an directory path. Allows cropping
    to 224x224 to be turned off for showing the image.
    
    Arguments:
    - image_path:
      Path to image as string
    Keyword Arguments:
    - cropping:
      Defines if cropping to 224x224 is being applies which
      is necessary for prediction, but not for showing the image (default: True)
    - device:
      device for the model (default: 'cuda:0' if available, else 'cpu')
      
    Returns:
    - image:
      Torch tensor of the image
    '''
    
    # Open the image
    image = Image.open(image_path)
    
    # Define transform as requested (default)
    if cropping:
        transform = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(mean=std_norm_mean, std=std_norm_std)
        ])
    else:
        transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(mean=std_norm_mean, std=std_norm_std)
        ])
    
    # Transform the image
    image = transform(image)
    image = image.to(device)
    
    return image

def _imshow(image, 
            ax=None, 
            title=None):
    
    '''
    Description:
    ============
    Helper function.
    Produces an matplotlib imshow from an torch image.
    ax for subplot and title of image can be defined.
    
    Arguments:
    - image:
      image as torch tensor
    Keyword Arguments:
    - ax:
      Can be defined if image shall be plotted as subplot (default: None)
    - title:
      Can be defined if title shall be printed as string
    Returns:
    - ax:
      matplotlib axis
    '''

    if ax is None:
        fig, ax = plt.subplots()
    
    # PyTorch tensors assume the color channel is the first dimension
    # but matplotlib assumes is the third dimension
    image = image.numpy().transpose((1, 2, 0))
    
    # Undo preprocessing
    mean = np.array(std_norm_mean)
    std = np.array(std_norm_std)
    image = std * image + mean
    
    # Image needs to be clipped between 0 and 1 or it looks like noise when displayed
    image = np.clip(image, 0, 1)
    
    # Showing the image
    ax.imshow(image)
    ax.axis('off')
    ax.title.set_text(title)
    
    return ax

def predict(image_path, 
            checkpoint_path='best_model.pth', 
            cat_to_name_path='cat_to_name.json',
            k=5,
            device=_std_device):
    
    '''
    Description:
    ============
    Given a model, an image path and a dict to map from classes to flower names, 
    this function predicts the top k classes with names. Prediction can be process
    on device
    
    Arguments:
    - image_path:
      Path to image as string
    Keyword Arguments:
    - checkpoint_path:
      Path to checkpoint containing model, criterion, optimizer and class_to_idx as str
      (default: best_model.pth)
    - cat_to_name_path:
      Path to json containing cat to name dict as str (default: cat_to_name.json)
    - k:
      Number of classes
    - device:
      device for the model (default: 'cuda:0' if available, else 'cpu')
    Returns:
    - top_probabilities:
      list of top k probabilities as float
    - top_classes:
      list of top k classes as int
    - top_names:
      list of top k names as str
    '''
    
    model, criterion, _, name_from_class = create_model_from_checkpoint(checkpoint_path=checkpoint_path, 
                                                                                cat_to_name_path=cat_to_name_path, 
                                                                                device=device)
    
    # Inferring top 5 probabilities and classes for given image
    with torch.no_grad():
        model.eval()
        image = _process_image(image_path, device=device)
        image = image.unsqueeze_(0)
        output = model.forward(image)
        probabilities = torch.exp(output)
        top_probabilities, top_classes = probabilities.topk(k, dim=1)
        model.train()
    
    # Retrieving list of top 5 probabilities, classes and class names
    top_probabilities = top_probabilities.numpy().tolist()[0]
    top_classes = top_classes.numpy().tolist()[0]
    
    # Retrieving names from top 5 classes
    if name_from_class:
        top_names = [name_from_class[top_class] for top_class in top_classes]
    else:
        top_names = None
    
    return top_probabilities, top_classes, top_names

def show_prediction(image_path, 
                    checkpoint_path='best_model.pth',
                    cat_to_name_path='cat_to_name.json', 
                    device=_std_device): 
    '''
    Description:
    ============
    Shows a bar chart with the top five predictions and prints picture with it.
    
    Arguments:
    - image_path:
      Path to image as string
    Keyword Arguments:
    - checkpoint_path:
      Path to checkpoint containing model, criterion, optimizer and class_to_idx as str
      (default: best_model.pth)
    - cat_to_name_path:
      Path to json containing cat to name dict as str (default: cat_to_name.json)
    - device:
      device for the model (default: 'cuda:0' if available, else 'cpu')
    '''
    
    # Calculating top 5 probabilities with class names
    top_probabilities, _, top_names = predict(image_path, 
                                              checkpoint_path=checkpoint_path, 
                                              cat_to_name_path=cat_to_name_path, 
                                              device=_std_device)
    
    # Setting of the chart for top 5 classes for image with probabilities
    fig, ax = plt.subplots(1, 2, figsize=(15,5))
    
    # Showing bar chart of top 5 classes 
    ax[0].barh(list(reversed(top_names)), list(reversed(top_probabilities)))
    ax[0].title.set_text('Top 5 classes for image (with probabilities)')
    
    # Retrieving the image and plotting it with name
    image = _process_image(image_path, cropping=False)
    ax[1] = _imshow(image, title='It\'s probably a(n) ' + top_names[0] + '!', ax=ax[1])

# MAIN FUNCTION

def main(path_to_image, checkpoint, category_names='cat_to_name.json', k=5, device='cpu'):
    '''
    Description:
    ============
    Function runs when predict.py is being executed.
    It reads data from an image, loads a checkpoint and calculates the according
    k highest probabilities and categories giving the category names when a mapping
    file is being passed. Can run on gpu when defined and available.
    
    Arguments:
    - path_to_image:
      Path to the image to predict as str.
    - checkpoint:
      Path to checkpoint/model as str.
    Keyword Arguments:
    - category_names:
      Path to mapping file as str.
      (default: cat_to_name.json)
    - k:
      Number of k highest probabilities and categories to calculate 
      as int (default: 5)
    - device:
      device for the model (default: 'cpu' if available, else 'cpu')
    '''
    
    top_probabilities, top_classes, top_names = predict(path_to_image, 
                                                        checkpoint_path=checkpoint, 
                                                        cat_to_name_path=category_names, 
                                                        k=k,
                                                        device=device)
    
    if top_names:
        print('Top', k, 'classes with probabilities:', *zip(top_classes, top_names), top_probabilities)
    else:
        print('Top', k, 'classes with probabilities:', top_classes, top_probabilities)
        
if __name__ == '__main__':
    
    parser = argparse.ArgumentParser()
    parser.add_argument('path_to_image', action='store', type=str,
                        help='Path to image')
    parser.add_argument('checkpoint', action='store', type=str,
                        help='Path to checkpoint')
    parser.add_argument('--top_k', action='store', type=int, default=1,
                        help='Number of classes with top probabilities to show')
    parser.add_argument('--category_names', action='store', type=str, default=None,
                        help='Name of the file to map classes to names (default: cat_to_name.json)')
    parser.add_argument('--gpu', action='store_true', 
                        help='Device (default: cpu, gpu possible only if available)')
    args = parser.parse_args()
    
    path_to_image = args.path_to_image
    checkpoint = args.checkpoint
    k = args.top_k
    category_names = args.category_names
    device = torch.device("cuda:0" if (torch.cuda.is_available() and args.gpu) else "cpu")  
    
    main(path_to_image, checkpoint, k=k, category_names=category_names, device=device)
# IMPORTS # Imports from the standard library # Imports from basic additional libraries # Pytorch imports # GLOBAL CONSTANTS # Constants regarding the dataset # Constants for torchvision dataset normalization # Constants for batches # Constants for model building # Constants for optimizer # Constants for training session # Constants for infrastructure # PROGRAM Description:     ============          Keyword arguments:     - checkpoint_path:       Path to checkpoint as str (default: best_model.pth)     - cat_to_name_path:       Path to cat_to_name_file as str (default: cat_to_name.json)     - device:       device for the model (default: 'cuda:0' if available, else 'cpu')     Returns:     - model:       torch model     - criterion       loss function for torch model     - optimizer       optimizer for torch model     - name_from_class       dict for retrieves names from classes # Load the checkpoint # Create model incl. criterion and optimizer # model.to(device) # Push model to device # Create mapping to names of classes Description:     ============     Helper function.     Retrieve an image as a torch tensor from an directory path. Allows cropping     to 224x224 to be turned off for showing the image.          Arguments:     - image_path:       Path to image as string     Keyword Arguments:     - cropping:       Defines if cropping to 224x224 is being applies which       is necessary for prediction, but not for showing the image (default: True)     - device:       device for the model (default: 'cuda:0' if available, else 'cpu')            Returns:     - image:       Torch tensor of the image # Open the image # Define transform as requested (default) # Transform the image Description:     ============     Helper function.     Produces an matplotlib imshow from an torch image.     ax for subplot and title of image can be defined.          Arguments:     - image:       image as torch tensor     Keyword Arguments:     - ax:       Can be defined if image shall be plotted as subplot (default: None)     - title:       Can be defined if title shall be printed as string     Returns:     - ax:       matplotlib axis # PyTorch tensors assume the color channel is the first dimension # but matplotlib assumes is the third dimension # Undo preprocessing # Image needs to be clipped between 0 and 1 or it looks like noise when displayed # Showing the image Description:     ============     Given a model, an image path and a dict to map from classes to flower names,      this function predicts the top k classes with names. Prediction can be process     on device          Arguments:     - image_path:       Path to image as string     Keyword Arguments:     - checkpoint_path:       Path to checkpoint containing model, criterion, optimizer and class_to_idx as str       (default: best_model.pth)     - cat_to_name_path:       Path to json containing cat to name dict as str (default: cat_to_name.json)     - k:       Number of classes     - device:       device for the model (default: 'cuda:0' if available, else 'cpu')     Returns:     - top_probabilities:       list of top k probabilities as float     - top_classes:       list of top k classes as int     - top_names:       list of top k names as str # Inferring top 5 probabilities and classes for given image # Retrieving list of top 5 probabilities, classes and class names # Retrieving names from top 5 classes Description:     ============     Shows a bar chart with the top five predictions and prints picture with it.          Arguments:     - image_path:       Path to image as string     Keyword Arguments:     - checkpoint_path:       Path to checkpoint containing model, criterion, optimizer and class_to_idx as str       (default: best_model.pth)     - cat_to_name_path:       Path to json containing cat to name dict as str (default: cat_to_name.json)     - device:       device for the model (default: 'cuda:0' if available, else 'cpu') # Calculating top 5 probabilities with class names # Setting of the chart for top 5 classes for image with probabilities # Showing bar chart of top 5 classes # Retrieving the image and plotting it with name # MAIN FUNCTION Description:     ============     Function runs when predict.py is being executed.     It reads data from an image, loads a checkpoint and calculates the according     k highest probabilities and categories giving the category names when a mapping     file is being passed. Can run on gpu when defined and available.          Arguments:     - path_to_image:       Path to the image to predict as str.     - checkpoint:       Path to checkpoint/model as str.     Keyword Arguments:     - category_names:       Path to mapping file as str.       (default: cat_to_name.json)     - k:       Number of k highest probabilities and categories to calculate        as int (default: 5)     - device:       device for the model (default: 'cpu' if available, else 'cpu')