import configs_fpred_patch.luna_c3 as patch_class_config

p_transform = patch_class_config.p_transform

data_prep_function = patch_class_config.partial(patch_class_config.data_prep_function,
                                                p_transform_augment=None,
                                                p_transform=p_transform,
                                                world_coord_system=False,
                                                luna_origin=None)

rng = patch_class_config.rng

# candidates after segmentations path
predictions_dir = utils.get_dir_path('model-predictions', pathfinder.METADATA_PATH)
segmentation_outputs_path = predictions_dir + '/%s' % seg_config_name
id2candidates_path = utils_lung.get_candidates_paths(segmentation_outputs_path)

# filter our those, who are already generated
predictions_dir = utils.get_dir_path('model-predictions', pathfinder.METADATA_PATH)
outputs_path = predictions_dir + '/dsb_c3_s2_p8a1_ls_elias'  # TODO write it here correctly
exclude_pids = []
if os.path.isdir(outputs_path):
    exclude_pids = os.listdir(outputs_path)
    exclude_pids = [utils_lung.extract_pid_filename(p) for p in exclude_pids]
    for pid in exclude_pids:
        id2candidates_path.pop(pid, None)

data_iterator = data_iterators.CandidatesDSBDataGenerator(data_path=pathfinder.DATA_PATH,
                                                          transform_params=p_transform,
                                                          data_prep_fun=data_prep_function,
                                                          id2candidates_path=id2candidates_path)

from functools import partial
import lasagne.layers.dnn as dnn
import theano.tensor as T
import utils
import utils_lung
import os

# TODO: import correct config here
candidates_config = 'dsb_c3_s2_p8a1_ls_elias'

restart_from_save = None
rng = np.random.RandomState(42)

predictions_dir = utils.get_dir_path('model-predictions', pathfinder.METADATA_PATH)
candidates_path = predictions_dir + '/%s' % candidates_config
id2candidates_path = utils_lung.get_candidates_paths(candidates_path)

# transformations
p_transform = {'patch_size': (48, 48, 48),
               'mm_patch_size': (48, 48, 48),
               'pixel_spacing': (1., 1., 1.)
               }
p_transform_augment = {
    'translation_range_z': [-5, 5],
    'translation_range_y': [-5, 5],
    'translation_range_x': [-5, 5],
    'rotation_range_z': [-10, 10],
    'rotation_range_y': [-10, 10],
    'rotation_range_x': [-10, 10]
}
n_candidates_per_patient = 8

import utils
from configuration import set_configuration
import utils_lung

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
pid2candidates_path = utils_lung.get_candidates_paths(outputs_path)
pid2candidates = {}
for k, v in pid2candidates_path.iteritems():
    pid2candidates[k] = utils.load_pkl(v)

pid2annotations = utils_lung.read_luna_annotations(pathfinder.LUNA_LABELS_PATH)

n_top = 1
tp_top_n = 0
fp_top_n = 0
tp = 0
n_pos = 0
idx1 = []
for k in pid2candidates.iterkeys():
    print '----------------------------------------'
    print k