#!/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', False, 'import resource\n'), (28, 'resource.setrlimit', 'resource.setrlimit', False, 'import resource\n'), (33, 'args.parse_args', 'parse_args', False, 'from args import parse_args\n'), (35, 'torch.device', 'torch.device', False, 'import torch\n'), (36, 'torch.cuda.manual_seed', 'torch.cuda.manual_seed', False, 'import torch\n'), (50, 'os.path.join', 'os.path.join', False, 'import os\n'), (51, 'os.path.join', 'os.path.join', False, 'import os\n'), (52, 'os.path.join', 'os.path.join', False, 'import os\n'), (55, 'os.makedirs', 'os.makedirs', False, 'import os\n'), (56, 'os.makedirs', 'os.makedirs', False, 'import os\n'), (57, 'os.makedirs', 'os.makedirs', False, 'import os\n'), (58, 'os.makedirs', 'os.makedirs', False, 'import os\n'), (59, 'os.makedirs', 'os.makedirs', False, 'import os\n'), (77, 'pickle.load', 'pickle.load', False, 'import pickle\n'), (78, 'os.path.join', 'os.path.join', False, 'import os\n'), (81, 'torch.utils.data.DataLoader', 'DataLoader', 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', False, 'import torch\n'), (88, 'matplotlib.pyplot.subplots', 'plt.subplots', True, 'import matplotlib.pyplot as plt\n'), (112, 'torch.no_grad', 'torch.no_grad', False, 'import torch\n'), (225, 'helperfunctions.stackall_Dict', 'stackall_Dict', False, 'from helperfunctions import mypause, stackall_Dict\n'), (226, 'numpy.stack', 'np.stack', True, 'import numpy as np\n'), (227, 'numpy.nanmean', 'np.nanmean', True, 'import numpy as np\n'), (240, 'pickle.dump', 'pickle.dump', False, 'import pickle\n'), (25, 'os.getcwd', 'os.getcwd', False, 'import os\n'), (64, 'os.path.join', 'os.path.join', False, 'import os\n'), (69, 'utils.get_nparams', 'get_nparams', False, 'from utils import get_nparams, get_predictions\n'), (113, 'tqdm.tqdm', 'tqdm.tqdm', False, 'import tqdm\n'), (136, 'utils.get_predictions', 'get_predictions', False, 'from utils import get_nparams, get_predictions\n'), (173, 'utils.unnormPts', 'unnormPts', False, 'from utils import getSeg_metrics, getPoint_metric, generateImageGrid, unnormPts\n'), (175, 'utils.unnormPts', 'unnormPts', False, 'from utils import getSeg_metrics, getPoint_metric, generateImageGrid, unnormPts\n'), (177, 'utils.unnormPts', 'unnormPts', False, 'from utils import getSeg_metrics, getPoint_metric, generateImageGrid, unnormPts\n'), (179, 'utils.unnormPts', 'unnormPts', False, 'from utils import getSeg_metrics, getPoint_metric, generateImageGrid, unnormPts\n'), (239, 'os.path.join', 'os.path.join', False, 'import os\n'), (228, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (229, 'numpy.nanmedian', 'np.nanmedian', True, 'import numpy as np\n'), (230, 'numpy.nanstd', 'np.nanstd', True, 'import numpy as np\n'), (231, 'numpy.nanmedian', 'np.nanmedian', True, 'import numpy as np\n'), (232, 'numpy.nanstd', 'np.nanstd', True, 'import numpy as np\n'), (233, 'numpy.nanmedian', 'np.nanmedian', True, 'import numpy as np\n'), (234, 'numpy.nanstd', 'np.nanstd', True, 'import numpy as np\n'), (235, 'numpy.nanmedian', 'np.nanmedian', True, 'import numpy as np\n'), (236, 'numpy.nanstd', 'np.nanstd', 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', True, 'import matplotlib.pyplot as plt\n'), (223, 'helperfunctions.mypause', 'mypause', False, 'from helperfunctions import mypause, stackall_Dict\n')]