#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import sys
import tqdm
import torch
import pickle
import resource
import numpy as np
import matplotlib.pyplot as plt

from args import parse_args
from modelSummary import model_dict
from pytorchtools import load_from_file
from torch.utils.data import DataLoader

from helperfunctions import mypause, stackall_Dict

from loss import get_seg2ptLoss

from utils import get_nparams, get_predictions
from utils import getSeg_metrics, getPoint_metric, generateImageGrid, unnormPts

sys.path.append(os.path.abspath(os.path.join(os.getcwd(), os.pardir)))

rlimit = resource.getrlimit(resource.RLIMIT_NOFILE)
resource.setrlimit(resource.RLIMIT_NOFILE, (2048*10, rlimit[1]))

#%%
if __name__ == '__main__':

    args = parse_args()

    device=torch.device("cuda")
    torch.cuda.manual_seed(12)
    if torch.cuda.device_count() > 1:
        print('Moving to a multiGPU setup.')
        args.useMultiGPU = True
    else:
        args.useMultiGPU = False
    torch.backends.cudnn.deterministic=False

    if args.model not in model_dict:
        print("Model not found.")
        print("valid models are: {}".format(list(model_dict.keys())))
        exit(1)

    LOGDIR = os.path.join(os.getcwd(), 'logs', args.model, args.expname)
    path2model = os.path.join(LOGDIR, 'weights')
    path2checkpoint = os.path.join(LOGDIR, 'checkpoints')
    path2writer = os.path.join(LOGDIR, 'TB.lock')
    path2op = os.path.join(os.getcwd(), 'op', str(args.curObj))

    os.makedirs(LOGDIR, exist_ok=True)
    os.makedirs(path2model, exist_ok=True)
    os.makedirs(path2checkpoint, exist_ok=True)
    os.makedirs(path2writer, exist_ok=True)
    os.makedirs(path2op, exist_ok=True)

    model = model_dict[args.model]

    netDict = load_from_file([args.loadfile,
                              os.path.join(path2checkpoint, 'checkpoint.pt')])
    startEp = netDict['epoch'] if 'epoch' in netDict.keys() else 0
    if 'state_dict' in netDict.keys():
        model.load_state_dict(netDict['state_dict'])

    print('Parameters: {}'.format(get_nparams(model)))
    model = model if not args.useMultiGPU else torch.nn.DataParallel(model)
    model = model.to(device).to(args.prec)

    f = open(os.path.join('curObjects',
                          'baseline',
                          'cond_'+str(args.curObj)+'.pkl'), 'rb')

    _, _, testObj = pickle.load(f)
    testObj.path2data = os.path.join(args.path2data, 'Datasets', 'All')
    testObj.augFlag = False

    testloader = DataLoader(testObj,
                            batch_size=args.batchsize,
                            shuffle=False,
                            num_workers=args.workers,
                            drop_last=False)

    if args.disp:
        fig, axs = plt.subplots(nrows=1, ncols=1)
    #%%
    accLoss = 0.0
    imCounter = 0

    ious = []

    dists_pupil_latent = []
    dists_pupil_seg = []

    dists_iris_latent = []
    dists_iris_seg = []

    model.eval()

    opDict = {'id':[], 'archNum': [], 'archName': [], 'code': [],
              'scores':{'iou':[], 'lat_dst':[], 'seg_dst':[]},
              'pred':{'pup_latent_c':[],
                      'pup_seg_c':[],
                      'iri_latent_c':[],
                      'iri_seg_c':[],
                      'mask':[]},
              'gt':{'pup_c':[], 'mask':[]}}

    with torch.no_grad():
        for bt, batchdata in enumerate(tqdm.tqdm(testloader)):
            img, labels, spatialWeights, distMap, pupil_center, iris_center, elNorm, cond, imInfo = batchdata
            out_tup = model(img.to(device).to(args.prec),
                            labels.to(device).long(),
                            pupil_center.to(device).to(args.prec),
                            elNorm.to(device).to(args.prec),
                            spatialWeights.to(device).to(args.prec),
                            distMap.to(device).to(args.prec),
                            cond.to(device).to(args.prec),
                            imInfo[:, 2].to(device).to(torch.long),
                            0.5)

            output, elOut, latent, loss = out_tup

            latent_pupil_center = elOut[:, 0:2].detach().cpu().numpy()
            latent_iris_center  = elOut[:, 5:7].detach().cpu().numpy()

            _, seg_pupil_center = get_seg2ptLoss(output[:, 2, ...].cpu(), pupil_center, temperature=4)
            _, seg_iris_center  = get_seg2ptLoss(-output[:, 0, ...].cpu(), iris_center, temperature=4)

            loss = loss if args.useMultiGPU else loss.mean()

            accLoss += loss.detach().cpu().item()
            predict = get_predictions(output)

            iou, iou_bySample = getSeg_metrics(labels.numpy(),
                                               predict.numpy(),
                                               cond[:, 1].numpy())[1:]

            latent_pupil_dist, latent_pupil_dist_bySample = getPoint_metric(pupil_center.numpy(),
                                                                            latent_pupil_center,
                                                                            cond[:,0].numpy(),
                                                                            img.shape[2:],
                                                                            True) # Unnormalizes the points

            seg_pupil_dist, seg_pupil_dist_bySample = getPoint_metric(pupil_center.numpy(),
                                                                      seg_pupil_center,
                                                                      cond[:,1].numpy(),
                                                                      img.shape[2:],
                                                                      True) # Unnormalizes the points

            latent_iris_dist, latent_iris_dist_bySample = getPoint_metric(iris_center.numpy(),
                                                                          latent_iris_center,
                                                                          cond[:,1].numpy(),
                                                                          img.shape[2:],
                                                                          True) # Unnormalizes the points

            seg_iris_dist, seg_iris_dist_bySample = getPoint_metric(iris_center.numpy(),
                                                                    seg_iris_center,
                                                                    cond[:,1].numpy(),
                                                                    img.shape[2:],
                                                                    True) # Unnormalizes the points

            dists_pupil_latent.append(latent_pupil_dist)
            dists_iris_latent.append(latent_iris_dist)
            dists_pupil_seg.append(seg_pupil_dist)
            dists_iris_seg.append(seg_iris_dist)

            ious.append(iou)

            pup_latent_c = unnormPts(latent_pupil_center,
                                     img.shape[2:])
            pup_seg_c = unnormPts(seg_pupil_center,
                                  img.shape[2:])
            iri_latent_c = unnormPts(latent_iris_center,
                                     img.shape[2:])
            iri_seg_c = unnormPts(seg_iris_center,
                                  img.shape[2:])

            dispI = generateImageGrid(img.numpy().squeeze(),
                                      predict.numpy(),
                                      elOut.detach().cpu().numpy().reshape(-1, 2, 5),
                                      pup_seg_c,
                                      cond.numpy(),
                                      override=True,
                                      heatmaps=False)

            for i in range(0, img.shape[0]):
                archNum = testObj.imList[imCounter, 1]
                opDict['id'].append(testObj.imList[imCounter, 0])
                opDict['code'].append(latent[i,...].detach().cpu().numpy())

                opDict['archNum'].append(archNum)
                opDict['archName'].append(testObj.arch[archNum])

                opDict['pred']['pup_latent_c'].append(pup_latent_c[i, :])
                opDict['pred']['pup_seg_c'].append(pup_seg_c[i, :])
                opDict['pred']['iri_latent_c'].append(iri_latent_c[i, :])
                opDict['pred']['iri_seg_c'].append(iri_seg_c[i, :])

                if args.test_save_op_masks:
                    opDict['pred']['mask'].append(predict[i,...].numpy().astype(np.uint8))

                opDict['scores']['iou'].append(iou_bySample[i, ...])
                opDict['scores']['lat_dst'].append(latent_pupil_dist_bySample[i, ...])
                opDict['scores']['seg_dst'].append(seg_pupil_dist_bySample[i, ...])

                opDict['gt']['pup_c'].append(pupil_center[i,...].numpy())

                if args.test_save_op_masks:
                    opDict['gt']['mask'].append(labels[i,...].numpy().astype(np.uint8))

                imCounter+=1

            if args.disp:
                if bt == 0:
                    h_im = plt.imshow(dispI.permute(1, 2, 0))
                    plt.pause(0.01)
                else:
                    h_im.set_data(dispI.permute(1, 2, 0))
                    mypause(0.01)

        opDict = stackall_Dict(opDict)
        ious = np.stack(ious, axis=0)
        ious = np.nanmean(ious, axis=0)
        print('mIoU: {}. IoUs: {}'.format(np.mean(ious), ious))
        print('Latent space PUPIL dist. Med: {}, STD: {}'.format(np.nanmedian(dists_pupil_latent),
                                                            np.nanstd(dists_pupil_latent)))
        print('Segmentation PUPIL dist. Med: {}, STD: {}'.format(np.nanmedian(dists_pupil_seg),
                                                            np.nanstd(dists_pupil_seg)))
        print('Latent space IRIS dist. Med: {}, STD: {}'.format(np.nanmedian(dists_iris_latent),
                                                           np.nanstd(dists_iris_latent)))
        print('Segmentation IRIS dist. Med: {}, STD: {}'.format(np.nanmedian(dists_iris_seg),
                                                           np.nanstd(dists_iris_seg)))

        print('--- Saving output directory ---')
        f = open(os.path.join(path2op, 'opDict.pkl'), 'wb')
        pickle.dump(opDict, f)
        f.close()

[(27, 'resource.getrlimit', 'resource.getrlimit', (['resource.RLIMIT_NOFILE'], {}), False, 'import resource\n'), (28, 'resource.setrlimit', 'resource.setrlimit', (['resource.RLIMIT_NOFILE', '(2048 * 10, rlimit[1])'], {}), False, 'import resource\n'), (33, 'args.parse_args', 'parse_args', ([], {}), False, 'from args import parse_args\n'), (35, 'torch.device', 'torch.device', (['"""cuda"""'], {}), False, 'import torch\n'), (36, 'torch.cuda.manual_seed', 'torch.cuda.manual_seed', (['(12)'], {}), False, 'import torch\n'), (50, 'os.path.join', 'os.path.join', (['LOGDIR', '"""weights"""'], {}), False, 'import os\n'), (51, 'os.path.join', 'os.path.join', (['LOGDIR', '"""checkpoints"""'], {}), False, 'import os\n'), (52, 'os.path.join', 'os.path.join', (['LOGDIR', '"""TB.lock"""'], {}), False, 'import os\n'), (55, 'os.makedirs', 'os.makedirs', (['LOGDIR'], {'exist_ok': '(True)'}), False, 'import os\n'), (56, 'os.makedirs', 'os.makedirs', (['path2model'], {'exist_ok': '(True)'}), False, 'import os\n'), (57, 'os.makedirs', 'os.makedirs', (['path2checkpoint'], {'exist_ok': '(True)'}), False, 'import os\n'), (58, 'os.makedirs', 'os.makedirs', (['path2writer'], {'exist_ok': '(True)'}), False, 'import os\n'), (59, 'os.makedirs', 'os.makedirs', (['path2op'], {'exist_ok': '(True)'}), False, 'import os\n'), (77, 'pickle.load', 'pickle.load', (['f'], {}), False, 'import pickle\n'), (78, 'os.path.join', 'os.path.join', (['args.path2data', '"""Datasets"""', '"""All"""'], {}), False, 'import os\n'), (81, 'torch.utils.data.DataLoader', 'DataLoader', (['testObj'], {'batch_size': 'args.batchsize', 'shuffle': '(False)', 'num_workers': 'args.workers', 'drop_last': '(False)'}), False, 'from torch.utils.data import DataLoader\n'), (37, 'torch.cuda.device_count', 'torch.cuda.device_count', ([], {}), False, 'import torch\n'), (49, 'os.getcwd', 'os.getcwd', ([], {}), False, 'import os\n'), (53, 'os.getcwd', 'os.getcwd', ([], {}), False, 'import os\n'), (70, 'torch.nn.DataParallel', 'torch.nn.DataParallel', (['model'], {}), False, 'import torch\n'), (88, 'matplotlib.pyplot.subplots', 'plt.subplots', ([], {'nrows': '(1)', 'ncols': '(1)'}), True, 'import matplotlib.pyplot as plt\n'), (112, 'torch.no_grad', 'torch.no_grad', ([], {}), False, 'import torch\n'), (225, 'helperfunctions.stackall_Dict', 'stackall_Dict', (['opDict'], {}), False, 'from helperfunctions import mypause, stackall_Dict\n'), (226, 'numpy.stack', 'np.stack', (['ious'], {'axis': '(0)'}), True, 'import numpy as np\n'), (227, 'numpy.nanmean', 'np.nanmean', (['ious'], {'axis': '(0)'}), True, 'import numpy as np\n'), (240, 'pickle.dump', 'pickle.dump', (['opDict', 'f'], {}), False, 'import pickle\n'), (25, 'os.getcwd', 'os.getcwd', ([], {}), False, 'import os\n'), (64, 'os.path.join', 'os.path.join', (['path2checkpoint', '"""checkpoint.pt"""'], {}), False, 'import os\n'), (69, 'utils.get_nparams', 'get_nparams', (['model'], {}), False, 'from utils import get_nparams, get_predictions\n'), (113, 'tqdm.tqdm', 'tqdm.tqdm', (['testloader'], {}), False, 'import tqdm\n'), (136, 'utils.get_predictions', 'get_predictions', (['output'], {}), False, 'from utils import get_nparams, get_predictions\n'), (173, 'utils.unnormPts', 'unnormPts', (['latent_pupil_center', 'img.shape[2:]'], {}), False, 'from utils import getSeg_metrics, getPoint_metric, generateImageGrid, unnormPts\n'), (175, 'utils.unnormPts', 'unnormPts', (['seg_pupil_center', 'img.shape[2:]'], {}), False, 'from utils import getSeg_metrics, getPoint_metric, generateImageGrid, unnormPts\n'), (177, 'utils.unnormPts', 'unnormPts', (['latent_iris_center', 'img.shape[2:]'], {}), False, 'from utils import getSeg_metrics, getPoint_metric, generateImageGrid, unnormPts\n'), (179, 'utils.unnormPts', 'unnormPts', (['seg_iris_center', 'img.shape[2:]'], {}), False, 'from utils import getSeg_metrics, getPoint_metric, generateImageGrid, unnormPts\n'), (239, 'os.path.join', 'os.path.join', (['path2op', '"""opDict.pkl"""'], {}), False, 'import os\n'), (228, 'numpy.mean', 'np.mean', (['ious'], {}), True, 'import numpy as np\n'), (229, 'numpy.nanmedian', 'np.nanmedian', (['dists_pupil_latent'], {}), True, 'import numpy as np\n'), (230, 'numpy.nanstd', 'np.nanstd', (['dists_pupil_latent'], {}), True, 'import numpy as np\n'), (231, 'numpy.nanmedian', 'np.nanmedian', (['dists_pupil_seg'], {}), True, 'import numpy as np\n'), (232, 'numpy.nanstd', 'np.nanstd', (['dists_pupil_seg'], {}), True, 'import numpy as np\n'), (233, 'numpy.nanmedian', 'np.nanmedian', (['dists_iris_latent'], {}), True, 'import numpy as np\n'), (234, 'numpy.nanstd', 'np.nanstd', (['dists_iris_latent'], {}), True, 'import numpy as np\n'), (235, 'numpy.nanmedian', 'np.nanmedian', (['dists_iris_seg'], {}), True, 'import numpy as np\n'), (236, 'numpy.nanstd', 'np.nanstd', (['dists_iris_seg'], {}), True, 'import numpy as np\n'), (46, 'modelSummary.model_dict.keys', 'model_dict.keys', ([], {}), False, 'from modelSummary import model_dict\n'), (220, 'matplotlib.pyplot.pause', 'plt.pause', (['(0.01)'], {}), True, 'import matplotlib.pyplot as plt\n'), (223, 'helperfunctions.mypause', 'mypause', (['(0.01)'], {}), False, 'from helperfunctions import mypause, stackall_Dict\n')]