import sys

from time import time

import pickle

from os import makedirs

from os.path import join

import argparse

import torch

from torch import nn, sigmoid

from torch.utils.data import DataLoader

from model import UNet3D

from dataset import MRIDataset

from utils import Report, transfer_weights

"""

Run inference on test set.

This script also saves voxel-wise inference results and labels on numpy arrays

for future reuse (without running on gpu).

"""

argv = sys.argv[1:]

parser = argparse.ArgumentParser(

    formatter_class=argparse.RawDescriptionHelpFormatter,

    prog='PROG',)

parser.add_argument('--data_dir',

                    type=str,

                    required=True,

                    help="path to images (should have subdirs: test)")

parser.add_argument('--model_path', '-m', type=str,

                    required=True,

                    help='path to a model snapshot')

parser.add_argument('--size', type=int, nargs='+', default=[128])

parser.add_argument('--dump_name',

                    type=str,

                    default='',

                    help='name dump file to avoid overwrting files.')

args = parser.parse_args(argv)

net = UNet3D(1, 1, use_bias=True, inplanes=16)

transfer_weights(net, args.model_path)

net.cuda()

net.train(False)

torch.no_grad()

test_dir = join(args.data_dir, 'test')

batch_size = 1

def inference(target_dir):

    volume_size = args.size*3 if len(args.size) == 1 else args.size

    dataloader = DataLoader(MRIDataset(target_dir,

                                       volume_size,

                                       sampling_mode='center',

                                       deterministic=True),

                            batch_size=1,

                            num_workers=4)

    input_paths = dataloader.dataset.inputs

    label_paths = dataloader.dataset.labels

    reporter = Report()

    preds_list = list()

    labels_list = list()

    sum_hd = 0

    sum_sd = 0

    num_voxels_hd = list()

    for i, (inputs, labels) in enumerate(dataloader):

        inputs = inputs.cuda()

        preds = sigmoid(net(inputs.detach()).detach())

        full_labels = dataloader.dataset._load_full_label(label_paths[i])

        full_preds = dataloader.dataset._project_full_label(input_paths[i],

                                                            preds.cpu())

        preds = full_preds

        labels = full_labels

        reporter.feed(preds, labels)

        temp_reporter = Report()

        temp_reporter.feed(preds, labels)

        sum_hd += temp_reporter.hard_dice()

        sum_sd += temp_reporter.soft_dice()

        num_voxels_hd.append(

            (labels.view(-1).sum().item(), temp_reporter.hard_dice()))

        del temp_reporter

        preds_list.append(preds.cpu())

        labels_list.append(labels.cpu())

        del inputs, labels, preds

    print("Micro Averaged Dice {}, {}".format(sum_hd / len(dataloader),

                                              sum_sd / len(dataloader)))

    if len(args.dump_name):

        # dump preds for visualization

        pickle.dump([input_paths, preds_list, labels_list],

                    open('preds_dump_{}.pickle'.format(args.dump_name), 'wb'))

    print(reporter)

    print(reporter.stats())

    preds = torch.stack(preds_list).view(-1).numpy()

    labels = torch.stack(labels_list).view(-1).numpy().astype(int)

    print(num_voxels_hd)

    return preds, labels

preds, labels = inference(test_dir)