import os
import numpy as np
import torch
import torch.nn as nn
import torchvision
from torchvision import datasets, models, transforms
from load_model import initialize_model
from PIL import Image
import json
import argparse
import os
import sys
from datetime import date
import onnxruntime

def preprocess(image_path, input_size):
    data_transforms = transforms.Compose([
        transforms.Resize(input_size),
        transforms.ToTensor(),
        transforms.Normalize([0, 0, 0], [1/255.0, 1/255.0, 1/255.0]),
        transforms.Normalize([0.5*256, 0.5*256, 0.5*256], [256.0, 256.0, 256.0])
        ])
    with torch.no_grad():
        img_data_pytorch = data_transforms(Image.open(image_path))
        img_data_pytorch = img_data_pytorch.unsqueeze(0)
    return img_data_pytorch.numpy()

def postprocess(pre_output):
    score = softmax(pre_output)
    labels = list(range(len(pre_output)))
    score_labels = list(zip(score, labels))
    score_labels.sort(reverse=True)
    score, labels = list(zip(*score_labels))
    
    return score, labels
    

def onnx_runner(image_path, model_path, class_id):
    sess = onnxruntime.InferenceSession(model_path)
    onnx_img_size_h = sess.get_inputs()[0].shape[2]
    onnx_img_size_w = sess.get_inputs()[0].shape[3]
    input_name = sess.get_inputs()[0].name
    input_size = (onnx_img_size_h, onnx_img_size_w)
    np_images = preprocess(image_path, input_size)
    np_images = np_images.astype(np.float32)
    pred_onnx = sess.run(None, {input_name: np_images })[0][0]

    score, labels = postprocess(pred_onnx)
    
    header = 'Label      Probability'
    itn_line = '{:10} {:8.3f} '   
    print(header)
    for i in range(len(score)):
        #print(itn_line.format( class_id[str(labels[i])], score[i]) )
        print(itn_line.format( str(labels[i]), score[i]) )

    return score, labels
     
def softmax(A):
    e = np.exp(A)
    return e / np.sum(e, axis=1, keepdims=True)    

def inference(backbone, image_path, class_id, device, model_def_path, pretrained_path, topk = None):
    
    num_classes = len(class_id)
    model_structure, input_size = initialize_model(backbone, num_classes, False, model_def_path)
    
    model_structure.load_state_dict(torch.load(pretrained_path))
    model = model_structure.eval()
    model = model.to(device)
    
    data_transforms = transforms.Compose([
        transforms.Resize(input_size),
        transforms.ToTensor(),
        transforms.Normalize([0, 0, 0], [1/255.0, 1/255.0, 1/255.0]),
        transforms.Normalize([0.5*256, 0.5*256, 0.5*256], [256.0, 256.0, 256.0])
        ])

    img_data_pytorch = data_transforms(Image.open(image_path))
    img_data_pytorch = img_data_pytorch.to(device)
    with torch.no_grad():
        if topk == None or topk > num_classes:
            topk = num_classes
        outputs = model(img_data_pytorch[None, ...]).topk(topk)
        scores = outputs[0].cpu().numpy()[0]
        probs = softmax(scores)
        preds = outputs[1].cpu().numpy()[0]
        
    header = 'Label      Probability'
    itn_line = '{:10} {:8.3f} '   
    print(header)
    for i in range(len(preds)):
        print(itn_line.format( class_id[str(preds[i])], probs[i]) )
        
    return probs, preds

def softmax(A):

    e = np.exp(A)
    return e / np.sum(e, keepdims=True)

def parse_args(args):
    """
    Parse the arguments.
    """
    today = str(date.today())

    parser = argparse.ArgumentParser(description='Simple training script for training a image classification network.')
    parser.add_argument('--img-path', type=str, help='Path to the image.')
    parser.add_argument('--backbone', help='Backbone model.', default='resnet18', type=str)
    parser.add_argument('--class_id_path', help='Path to the class id mapping file.', default='./eval_utils/class_id.json')
    parser.add_argument('--gpu', help='Id of the GPU to use (as reported by nvidia-smi). (-1 for cpu)',type=int,default=-1)
    parser.add_argument('--model-def-path', type=str, help='Path to pretrained model definition', default=None )
    parser.add_argument('--snapshot', help='Path to the pretrained models.')
    parser.add_argument('--save-path', help='Path to the classification result.', default='inference_result.json')
    parser.add_argument('--onnx', help='inference onnx model',action='store_true')
    
    print(vars(parser.parse_args(args)))
    return check_args(parser.parse_args(args))

def check_args(parsed_args):
    """ Function to check for inherent contradictions within parsed arguments.
    Args
        parsed_args: parser.parse_args()
    Returns
        parsed_args
    """
    if parsed_args.gpu >= 0 and torch.cuda.is_available() == False:
        raise ValueError("No gpu is available")
    return parsed_args                                     
    
def main(args=None):
    # parse arguments
    if args is None:
        args = sys.argv[1:]

    args = parse_args(args)
    device = "cuda:"+str(args.gpu) if args.gpu >= 0 else "cpu"
    with open(args.class_id_path,'r') as fp:
        class_id = json.load(fp)

    # Inference 
    if args.onnx:
        probs, preds = onnx_runner(args.img_path, args.snapshot, class_id)
    else:
        probs, preds = inference(args.backbone, args.img_path, class_id, device, args.model_def_path, args.snapshot)
    res = {}
    res['img_path'] = os.path.abspath(args.img_path)
    res['0_0'] = []

    for i in range(len(probs)):
        res['0_0'].append([ float(probs[i]), int(preds[i]) ])
    
    with open(args.save_path, 'w') as fp:
        json.dump(res, fp)


if __name__ == '__main__':
    main()