import numpy as np

import utils_lung

import pathfinder

import utils

# 6% to 28% for nodules 5 to 10 mm,

prob5 = (0.01+0.06)/2.

slope10 = (0.28-prob5) / (10.-5.)

offset10 = prob5 - slope10*5.

slope20 = (0.64-0.28) / (20.-10.)

offset20 = 0.28 - slope20*10.

# and 64% to 82% for nodules >20 mm in diameter

slope25 = (0.82-0.64) / (25.-20.)

offset25 = 0.64 - slope25*20.

slope30 = (0.93-0.82) / (30.-25.)

offset30 = 0.82 - slope30*25.

# For nodules more than 3 cm in diameter, 93% to 97% are malignant

slope40 = (0.97-0.93) / (40.-30.)

offset40 = 0.93 - slope40*30.

def diameter_to_prob(diam):

    # The prevalence of malignancy is 0% to 1% for nodules <5 mm,

    if diam < 5:

        p = prob5*diam/5.

    elif diam < 10:

        p = slope10*diam+offset10

    elif diam < 20:

        p = slope20*diam+offset20

    elif diam < 25:

        p = slope25*diam+offset25

    elif diam < 30:

        p = slope30*diam+offset30

    else:

        p = slope40 * diam + offset40

    return np.clip(p ,0.,1.)

class LunaDataGenerator(object):

    def __init__(self, data_path, transform_params, data_prep_fun, rng,

                 random, infinite, patient_ids=None, **kwargs):

        self.patient_ids = patient_ids

        if patient_ids:

            self.patient_paths = [data_path + '/' + p + '.mhd' for p in patient_ids]

        else:

            patient_paths = utils_lung.get_patient_data_paths(data_path)

            self.patient_paths = [p for p in patient_paths if '.mhd' in p]

        self.id2annotations = utils_lung.read_luna_annotations(pathfinder.LUNA_LABELS_PATH)

        self.nsamples = len(self.patient_paths)

        self.data_path = data_path

        self.rng = rng

        self.random = random

        self.infinite = infinite

        self.data_prep_fun = data_prep_fun

        self.transform_params = transform_params

    def generate(self):

        while True:

            rand_idxs = np.arange(self.nsamples)

            if self.random:

                self.rng.shuffle(rand_idxs)

            for pos in xrange(0, len(rand_idxs)):

                idx = rand_idxs[pos]

                patient_path = self.patient_paths[idx]

                pid = utils_lung.extract_pid_filename(patient_path)

                img, origin, pixel_spacing = utils_lung.read_mhd(patient_path)

                x, y, annotations, tf_matrix = self.data_prep_fun(data=img,

                                                                  pixel_spacing=pixel_spacing,

                                                                  luna_annotations=

                                                                  self.id2annotations[pid],

                                                                  luna_origin=origin)

                x = np.float32(x)[None, None, :, :, :]

                y = np.float32(y)[None, None, :, :, :]

                yield x, y, None, annotations, tf_matrix, pid

            if not self.infinite:

                break

class LunaSimpleDataGenerator(object):

    def __init__(self, data_path, patient_ids=None, **kwargs):

        self.patient_ids = patient_ids

        self.data_path = data_path

        self.file_extension = '.pkl' if 'pkl' in data_path else '.mhd'

        if patient_ids:

            self.patient_paths = [data_path + '/' + p + self.file_extension for p in patient_ids]

        else:

            patient_paths = utils_lung.get_patient_data_paths(data_path)

            self.patient_paths = [p for p in patient_paths if self.file_extension in p]

        self.nsamples = len(self.patient_paths)

        print self.data_path

    def generate(self):

        for patient_path in self.patient_paths:

            pid = utils_lung.extract_pid_filename(patient_path)

            img, origin, pixel_spacing = utils_lung.read_pkl(patient_path) \

                if self.file_extension == '.pkl' else utils_lung.read_mhd(patient_path)

            x = np.float32(img)

            yield x, pid

class LunaScanPositiveDataGenerator(LunaDataGenerator):

    def __init__(self, data_path, transform_params, data_prep_fun, rng,

                 random, infinite, patient_ids=None, **kwargs):

        super(LunaScanPositiveDataGenerator, self).__init__(data_path, transform_params, data_prep_fun, rng,

                                                            random, infinite, patient_ids, **kwargs)

        patient_ids_all = [utils_lung.extract_pid_filename(p) for p in self.patient_paths]

        patient_ids_pos = [pid for pid in patient_ids_all if pid in self.id2annotations.keys()]

        self.patient_paths = [data_path + '/' + p + '.mhd' for p in patient_ids_pos]

        self.nsamples = len(self.patient_paths)

class LunaScanPositiveLungMaskDataGenerator(LunaDataGenerator):

    def __init__(self, data_path, batch_size, transform_params, data_prep_fun, rng,

                 full_batch, random, infinite, patient_ids=None, **kwargs):

        super(LunaScanPositiveLungMaskDataGenerator, self).__init__(data_path, transform_params,

                                                                    data_prep_fun, rng,

                                                                    random, infinite, patient_ids, **kwargs)

    def generate(self):

        while True:

            rand_idxs = np.arange(self.nsamples)

            if self.random:

                self.rng.shuffle(rand_idxs)

            for pos in xrange(0, len(rand_idxs)):

                idx = rand_idxs[pos]

                patient_path = self.patient_paths[idx]

                pid = utils_lung.extract_pid_filename(patient_path)

                img, origin, pixel_spacing = utils_lung.read_mhd(patient_path)

                x, y, lung_mask, annotations, tf_matrix = self.data_prep_fun(data=img,

                                                                             pixel_spacing=pixel_spacing,

                                                                             luna_annotations=

                                                                             self.id2annotations[pid],

                                                                             luna_origin=origin)

                x = np.float32(x)[None, None, :, :, :]

                y = np.float32(y)[None, None, :, :, :]

                lung_mask = np.float32(lung_mask)[None, None, :, :, :]

                yield x, y, lung_mask, annotations, tf_matrix, pid

            if not self.infinite:

                break

class LunaScanMaskPositiveDataGenerator(LunaDataGenerator):

    def __init__(self, data_path, seg_data_path, batch_size, transform_params, data_prep_fun, rng,

                 full_batch, random, infinite, patient_ids=None, **kwargs):

        super(LunaScanMaskPositiveDataGenerator, self).__init__(data_path, transform_params,

                                                                    data_prep_fun, rng,

                                                                    random, infinite, patient_ids, **kwargs)

        self.seg_data_path = seg_data_path

        self.mask_paths = [seg_data_path + '/' + p + '.mhd' for p in self.patient_ids]

    def generate(self):

        while True:

            rand_idxs = np.arange(self.nsamples)

            if self.random:

                self.rng.shuffle(rand_idxs)

            for pos in xrange(0, len(rand_idxs)):

                idx = rand_idxs[pos]

                ct_scan_path = self.patient_paths[idx]

                mask_path = self.mask_paths[idx]

                pid = utils_lung.extract_pid_filename(ct_scan_path)

                ct_scan, ct_origin, ct_pixel_spacing = utils_lung.read_mhd(ct_scan_path)

                mask, mask_origin, mask_pixel_spacing = utils_lung.read_mhd(mask_path)

                assert(sum(abs(ct_origin-mask_origin)) < 1e-9)

                assert(sum(abs(ct_pixel_spacing-mask_pixel_spacing)) < 1e-9)

                ct, lung_mask, annotations, tf_matrix = self.data_prep_fun(ct_scan=ct_scan, mask=mask,

                                                                             pixel_spacing=ct_pixel_spacing,

                                                                             luna_annotations=

                                                                             self.id2annotations[pid],

                                                                             luna_origin=ct_origin)

                ct = np.float32(ct)[None, None, :, :, :]

                lung_mask = np.float32(lung_mask)[None, None, :, :, :]

                yield ct, lung_mask, annotations, tf_matrix, pid

            if not self.infinite:

                break

#for lung segmentation, does not work yet

class PatchLunaDataGenerator(object):

    def __init__(self, ct_data_path, seg_data_path, batch_size, transform_params, data_prep_fun, rng,

                 full_batch, random, infinite, patient_ids=None, **kwargs):

        if patient_ids:

            self.patient_ids = patient_ids

            #self.patient_paths = [data_path + '/' + p + '.mhd' for p in patient_ids]

        else:

            patient_paths = utils_lung.get_patient_data_paths(data_path)

            #self.patient_paths = [p for p in patient_paths if '.mhd' in p]

            self.patient_ids = [utils_lung.extract_pid_filename(p) for p in self.patient_paths]\

        self.nsamples = len(self.patient_ids)

        self.ct_data_path = ct_data_path

        self.seg_data_path = seg_data_path

        self.rng = rng

        self.random = random

        self.infinite = infinite

        self.data_prep_fun = data_prep_fun

        self.transform_params = transform_params

        self.batch_size = batch_size

        self.full_batch = full_batch

    def generate(self):

        while True:

            rand_idxs = np.arange(self.nsamples)

            if self.random:

                self.rng.shuffle(rand_idxs)

            for pos in xrange(0, len(rand_idxs), self.batch_size):

                idxs_batch = rand_idxs[pos:pos + self.batch_size]

                nb = len(idxs_batch)

                # allocate batches

                x_batch = np.zeros((nb, 1) + self.transform_params['patch_size'], dtype='float32')

                y_batch = np.zeros((nb, 1) + self.transform_params['patch_size'], dtype='float32')

                patients_ids = []

                for i, idx in enumerate(idxs_batch):

                    pid = self.patient_ids[idx]

                    ct_path = self.ct_data_path + pid + '.mhd'

                    seg_path = self.seg_data_path + pid + '.mhd'

                    patients_ids.append(pid)

                    ct_img, ct_origin, ct_pixel_spacing = utils_lung.read_mhd(ct_path)

                    seg_img, seg_origin, seg_pixel_spacing = utils_lung.read_mhd(seg_path)

                    assert(np.sum(ct_origin-seg_origin) <  1e-9)

                    assert(np.sum(ct_pixel_spacing-seg_pixel_spacing) <  1e-9)

                    print 'ct_img.shape', ct_img.shape

                    print 'seg_img.shape', seg_img.shape

                    w,h,d = self.transform_params['patch_size']

                    patch_center = [self.rng.randint(w/2, ct_img.shape[0]-w/2),

                                    self.rng.randint(h/2, ct_img.shape[1]-h/2),

                                    self.rng.randint(d/2, ct_img.shape[1]-d/2)]

                    print patch_center

                    x_batch[i, 0, :, :, :], y_batch[i, 0, :, :, :]  = self.data_prep_fun(ct_img=ct_img, seg_img=seg_img,

                                                                    patch_center=patch_center,

                                                                    pixel_spacing=ct_pixel_spacing,

                                                                    luna_origin=ct_origin)

                    # y_batch[i, 0, :, :, :],  = self.data_prep_fun(data=seg_img,

                    #                                                 patch_center=patch_center,

                    #                                                 pixel_spacing=seg_pixel_spacing,

                    #                                                 luna_origin=seg_origin)

                if self.full_batch:

                    if nb == self.batch_size:

                        yield x_batch, y_batch, patients_ids

                else:

                    yield x_batch, y_batch, patients_ids

            if not self.infinite:

                break

#works, tested

class LunaScanDataGenerator(object):

    def __init__(self, ct_data_path, seg_data_path, patient_ids=None, **kwargs):

        if patient_ids:

            self.patient_ids = patient_ids

            #self.patient_paths = [data_path + '/' + p + '.mhd' for p in patient_ids]

        else:

            patient_paths = utils_lung.get_patient_data_paths(ct_data_path)

            #self.patient_paths = [p for p in patient_paths if '.mhd' in p]

            self.patient_ids = [utils_lung.extract_pid_filename(p) for p in self.patient_paths]\

        self.nsamples = len(self.patient_ids)

        self.ct_data_path = ct_data_path

        self.seg_data_path = seg_data_path

    def generate(self):

        for pid in self.patient_ids:

            ct_path = self.ct_data_path + pid + '.mhd'

            seg_path = self.seg_data_path + pid + '.mhd'

            ct_img, ct_origin, ct_pixel_spacing = utils_lung.read_mhd(ct_path)

            seg_img, seg_origin, seg_pixel_spacing = utils_lung.read_mhd(seg_path)

            assert(np.sum(ct_origin-seg_origin) <  1e-9)

            assert(np.sum(ct_pixel_spacing-seg_pixel_spacing) <  1e-9)

            print 'ct_img.shape', ct_img.shape

            print 'seg_img.shape', seg_img.shape

            yield ct_img, seg_img, pid

class PatchPositiveLunaDataGenerator(object):

    def __init__(self, data_path, batch_size, transform_params, data_prep_fun, rng,

                 full_batch, random, infinite, patient_ids=None, **kwargs):

        self.id2annotations = utils_lung.read_luna_annotations(pathfinder.LUNA_LABELS_PATH)

        if patient_ids:

            self.patient_paths = [data_path + '/' + p + '.mhd' for p in patient_ids]

        else:

            patient_paths = utils_lung.get_patient_data_paths(data_path)

            self.patient_paths = [p for p in patient_paths if '.mhd' in p]

        patient_ids_all = [utils_lung.extract_pid_filename(p) for p in self.patient_paths]

        patient_ids_pos = [pid for pid in patient_ids_all if pid in self.id2annotations.keys()]

        self.patient_paths = [data_path + '/' + p + '.mhd' for p in patient_ids_pos]

        self.nsamples = len(self.patient_paths)

        self.data_path = data_path

        self.rng = rng

        self.random = random

        self.infinite = infinite

        self.data_prep_fun = data_prep_fun

        self.transform_params = transform_params

        self.batch_size = batch_size

        self.full_batch = full_batch

    def generate(self):

        while True:

            rand_idxs = np.arange(self.nsamples)

            if self.random:

                self.rng.shuffle(rand_idxs)

            for pos in xrange(0, len(rand_idxs), self.batch_size):

                idxs_batch = rand_idxs[pos:pos + self.batch_size]

                nb = len(idxs_batch)

                # allocate batches

                x_batch = np.zeros((nb, 1) + self.transform_params['patch_size'], dtype='float32')

                y_batch = np.zeros((nb, 1) + self.transform_params['patch_size'], dtype='float32')

                patients_ids = []

                for i, idx in enumerate(idxs_batch):

                    patient_path = self.patient_paths[idx]

                    id = utils_lung.extract_pid_filename(patient_path)

                    patients_ids.append(id)

                    img, origin, pixel_spacing = utils_lung.read_mhd(patient_path)

                    patient_annotations = self.id2annotations[id]

                    patch_center = patient_annotations[self.rng.randint(len(patient_annotations))]

                    x_batch[i, 0, :, :, :], y_batch[i, 0, :, :, :] = self.data_prep_fun(data=img,

                                                                                        patch_center=patch_center,

                                                                                        pixel_spacing=pixel_spacing,

                                                                                        luna_annotations=patient_annotations,

                                                                                        luna_origin=origin)

                if self.full_batch:

                    if nb == self.batch_size:

                        yield x_batch, y_batch, patients_ids

                else:

                    yield x_batch, y_batch, patients_ids

            if not self.infinite:

                break

class ValidPatchPositiveLunaDataGenerator(object):

    def __init__(self, data_path, transform_params, patient_ids, data_prep_fun, **kwargs):

        id2positive_annotations = utils_lung.read_luna_annotations(pathfinder.LUNA_LABELS_PATH)

        self.id2positive_annotations = {}

        self.id2patient_path = {}

        n_positive = 0

        for pid in patient_ids:

            if pid in id2positive_annotations:

                self.id2positive_annotations[pid] = id2positive_annotations[pid]

                n_pos = len(id2positive_annotations[pid])

                self.id2patient_path[pid] = data_path + '/' + pid + '.mhd'

                n_positive += n_pos

        self.nsamples = n_positive

        self.data_path = data_path

        self.data_prep_fun = data_prep_fun

        self.transform_params = transform_params

    def generate(self):

        for pid in self.id2positive_annotations.iterkeys():

            for patch_center in self.id2positive_annotations[pid]:

                patient_path = self.id2patient_path[pid]

                img, origin, pixel_spacing = utils_lung.read_mhd(patient_path)

                patient_annotations = self.id2positive_annotations[pid]

                x_batch, y_batch = self.data_prep_fun(data=img,

                                                      patch_center=patch_center,

                                                      pixel_spacing=pixel_spacing,

                                                      luna_annotations=patient_annotations,

                                                      luna_origin=origin)

                x_batch = np.float32(x_batch)[None, None, :, :, :]

                y_batch = np.float32(y_batch)[None, None, :, :, :]

                yield x_batch, y_batch, [pid]

class CandidatesLunaDataGenerator(object):

    def __init__(self, data_path, batch_size, transform_params, patient_ids, data_prep_fun, rng,

                 full_batch, random, infinite, positive_proportion, **kwargs):

        id2positive_annotations = utils_lung.read_luna_annotations(pathfinder.LUNA_LABELS_PATH)

        id2negative_annotations = utils_lung.read_luna_negative_candidates(pathfinder.LUNA_CANDIDATES_PATH)

        self.file_extension = '.pkl' if 'pkl' in data_path else '.mhd'

        self.id2positive_annotations = {}

        self.id2negative_annotations = {}

        self.patient_paths = []

        n_positive, n_negative = 0, 0

        for pid in patient_ids:

            if pid in id2positive_annotations:

                self.id2positive_annotations[pid] = id2positive_annotations[pid]

                self.id2negative_annotations[pid] = id2negative_annotations[pid]

                self.patient_paths.append(data_path + '/' + pid + self.file_extension)

                n_positive += len(id2positive_annotations[pid])

                n_negative += len(id2negative_annotations[pid])

        print 'n positive', n_positive

        print 'n negative', n_negative

        self.nsamples = len(self.patient_paths)

        print 'n patients', self.nsamples

        self.data_path = data_path

        self.batch_size = batch_size

        self.rng = rng

        self.full_batch = full_batch

        self.random = random

        self.infinite = infinite

        self.data_prep_fun = data_prep_fun

        self.transform_params = transform_params

        self.positive_proportion = positive_proportion

    def generate(self):

        while True:

            rand_idxs = np.arange(self.nsamples)

            if self.random:

                self.rng.shuffle(rand_idxs)

            for pos in xrange(0, len(rand_idxs), self.batch_size):

                idxs_batch = rand_idxs[pos:pos + self.batch_size]

                nb = len(idxs_batch)

                # allocate batches

                x_batch = np.zeros((nb, 1) + self.transform_params['patch_size'], dtype='float32')

                y_batch = np.zeros((nb, 1), dtype='float32')

                patients_ids = []

                for i, idx in enumerate(idxs_batch):

                    patient_path = self.patient_paths[idx]

                    id = utils_lung.extract_pid_filename(patient_path, self.file_extension)

                    patients_ids.append(id)

                    img, origin, pixel_spacing = utils_lung.read_pkl(patient_path) \

                        if self.file_extension == '.pkl' else utils_lung.read_mhd(patient_path)

                    if i < np.rint(self.batch_size * self.positive_proportion):

                        patient_annotations = self.id2positive_annotations[id]

                    else:

                        patient_annotations = self.id2negative_annotations[id]

                    patch_center = patient_annotations[self.rng.randint(len(patient_annotations))]

                    y_batch[i] = float(patch_center[-1] > 0)

                    x_batch[i, 0, :, :, :] = self.data_prep_fun(data=img,

                                                                patch_center=patch_center,

                                                                pixel_spacing=pixel_spacing,

                                                                luna_origin=origin)

                if self.full_batch:

                    if nb == self.batch_size:

                        yield x_batch, y_batch, patients_ids

                else:

                    yield x_batch, y_batch, patients_ids

            if not self.infinite:

                break

class CandidatesLunaDataGenerator(object):

    def __init__(self, data_path, batch_size, transform_params, patient_ids, data_prep_fun, rng,

                 full_batch, random, infinite, positive_proportion, return_malignancy=False, **kwargs):

        id2positive_annotations = utils_lung.read_luna_annotations(pathfinder.LUNA_LABELS_PATH)

        id2negative_annotations = utils_lung.read_luna_negative_candidates(pathfinder.LUNA_CANDIDATES_PATH)

        self.file_extension = '.pkl' if 'pkl' in data_path else '.mhd'

        self.id2positive_annotations = {}

        self.id2negative_annotations = {}

        self.patient_paths = []

        n_positive, n_negative = 0, 0

        for pid in patient_ids:

            if pid in id2positive_annotations:

                self.id2positive_annotations[pid] = id2positive_annotations[pid]

                self.id2negative_annotations[pid] = id2negative_annotations[pid]

                self.patient_paths.append(data_path + '/' + pid + self.file_extension)

                n_positive += len(id2positive_annotations[pid])

                n_negative += len(id2negative_annotations[pid])

        print 'n positive', n_positive

        print 'n negative', n_negative

        self.nsamples = len(self.patient_paths)

        print 'n patients', self.nsamples

        self.data_path = data_path

        self.batch_size = batch_size

        self.rng = rng

        self.full_batch = full_batch

        self.random = random

        self.infinite = infinite

        self.data_prep_fun = data_prep_fun

        self.transform_params = transform_params

        self.positive_proportion = positive_proportion

        self.return_malignancy = return_malignancy

    def generate(self):

        while True:

            rand_idxs = np.arange(self.nsamples)

            if self.random:

                self.rng.shuffle(rand_idxs)

            for pos in xrange(0, len(rand_idxs), self.batch_size):

                idxs_batch = rand_idxs[pos:pos + self.batch_size]

                nb = len(idxs_batch)

                # allocate batches

                x_batch = np.zeros((nb,) + self.transform_params['patch_size'], dtype='float32')

                y_batch = np.zeros((nb,), dtype='float32')

                patients_ids = []

                for i, idx in enumerate(idxs_batch):

                    patient_path = self.patient_paths[idx]

                    id = utils_lung.extract_pid_filename(patient_path, self.file_extension)

                    patients_ids.append(id)

                    img, origin, pixel_spacing = utils_lung.read_pkl(patient_path) \

                        if self.file_extension == '.pkl' else utils_lung.read_mhd(patient_path)

                    if i < np.rint(self.batch_size * self.positive_proportion):

                        patient_annotations = self.id2positive_annotations[id]

                    else:

                        patient_annotations = self.id2negative_annotations[id]

                    patch_center = patient_annotations[self.rng.randint(len(patient_annotations))]

                    if self.return_malignancy:

                        y_batch[i] = np.float32(diameter_to_prob(patch_center[-1]))

                    else:

                        y_batch[i] = float(patch_center[-1] > 0) 

                    x_batch[i, :, :, :] = self.data_prep_fun(data=img,

                                                                patch_center=patch_center,

                                                                pixel_spacing=pixel_spacing,

                                                                luna_origin=origin)

                if self.full_batch:

                    if nb == self.batch_size:

                        yield x_batch, y_batch, patients_ids

                else:

                    yield x_batch, y_batch, patients_ids

            if not self.infinite:

                break

class CandidatesLunaValidDataGenerator(object):

    def __init__(self, data_path, transform_params, patient_ids, data_prep_fun, return_malignancy=False, **kwargs):

        rng = np.random.RandomState(42)  # do not change this!!!

        id2positive_annotations = utils_lung.read_luna_annotations(pathfinder.LUNA_LABELS_PATH)

        id2negative_annotations = utils_lung.read_luna_negative_candidates(pathfinder.LUNA_CANDIDATES_PATH)

        self.file_extension = '.pkl' if 'pkl' in data_path else '.mhd'

        self.id2positive_annotations = {}

        self.id2negative_annotations = {}

        self.id2patient_path = {}

        n_positive, n_negative = 0, 0

        for pid in patient_ids:

            if pid in id2positive_annotations:

                self.id2positive_annotations[pid] = id2positive_annotations[pid]

                negative_annotations = id2negative_annotations[pid]

                n_pos = len(id2positive_annotations[pid])

                n_neg = len(id2negative_annotations[pid])

                neg_idxs = rng.choice(n_neg, size=n_pos, replace=False)

                negative_annotations_selected = []

                for i in neg_idxs:

                    negative_annotations_selected.append(negative_annotations[i])

                self.id2negative_annotations[pid] = negative_annotations_selected

                self.id2patient_path[pid] = data_path + '/' + pid + self.file_extension

                n_positive += n_pos

                n_negative += n_pos

        print 'n positive', n_positive

        print 'n negative', n_negative

        self.nsamples = len(self.id2patient_path)

        self.data_path = data_path

        self.rng = rng

        self.data_prep_fun = data_prep_fun

        self.transform_params = transform_params

        self.return_malignancy = return_malignancy

    def generate(self):

        for pid in self.id2positive_annotations.iterkeys():

            for patch_center in self.id2positive_annotations[pid]:

                patient_path = self.id2patient_path[pid]

                img, origin, pixel_spacing = utils_lung.read_pkl(patient_path) \

                    if self.file_extension == '.pkl' else utils_lung.read_mhd(patient_path)

                if self.return_malignancy:

                    y_batch = np.array([diameter_to_prob(patch_center[-1])], dtype='float32')

                else:

                    y_batch = np.array([1.], dtype='float32')

                x_batch = np.float32(self.data_prep_fun(data=img,

                                                        patch_center=patch_center,

                                                        pixel_spacing=pixel_spacing,

                                                        luna_origin=origin))[None, :, :, :]

                yield x_batch, y_batch, [pid]

            for patch_center in self.id2negative_annotations[pid]:

                patient_path = self.id2patient_path[pid]

                img, origin, pixel_spacing = utils_lung.read_pkl(patient_path) \

                    if self.file_extension == '.pkl' else utils_lung.read_mhd(patient_path)

                y_batch = np.array([0.], dtype='float32')

                x_batch = np.float32(self.data_prep_fun(data=img,

                                                        patch_center=patch_center,

                                                        pixel_spacing=pixel_spacing,

                                                        luna_origin=origin))[None, :, :, :]

                yield x_batch, y_batch, [pid]

class FixedCandidatesLunaDataGenerator(object):

    def __init__(self, data_path, transform_params, id2candidates_path, data_prep_fun, top_n=None):

        self.file_extension = '.pkl' if 'pkl' in data_path else '.mhd'

        self.id2candidates_path = id2candidates_path

        self.id2patient_path = {}

        for pid in id2candidates_path.keys():

            self.id2patient_path[pid] = data_path + '/' + pid + self.file_extension

        self.nsamples = len(self.id2patient_path)

        self.data_path = data_path

        self.data_prep_fun = data_prep_fun

        self.transform_params = transform_params

        self.top_n = top_n

    def generate(self):

        for pid in self.id2candidates_path.iterkeys():

            patient_path = self.id2patient_path[pid]

            print 'PATIENT', pid

            candidates = utils.load_pkl(self.id2candidates_path[pid])

            if self.top_n is not None:

                candidates = candidates[:self.top_n]

                print candidates

            print 'n blobs', len(candidates)

            img, origin, pixel_spacing = utils_lung.read_pkl(patient_path) \

                if self.file_extension == '.pkl' else utils_lung.read_mhd(patient_path)

            for candidate in candidates:

                y_batch = np.array(candidate, dtype='float32')

                patch_center = candidate[:3]

                x_batch = np.float32(self.data_prep_fun(data=img,

                                                        patch_center=patch_center,

                                                        pixel_spacing=pixel_spacing,

                                                        luna_origin=origin))[None, None, :, :, :]

                yield x_batch, y_batch, [pid]

class CandidatesLunaSizeDataGenerator(object):

    def __init__(self, data_path, batch_size, transform_params, patient_ids, data_prep_fun, rng,

                 full_batch, random, infinite, positive_proportion, **kwargs):

        id2positive_annotations = utils_lung.read_luna_annotations(pathfinder.LUNA_LABELS_PATH)

        id2negative_annotations = utils_lung.read_luna_negative_candidates(pathfinder.LUNA_CANDIDATES_PATH)

        self.file_extension = '.pkl' if 'pkl' in data_path else '.mhd'

        self.id2positive_annotations = {}

        self.id2negative_annotations = {}

        self.patient_paths = []

        n_positive, n_negative = 0, 0

        for pid in patient_ids:

            if pid in id2positive_annotations:

                self.id2positive_annotations[pid] = id2positive_annotations[pid]

                self.id2negative_annotations[pid] = id2negative_annotations[pid]

                self.patient_paths.append(data_path + '/' + pid + self.file_extension)

                n_positive += len(id2positive_annotations[pid])

                n_negative += len(id2negative_annotations[pid])

        print 'n positive', n_positive

        print 'n negative', n_negative

        self.nsamples = len(self.patient_paths)

        print 'n patients', self.nsamples

        self.data_path = data_path

        self.batch_size = batch_size

        self.rng = rng

        self.full_batch = full_batch

        self.random = random

        self.infinite = infinite

        self.data_prep_fun = data_prep_fun

        self.transform_params = transform_params

        self.positive_proportion = positive_proportion

    def generate(self):

        while True:

            rand_idxs = np.arange(self.nsamples)

            if self.random:

                self.rng.shuffle(rand_idxs)

            for pos in xrange(0, len(rand_idxs), self.batch_size):

                idxs_batch = rand_idxs[pos:pos + self.batch_size]

                nb = len(idxs_batch)

                # allocate batches

                x_batch = np.zeros((nb, 1) + self.transform_params['patch_size'], dtype='float32')

                y_batch = np.zeros((nb, 1), dtype='float32')

                patients_ids = []

                for i, idx in enumerate(idxs_batch):

                    patient_path = self.patient_paths[idx]

                    id = utils_lung.extract_pid_filename(patient_path, self.file_extension)

                    patients_ids.append(id)

                    img, origin, pixel_spacing = utils_lung.read_pkl(patient_path) \

                        if self.file_extension == '.pkl' else utils_lung.read_mhd(patient_path)

                    if i < np.rint(self.batch_size * self.positive_proportion):

                        patient_annotations = self.id2positive_annotations[id]

                    else:

                        patient_annotations = self.id2negative_annotations[id]

                    patch_center = patient_annotations[self.rng.randint(len(patient_annotations))]

                    y_batch[i] = float(patch_center[-1])

                    x_batch[i, 0, :, :, :] = self.data_prep_fun(data=img,

                                                                patch_center=patch_center,

                                                                pixel_spacing=pixel_spacing,

                                                                luna_origin=origin)

                if self.full_batch:

                    if nb == self.batch_size:

                        yield x_batch, y_batch, patients_ids

                else:

                    yield x_batch, y_batch, patients_ids

            if not self.infinite:

                break

class CandidatesLunaSizeValidDataGenerator(object):

    def __init__(self, data_path, transform_params, patient_ids, data_prep_fun, **kwargs):

        rng = np.random.RandomState(42)  # do not change this!!!

        id2positive_annotations = utils_lung.read_luna_annotations(pathfinder.LUNA_LABELS_PATH)

        id2negative_annotations = utils_lung.read_luna_negative_candidates(pathfinder.LUNA_CANDIDATES_PATH)

        self.file_extension = '.pkl' if 'pkl' in data_path else '.mhd'

        self.id2positive_annotations = {}

        self.id2negative_annotations = {}

        self.id2patient_path = {}

        n_positive, n_negative = 0, 0

        for pid in patient_ids:

            if pid in id2positive_annotations:

                self.id2positive_annotations[pid] = id2positive_annotations[pid]

                negative_annotations = id2negative_annotations[pid]

                n_pos = len(id2positive_annotations[pid])

                n_neg = len(id2negative_annotations[pid])

                neg_idxs = rng.choice(n_neg, size=n_pos, replace=False)

                negative_annotations_selected = []

                for i in neg_idxs:

                    negative_annotations_selected.append(negative_annotations[i])

                self.id2negative_annotations[pid] = negative_annotations_selected

                self.id2patient_path[pid] = data_path + '/' + pid + self.file_extension

                n_positive += n_pos

                n_negative += n_pos

        print 'n positive', n_positive

        print 'n negative', n_negative

        self.nsamples = len(self.id2patient_path)

        self.data_path = data_path

        self.rng = rng

        self.data_prep_fun = data_prep_fun

        self.transform_params = transform_params

    def generate(self):

        for pid in self.id2positive_annotations.iterkeys():

            for patch_center in self.id2positive_annotations[pid]:

                patient_path = self.id2patient_path[pid]

                img, origin, pixel_spacing = utils_lung.read_pkl(patient_path) \

                    if self.file_extension == '.pkl' else utils_lung.read_mhd(patient_path)

                y_batch = np.array([[float(patch_center[-1])]], dtype='float32')

                x_batch = np.float32(self.data_prep_fun(data=img,

                                                        patch_center=patch_center,

                                                        pixel_spacing=pixel_spacing,

                                                        luna_origin=origin))[None, None, :, :, :]

                yield x_batch, y_batch, [pid]

            for patch_center in self.id2negative_annotations[pid]:

                patient_path = self.id2patient_path[pid]

                img, origin, pixel_spacing = utils_lung.read_pkl(patient_path) \

                    if self.file_extension == '.pkl' else utils_lung.read_mhd(patient_path)

                y_batch = np.array([[0.]], dtype='float32')

                x_batch = np.float32(self.data_prep_fun(data=img,

                                                        patch_center=patch_center,

                                                        pixel_spacing=pixel_spacing,

                                                        luna_origin=origin))[None, None, :, :, :]

                yield x_batch, y_batch, [pid]

class CandidatesLunaSizeBinDataGenerator(object):

    def __init__(self, data_path, batch_size, transform_params, patient_ids, data_prep_fun, rng,

                 full_batch, random, infinite, positive_proportion, bin_borders = [4,8,20,50], **kwargs):

        id2positive_annotations = utils_lung.read_luna_annotations(pathfinder.LUNA_LABELS_PATH)

        id2negative_annotations = utils_lung.read_luna_negative_candidates(pathfinder.LUNA_CANDIDATES_PATH)

        self.file_extension = '.pkl' if 'pkl' in data_path else '.mhd'

        self.id2positive_annotations = {}

        self.id2negative_annotations = {}

        self.patient_paths = []

        n_positive, n_negative = 0, 0

        for pid in patient_ids:

            if pid in id2positive_annotations:

                self.id2positive_annotations[pid] = id2positive_annotations[pid]

                self.id2negative_annotations[pid] = id2negative_annotations[pid]

                self.patient_paths.append(data_path + '/' + pid + self.file_extension)

                n_positive += len(id2positive_annotations[pid])

                n_negative += len(id2negative_annotations[pid])

        print 'n positive', n_positive

        print 'n negative', n_negative

        self.nsamples = len(self.patient_paths)

        print 'n patients', self.nsamples

        self.data_path = data_path

        self.batch_size = batch_size

        self.rng = rng

        self.full_batch = full_batch

        self.random = random

        self.infinite = infinite

        self.data_prep_fun = data_prep_fun

        self.transform_params = transform_params

        self.positive_proportion = positive_proportion

        self.bin_borders = bin_borders

    def generate(self):

        while True:

            rand_idxs = np.arange(self.nsamples)

            if self.random:

                self.rng.shuffle(rand_idxs)

            for pos in xrange(0, len(rand_idxs), self.batch_size):

                idxs_batch = rand_idxs[pos:pos + self.batch_size]

                nb = len(idxs_batch)

                # allocate batches

                x_batch = np.zeros((nb,) + self.transform_params['patch_size'], dtype='float32')

                y_batch = np.zeros((nb,), dtype='float32')

                patients_ids = []

                for i, idx in enumerate(idxs_batch):

                    patient_path = self.patient_paths[idx]

                    id = utils_lung.extract_pid_filename(patient_path, self.file_extension)

                    patients_ids.append(id)

                    img, origin, pixel_spacing = utils_lung.read_pkl(patient_path) \

                        if self.file_extension == '.pkl' else utils_lung.read_mhd(patient_path)

                    if i < np.rint(self.batch_size * self.positive_proportion):

                        patient_annotations = self.id2positive_annotations[id]

                    else:

                        patient_annotations = self.id2negative_annotations[id]

                    patch_center = patient_annotations[self.rng.randint(len(patient_annotations))]

                    diameter = patch_center[-1]

                    if diameter > 0.:

                        ybin = 0

                        for idx, border in enumerate(self.bin_borders):

                            if diameter<border:

                                ybin = idx

                                break                            

                        y_batch[i] = 1. + ybin

                    else:

                        y_batch[i] = 0. 

                    #print 'y_batch[i]', y_batch[i], 'diameter', diameter

                    x_batch[i, :, :, :] = self.data_prep_fun(data=img,

                                                                patch_center=patch_center,

                                                                pixel_spacing=pixel_spacing,

                                                                luna_origin=origin)

                if self.full_batch:

                    if nb == self.batch_size:

                        yield x_batch, y_batch, patients_ids

                else:

                    yield x_batch, y_batch, patients_ids

            if not self.infinite:

                break

class CandidatesLunaSizeBinValidDataGenerator(object):

    def __init__(self, data_path, transform_params, patient_ids, data_prep_fun, bin_borders = [4,8,20,50], **kwargs):

        rng = np.random.RandomState(42)  # do not change this!!!

        id2positive_annotations = utils_lung.read_luna_annotations(pathfinder.LUNA_LABELS_PATH)

        id2negative_annotations = utils_lung.read_luna_negative_candidates(pathfinder.LUNA_CANDIDATES_PATH)

        self.file_extension = '.pkl' if 'pkl' in data_path else '.mhd'

        self.id2positive_annotations = {}

        self.id2negative_annotations = {}

        self.id2patient_path = {}

        n_positive, n_negative = 0, 0

        for pid in patient_ids:

            if pid in id2positive_annotations:

                self.id2positive_annotations[pid] = id2positive_annotations[pid]

                negative_annotations = id2negative_annotations[pid]

                n_pos = len(id2positive_annotations[pid])

                n_neg = len(id2negative_annotations[pid])

                neg_idxs = rng.choice(n_neg, size=n_pos, replace=False)

                negative_annotations_selected = []

                for i in neg_idxs:

                    negative_annotations_selected.append(negative_annotations[i])

                self.id2negative_annotations[pid] = negative_annotations_selected

                self.id2patient_path[pid] = data_path + '/' + pid + self.file_extension

                n_positive += n_pos

                n_negative += n_pos

        print 'n positive', n_positive

        print 'n negative', n_negative

        self.nsamples = len(self.id2patient_path)

        self.data_path = data_path

        self.rng = rng

        self.data_prep_fun = data_prep_fun

        self.transform_params = transform_params

        self.bin_borders = bin_borders

    def generate(self):

        for pid in self.id2positive_annotations.iterkeys():

            for patch_center in self.id2positive_annotations[pid]:

                patient_path = self.id2patient_path[pid]

                img, origin, pixel_spacing = utils_lung.read_pkl(patient_path) \

                    if self.file_extension == '.pkl' else utils_lung.read_mhd(patient_path)

                diameter = patch_center[3]                        

                ybin = 0

                for idx, border in enumerate(self.bin_borders):

                    if diameter<border:

                        ybin = idx

                        break  

                y_batch = np.array([1. + ybin], dtype='float32')

                x_batch = np.float32(self.data_prep_fun(data=img,

                                                        patch_center=patch_center,

                                                        pixel_spacing=pixel_spacing,

                                                        luna_origin=origin))[None, :, :, :]

                yield x_batch, y_batch, [pid]

            for patch_center in self.id2negative_annotations[pid]:

                patient_path = self.id2patient_path[pid]

                img, origin, pixel_spacing = utils_lung.read_pkl(patient_path) \

                    if self.file_extension == '.pkl' else utils_lung.read_mhd(patient_path)

                y_batch = np.array([0.], dtype='float32')

                x_batch = np.float32(self.data_prep_fun(data=img,

                                                        patch_center=patch_center,

                                                        pixel_spacing=pixel_spacing,

                                                        luna_origin=origin))[None, :, :, :]

                yield x_batch, y_batch, [pid]

class CandidatesLunaPropsDataGenerator(object):

    def __init__(self, data_path, batch_size, transform_params, patient_ids, data_prep_fun, rng,

                 full_batch, random, infinite, 

                 positive_proportion,

                 order_objectives,

                 property_type,

                 property_bin_borders = None,

                 return_enable_target_vector = False, **kwargs):

        id2positive_annotations = utils_lung.read_luna_annotations(pathfinder.LUNA_LABELS_PATH)

        id2negative_annotations = utils_lung.read_luna_negative_candidates(pathfinder.LUNA_CANDIDATES_PATH)

        self.file_extension = '.pkl' if 'pkl' in data_path else '.mhd'

        self.id2positive_annotations = {}

        self.id2negative_annotations = {}

        self.all_pids = patient_ids

        self.pos_pids = []

        self.neg_pids = []

        n_positive, n_negative = 0, 0

        for pid in patient_ids:

            if pid in id2positive_annotations:

                self.id2positive_annotations[pid] = id2positive_annotations[pid]

                self.id2negative_annotations[pid] = id2negative_annotations[pid]

                self.pos_pids.append(pid)

                n_positive += len(id2positive_annotations[pid])

                n_negative += len(id2negative_annotations[pid])

            elif pid in id2negative_annotations:

                self.id2negative_annotations[pid] = id2negative_annotations[pid]

                self.neg_pids.append(pid)

                n_negative += len(id2negative_annotations[pid])

            else:

                print 'WARNING something weird happens'

        print 'n positive', n_positive

        print 'n negative', n_negative

        self.n_neg_cans = n_negative

        self.n_pos_cans = n_positive

        self.n_pos_pids = len(self.pos_pids)

        self.n_neg_pids = len(self.neg_pids)

        self.nsamples = len(self.all_pids)

        print 'n patients', self.nsamples

        self.data_path = data_path

        self.batch_size = batch_size

        self.rng = rng

        self.full_batch = full_batch

        self.random = random

        self.infinite = infinite

        self.data_prep_fun = data_prep_fun

        self.transform_params = transform_params

        self.positive_proportion = positive_proportion

        self.order_objectives = order_objectives

        self.property_bin_borders = property_bin_borders

    self.property_type = property_type

        #self.return_enable_target_vector = return_enable_target_vector

    def L2(self, a,b):

        return ((a[0]-b[0])**2 + (a[1]-b[1])**2 + (a[2]-b[2])**2)**(0.5)

    def build_ground_truth_vector(self, pid, patch_center):

        properties={}

        feature_vector = np.zeros((len(self.order_objectives)), dtype='float32')

        enable_target_vector = np.zeros((len(self.order_objectives)), dtype='float32')

        diameter = patch_center[-1]

        is_nodule  = diameter>0.01

        properties['nodule'] = np.float32(is_nodule)

        if is_nodule:

            if 'size' in self.property_bin_borders:

                properties['size'] = np.digitize(diameter, self.property_bin_borders['size'])

            else:

                properties['size'] = diameter

            patient = utils_lung.read_patient_annotations_luna(pid, pathfinder.LUNA_NODULE_ANNOTATIONS_PATH)

            #find the nodules in the doctor's annotations

            nodule_characteristics = []

            for doctor in patient:

                for nodule in doctor:

                    if "centroid_xyz" in nodule:

                        dist = self.L2(patch_center[:3],nodule["centroid_xyz"][::-1])

                        if  dist < 5:

                            #print 'found a very close nodule at', dist, ': ', patch_center[:3]

                            nodule_characteristics.append(nodule['characteristics'])

            if len(nodule_characteristics)==0:

                print 'WARNING: no nodule found in doctor annotations for ', patch_center

            else:

                #calculate the median property values

                for prop in nodule_characteristics[0]:

                    if prop in self.order_objectives:

                        prop_values = []

                        for nchar in nodule_characteristics:

                            prop_values.append(float(nchar[prop]))

                            random_value = self.rng.choice(np.array(prop_values))

                            if prop in self.property_bin_borders:

                                properties[prop] = np.digitize(random_value, self.property_bin_borders[prop])

                            else:      

                                if self.property_type:

                                    if self.property_type[prop] == 'bounded_continuous':

                                        properties[prop] = (random_value-1) / 4.

                                    else:

                                        properties[prop] = random_value-1

                                else:

                                    raise

        for idx, prop in enumerate(self.order_objectives):

            if prop in properties:

                feature_vector[idx] = properties[prop]

                enable_target_vector[idx] = 1.

        return feature_vector, enable_target_vector

    def generate(self):

        while True:

            # Construct pid set with

            rand_pos_idxs = np.arange(self.n_pos_pids)

            rand_neg_idxs = np.arange(self.n_neg_pids)

            ptr_pos_idcs = 0

            ptr_neg_idcs = 0

            if self.random: