#Import libraries
import torch
import torchvision.models as models
import json

#Import User Defined libraries
from neural_network_model import initialize_existing_models, build_custom_models, set_parameter_requires_grad
from utilities import process_image, get_input_args_predict

def predict(image_path, model, topk=5):
    
    
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model.to(device)
    model.eval()
    
    tensor_img = torch.FloatTensor(process_image(image_path))
    tensor_img = tensor_img.unsqueeze(0)
    tensor_img = tensor_img.to(device)
    log_ps = model(tensor_img)
    result = log_ps.topk(topk)
    if torch.cuda.is_available(): #gpu Move it from gpu to cpu for numpy
        ps = torch.exp(result[0].data).cpu().numpy()[0] 
        idxs = result[1].data.cpu().numpy()[0]
    else: #cpu Keep it on cpu for nump
        ps = torch.exp(result[0].data).numpy()[0]
        idxs = result[1].data.numpy()[0]

    return (ps, idxs)

def process_input(image_path):
    #0. Get user inputs
    #in_arg = vars(get_input_args_predict())
    #print("User arguments/hyperparameters or default used are as below")
    #print(in_arg)
    in_arg = {}
    in_arg['gpu'] = 'gpu'
    in_arg['save_dir'] = 'checkpoint-densenet121.pth'
    in_arg['path'] = image_path
    in_arg['top_k'] = 5
    
    #1. Get device for prediction and Load model from checkpoint along with some other information
    if in_arg['gpu'] == 'gpu' and torch.cuda.is_available():
        device = torch.device("cuda")
        checkpoint = torch.load(in_arg['save_dir'])
    else:
        device = "cpu"
        checkpoint = torch.load(in_arg['save_dir'], map_location = device)
    #print(f"Using {device} device for predicting/inference")
    
    if checkpoint['arch_type'] == 'existing':
        model_ft, input_size = initialize_existing_models(checkpoint['arch'], checkpoint['arch_type'], len(checkpoint['class_to_idx']),
                                                          checkpoint['feature_extract'], checkpoint['hidden_units'], use_pretrained=False)
    elif checkpoint['arch_type'] == 'custom':
        model_ft = build_custom_models(checkpoint['arch'], checkpoint['arch_type'], len(checkpoint['class_to_idx']), checkpoint['feature_extract'], 
                                       checkpoint['hidden_units'], use_pretrained=True)
    else:
        #print("Nothing to predict")
        exit()
        
       
    model_ft.class_to_idx = checkpoint['class_to_idx']
    model_ft.gpu_or_cpu = checkpoint['gpu_or_cpu']
    model_ft.load_state_dict(checkpoint['state_dict'])
    model_ft.to(device)
    
    #Predict
    # Get the prediction by passing image and other user preferences through the model
    probs, idxs  = predict(image_path = in_arg['path'], model = model_ft, topk = in_arg['top_k'])
    
    # Swap class to index mapping with index to class mapping and then map the classes to flower category labels using the json file
    idx_to_class = {v: k for k, v in model_ft.class_to_idx.items()}
    with open('cat_to_name.json','r') as f:
        cat_to_name = json.load(f)
    names = list(map(lambda x: cat_to_name[f"{idx_to_class[x]}"],idxs))
    
    # Display final prediction and Top k most probable flower categories                               
    #print("-"*60)
    #print("                    PREDICTION")
    #print("-"*60)
    #print("Image provided        : {}" .format(in_arg['path']))
    #print("Predicted Flower Name : {} (Class {} and Index {})" .format(names[0].upper(), idx_to_class[idxs[0]], idxs[0] ))
    #print("Model used            : {}" .format(checkpoint['arch']))
    #print(f"The top {in_arg['top_k']} probabilities of the flower names")
    #for name, prob in zip(names, probs):
    #    length = 30 - len(name)
    #    print(f"{name.title()}{' '*length}{round(prob*100,2)}%")