import sys

sys.path.append('architectures/deeplab_3D/')

sys.path.append('architectures/unet_3D/')

sys.path.append('architectures/hrnet_3D/')

sys.path.append('architectures/experiment_nets_3D/')

sys.path.append('utils/')

import os

from os import walk

import torch

import torch.nn as nn

from torch.autograd import Variable

import torch.backends.cudnn as cudnn

import torch.nn.functional as F

import torch.optim as optim

import numpy as np

import scipy.misc

import os

from tqdm import *

import random

from random import randint

from docopt import docopt

import deeplab_resnet_3D

import unet_3D

import highresnet_3D

import exp_net_3D

import lossF

import PP

import augmentations as AUG

import evalF as EF

import evalFP as EFP

import evalMetrics as METRICS

docstr = """Write something here

Usage:

    evalpyt.py [options]

Options:

    -h, --help                  Print this message

    --visualize                 view outputs of each sketch

    --evalMethod=<int>          0 for evaluation of model by whole image, 1 for patches [default: 1]

    --patchPredSize=<int>       If evaluating model with patches, the size of the patch [default: 60]

    --evalMetric=<str>          'iou','dice',only iou supported right now [default: iou]

    --snapPrefix=<str>          Snapshot prefix. a_1000.pth, a_2000.pth, a is prefix [default: HR3Dadice_1_2017-07-16-18-32_iter]

    --singleEval                Evaluate a single model

    --postFix=<str>             Postfix [default: _200x200x100orig]

    --resultsDir=<str>          Path to save evaluation results and predictions to [default: eval_results/]

    --predictionsPath=<str>     predictions path [default: 1]

    --snapshotPath=<str>        Snapshot path [default: models/snapshots/]

    --mainFolderPath=<str>      Main folder path [default: ../Data/MS2017b/]

    --NoLabels=<int>            The number of different labels in training data [default: 2]

    --gpu0=<int>                GPU number [default: 0]

    --useGPU=<int>              Use GPU [default: 0]

    --testMode                  Enable test model (no evaluation, only predictions)

    --modelPath=<str>           Full model path to test if only 1 model (test mode or singleEval mode use this) [default: None]

    --iterRange=<str>           Range of num iters [default: 1-21]

    --iterStep=<int>            Step size of iters [default: 1]

    --testAugm                  Apply test time augmentations

    --extraPatch=<int>          Extra patch size [default: 5]

"""

args = docopt(docstr, version='v0.1')

print(args)

eval_method = int(args['--evalMethod'])

patch_pred_size = int(args['--patchPredSize'])

eval_metric = args['--evalMetric']

snap_prefix = args['--snapPrefix']

results_dir = args['--resultsDir']

snapshots_path = args['--snapshotPath']

postfix = args['--postFix']

main_folder_path = args['--mainFolderPath']

num_labels = int(args['--NoLabels'])

gpu0 = int(args['--gpu0'])

useGPU = int(args['--useGPU'])

test_mode = args['--testMode']

model_path = args['--modelPath']

iter_range = args['--iterRange']

iter_step = int(args['--iterStep'])

iter_low, iter_high = int(iter_range.split('-')[0]), int(iter_range.split('-')[1])

eval_list = main_folder_path + 'val' + postfix + '.txt'

test_augm = args['--testAugm']

single_eval = args['--singleEval']

extra_patch = int(args['--extraPatch'])

if single_eval or test_mode:

    models_path = model_path

else:

    models_path = snap_prefix

if test_mode:

    if not os.path.exists('temp_preds/'):

        os.makedirs('temp_preds/')

else:

    if not os.path.exists(results_dir):

        print('Creating directory at:' , results_dir)

        os.makedirs(results_dir)

    results_file = open(os.path.join(results_dir, os.path.splitext(os.path.basename(models_path))[0] + '_ep_' + str(extra_patch) + '.txt'), 'w+')

if num_labels == 2:

    onlyLesions = True

else:

    onlyLesions = False

num_labels2 = 209

def modelInit():

    isPriv = False

    f_name = models_path.split('/')[-1]

    #load model

    if 'EXP3D' in f_name:

        experiment = f_name.replace('EXP3D_', '').replace('.pth', '').split('_')

        experiment = '_'.join(experiment[0:3])

        dilation_arr, isPriv, withASPP = PP.getExperimentInfo(experiment)

        model = exp_net_3D.getExpNet(num_labels, dilation_arr, isPriv, NoLabels2 = num_labels2, withASPP = withASPP)

    elif 'HR3D' in f_name:

        model = highresnet_3D.getHRNet(num_labels)

    elif 'DL3D' in f_name:

        model = deeplab_resnet_3D.Res_Deeplab(num_labels)

    elif 'UNET3D' in  f_name:

        model = unet_3D.UNet3D(1, num_labels)

    else:

        print('No model available for this .pth')

        sys.exit()

    model.eval()

    if useGPU:

        model.cuda(gpu0)

    return model, isPriv

def evalModel(model):

    img_list = open(eval_list).readlines()

    if test_mode:

        if models_path == 'None':

            print('Insert model path if you are testing this model')

            sys.exit()

        model = loadSnapshot(model, models_path)

        for img_str in img_list:

            img_str = img_str.rstrip()

            img, gt, out, affine = EF.predict(os.path.join(main_folder_path, img_str),

                                                        model, num_labels, postfix, main_folder_path, eval_method, 

                                                        gpu0, useGPU, patch_size = patch_pred_size, test_augm = test_augm, extra_patch = extra_patch)

            #save prediction

            save_path = os.path.join('temp_preds', 'pred_' + img_str.split('/')[-3] + '_s' + str(gt.shape[0]) + '.nii.gz')

            PP.saveScan(out, affine, save_path)

    else:

        if single_eval:

            r = range(1)

        else:

            r = range(iter_low, iter_high, iter_step)

        for iter in r:

            counter = 0

            if single_eval:

                model = loadSnapshot(model, models_path)

            else:

                model = loadSnapshot(model, os.path.join(snapshots_path, models_path + '_' + str(iter*1000) + '.pth'))

            r_list_iou = []

            r_list_dice = []

            r_list_recall = []

            r_list_precision = []

            for img_str in img_list:

                img_str = img_str.rstrip()

                img, gt, out, _ = EF.predict(os.path.join(main_folder_path, img_str),

                                                                model, num_labels, postfix, main_folder_path,

                                                                eval_method, gpu0, useGPU,  patch_size = patch_pred_size, test_augm = test_augm, extra_patch = extra_patch)

                result_iou = METRICS.metricEval('iou', out, gt, num_labels)

                result_dice = METRICS.metricEval('dice', out, gt, num_labels)

                result_recall = METRICS.metricEval('recall', out, gt, num_labels)

                result_precision = METRICS.metricEval('precision', out, gt, num_labels)

                r_list_iou.append(result_iou)

                r_list_dice.append(result_dice)

                r_list_recall.append(result_recall)

                r_list_precision.append(result_precision)

                counter += 1

                print "Model Iter {:5d} Progress: {:4d}/{:4d} iou {:1.4f} dice {:1.4f} recall {:1.4f} precision {:1.4f}  \r".format(iter * 1000, counter, len(img_list), result_iou, result_dice, result_recall, result_precision),

                sys.stdout.flush()

            avg_iou = np.sum(np.asarray(r_list_iou))/len(r_list_iou)

            avg_dice = np.sum(np.asarray(r_list_dice))/len(r_list_dice)

            avg_recall = np.sum(np.asarray(r_list_recall))/len(r_list_recall)

            avg_precision = np.sum(np.asarray(r_list_precision))/len(r_list_precision)

            results_file.write('Iterations: {:5d} iou: {:1.4f} dice: {:1.4f} recall: {:1.4f} precision: {:1.4f} \n'.format(iter*1000, avg_iou, avg_dice, avg_recall, avg_precision))

            print('Done!')

        results_file.close()

def evalModelPriv(model):

    img_list = open(eval_list).readlines()

    if test_mode:

        if models_path == 'None':

            print('Insert model path if you are testing this model')

            sys.exit()

        model = loadSnapshot(model, models_path)

        for img_str in img_list:

            img_str = img_str.rstrip()

            img, gif, out1, gt, out2, affine = EFP.predict(os.path.join(main_folder_path, img_str), 

                                                        model, num_labels, num_labels2, postfix, main_folder_path, eval_method, 

                                                        gpu0, useGPU, patch_size = patch_pred_size, test_augm = test_augm, extra_patch = extra_patch)          

            #save prediction

            save_path = os.path.join('temp_preds', 'pred_' + img_str.split('/')[-3] + '_s' + str(gt.shape[0]) + '.nii.gz')

            PP.saveScan(out2, affine, save_path)

    else:

        if single_eval:

            r = range(1)

        else:

            r = range(iter_low, iter_high, iter_step)

        for iter in r:

            if single_eval:

                model = loadSnapshot(model, models_path)

            else:

                model = loadSnapshot(model, os.path.join(snapshots_path, models_path + '_' + str(iter*1000) + '.pth'))

            counter = 0

            r_list_iou_main = []

            r_list_dice_main = []

            r_list_recall_main = []

            r_list_precision_main = []

            r_list_iou_sec = []

            v = 0

            v_priv = 0

            for img_str in img_list:

                img_str = img_str.rstrip()

                img, gt1, out1, gt2, out2, _ = EFP.predict(os.path.join(main_folder_path, img_str),

                                                                model, num_labels, num_labels2, postfix, main_folder_path, 

                                                                eval_method, gpu0, useGPU,  patch_size = patch_pred_size, test_augm = test_augm, extra_patch = extra_patch)

                result_iou_main = METRICS.metricEval('iou', out2, gt2, num_labels)

                result_dice_main = METRICS.metricEval('dice', out2, gt2, num_labels)

                result_recall_main = METRICS.metricEval('recall', out2, gt2, num_labels)

                result_precision_main = METRICS.metricEval('precision', out2, gt2, num_labels)

                result_iou_sec = METRICS.metricEval('iou', out1, gt1, num_labels2)

                r_list_iou_main.append(result_iou_main)

                r_list_dice_main.append(result_dice_main)

                r_list_recall_main.append(result_recall_main)

                r_list_precision_main.append(result_precision_main)

                r_list_iou_sec.append(result_iou_sec)

                counter += 1

                print "Model Iter | {:5d} | Progress: | {:4d}/{:4d} | Last result {:1.4f}    \r".format(iter * 1000, counter, len(img_list), result_iou_main),

                sys.stdout.flush()

            avg_iou = np.sum(np.asarray(r_list_iou_main))/len(r_list_iou_main)

            avg_dice = np.sum(np.asarray(r_list_dice_main))/len(r_list_dice_main)

            avg_recall = np.sum(np.asarray(r_list_recall_main))/len(r_list_recall_main)

            avg_precision = np.sum(np.asarray(r_list_precision_main))/len(r_list_precision_main)

            avg_iou_sec =  np.sum(np.asarray(r_list_iou_sec))/len(r_list_iou_sec)

            results_file.write('Iterations: {:5d} iou: {:1.4f} dice: {:1.4f} recall: {:1.4f} precision: {:1.4f} iou_secondary: {:1.4f} \n'.format(iter*1000, avg_iou, avg_dice, avg_recall, avg_precision, avg_iou_sec))

        print('Done!')

        results_file.close()

def loadSnapshot(model, path):

    if useGPU:

        #loading on GPU when model was saved on GPU

        saved_state_dict = torch.load(path)

    else:

        #loading on CPU when model was saved on GPU

        saved_state_dict = torch.load(path, map_location=lambda storage, loc: storage)

    model.load_state_dict(saved_state_dict)

    return model

if __name__ == "__main__":

    model, with_priv = modelInit()

    if with_priv:

        evalModelPriv(model)

    else:

        evalModel(model)