import os
import cv2
import numpy as np
from random import randrange

import utils
import config
import lung_segmentation as ls


segmentation_algo = ls.get_segmentation_algorithm()


class PatientImageLoader(object):
    def __init__(self, images_dir):
        self._images_input = images_dir or config.SEGMENTED_LUNGS_DIR
        self._augment = False

    def load_scans(self, patient):
        if 'augm' in patient:
            self._augment = True
            patient = patient.split('-')[0]
        return utils.load_patient_image(self._images_input, patient)

    @property
    def images_input(self):
        return self._images_input

    @property
    def name(self):
        return 'base_image_loader'


# Tests with the mean scans loader are not using
# lung segmentation, only compressed sorted slices in HU units.
class MeanScansLoader(PatientImageLoader):
    def __init__(self, images_dir=None):
        super(MeanScansLoader, self).__init__(images_dir)

    def load_scans(self, patient):
        image = utils.load_patient_image(self._images_input, patient)
        image = utils.resize(image)

        return utils.get_mean_chunk_slices(image)

    @property
    def name(self):
        return 'mean_scans_loader'


class SegmentedGaussianLungsLoader(PatientImageLoader):
    def __init__(self, images_dir=config.SEGMENTED_LUNGS_DIR):
        super(SegmentedGaussianLungsLoader, self).__init__(images_dir)

    def process_scans(self, image):
        image = np.stack([cv2.GaussianBlur(scan, (5, 5), 0) for scan in image])
        image = utils.resize(image)

        return utils.trim_pad_slices(image, pad_with_existing=False)

    def load_scans(self, patient):
        image = utils.load_patient_image(self._images_input, patient)
        return self.process_scans(image)

    @property
    def name(self):
        return 'segmented_gaussian_lungs_loader'


# Default loader
class SegmentedLungsScansLoader(PatientImageLoader):
    def __init__(self, images_dir=config.SEGMENTED_LUNGS_DIR):
        super(SegmentedLungsScansLoader, self).__init__(images_dir)

    def process_scans(self, image):
        image = segmentation_algo.get_slices_with_nodules(image)
        image = utils.resize(image)

        if self._augment:
            angle = randrange(-15, 15)
            image = utils.rotate_scans(image, angle)

        return utils.trim_pad_slices(image, pad_with_existing=True)

    def load_scans(self, patient):
        image = super(SegmentedLungsScansLoader, self).load_scans(patient)
        return self.process_scans(image)

    @property
    def name(self):
        return 'segmented_lungs_loader'


class NodulesScansLoader(PatientImageLoader):
    def __init__(self, images_dir=config.SEGMENTED_LUNGS_DIR):
        super(NodulesScansLoader, self).__init__(images_dir)

    def process_scans(self, patient):
        image = utils.load_patient_image(self._images_input, patient)
        nodules = segmentation_algo.get_lung_nodules_candidates(image)
        nodules = utils.resize(nodules)

        return utils.trim_pad_slices(nodules)

    def load_scans(self, patient):
        return self.process_scans(patient)

    @property
    def name(self):
        return 'nodules_scans_loader'