import sys
import lasagne as nn
import numpy as np
import theano
import pathfinder
import utils
from configuration import config, set_configuration
from utils_plots import plot_slice_3d_3
import theano.tensor as T
import utils_lung
import blobs_detection
import logger
from collections import defaultdict
import glob
import data_transforms

theano.config.warn_float64 = 'raise'

if len(sys.argv) < 2:
    sys.exit("Usage: test_luna_scan.py <configuration_name>")

config_name = sys.argv[1]
set_configuration('configs_fpred_scan', config_name)

# predictions path
predictions_dir = utils.get_dir_path('model-predictions', pathfinder.METADATA_PATH)
outputs_path = predictions_dir + '/%s' % config_name
outputs_img_path = predictions_dir + '/%s_img' % config_name
utils.auto_make_dir(outputs_img_path)

blob_files = sorted(glob.glob(outputs_path + '/*.pkl'))

p_transform = {'patch_size': (64, 64, 64),
               'mm_patch_size': (64, 64, 64),
               'pixel_spacing': (1., 1., 1.)
               }

for p in blob_files:
    pid = utils_lung.extract_pid_filename(p, '.pkl')
    blobs = utils.load_pkl(p)
    blobs = np.asarray(sorted(blobs, key=lambda x: x[-1], reverse=True))

    img, pixel_spacing = utils_lung.read_dicom_scan(pathfinder.DATA_PATH + '/' + pid)
    print pid
    for blob in blobs[:10]:
        patch_center = blob[:3]
        p1 = blob[-1]
        print p1
        x, _ = data_transforms.transform_patch3d(data=img,
                                                 luna_annotations=None,
                                                 patch_center=patch_center,
                                                 p_transform=p_transform,
                                                 pixel_spacing=pixel_spacing,
                                                 luna_origin=None,
                                                 world_coord_system=False)

        plot_slice_3d_3(input=x, mask=x, prediction=x,
                        axis=0, pid='-'.join([str(pid), str(p1)]),
                        img_dir=outputs_img_path, idx=np.array(x[0, 0].shape) / 2)
        # print 'saved'