import time

from tqdm import tqdm

import numpy as np

import pydicom as dicom

import os

from scipy import ndimage

import matplotlib.pyplot as plt

from pathlib import Path

from skimage import measure

from collections import defaultdict

from sys import argv

from random import shuffle

from lungs.utils import apply_window

# The pixel size/coarseness of the scan differs from scan to scan (e.g. the distance between slices may differ), which can hurt performance of 

# CNN approaches. We can deal with this by isomorphic resampling.

# Below is code to load a scan, which consists of multiple slices, which we simply save in a Python list. Every folder in the dataset is one 

# scan (so one patient). One metadata field is missing, the pixel size in the Z direction, which is the slice thickness. 

# Fortunately we can infer this, and we add this to the metadata.

def is_dcm_file(path):

    name, ext = os.path.splitext(path)

    # DICOM file extension

    if 'dcm' in ext:

        return True

    # NLST file format

    if name.isdigit() and not ext:

        return True

    return False

# Load a volume from the given folder path

def load_scan(path):

    slices = [dicom.read_file(path + '/' + scan) for scan in os.listdir(path) if is_dcm_file(scan)]

    slices.sort(key = lambda x: float(x.ImagePositionPatient[2]))

    if not slices[0].SliceThickness:

        try:

            slice_thickness = np.abs(slices[0].ImagePositionPatient[2] - slices[1].ImagePositionPatient[2])

        except:

            slice_thickness = np.abs(slices[0].SliceLocation - slices[1].SliceLocation)

        for s in slices:

            s.SliceThickness = slice_thickness  

    return slices

# The unit of measurement in CT scans is the **Hounsfield Unit (HU)**, which is a measure of radiodensity. 

# CT scanners are carefully calibrated to accurately measure this.  From Wikipedia:

# By default however, the returned values are not in this unit. Let's fix this.

# Some scanners have cylindrical scanning bounds, but the output volume is square. 

# The pixels that fall outside of these bounds get the fixed value -2000. The first step is setting these values to 0, which currently corresponds to air. 

# Next, let's go back to HU units, by multiplying with the rescale slope and adding the intercept (which are conveniently stored in the metadata of the scans!).

def get_pixels_hu(slices):

    volume = np.stack([s.pixel_array for s in slices])

    # Convert to int16 (from sometimes int16), 

    # should be possible as values should always be low enough (<32k)

    volume = volume.astype(np.int16)

    # Set outside-of-scan pixels to 0

    # The intercept is usually -1024, so air is approximately 0

    volume[volume == -2000] = 0

    # Convert to Hounsfield units (HU)

    for slice_number in range(len(slices)):

        intercept = slices[slice_number].RescaleIntercept

        slope = slices[slice_number].RescaleSlope

        if slope != 1:

            volume[slice_number] = slope * volume[slice_number].astype(np.float64)

            volume[slice_number] = volume[slice_number].astype(np.int16)

        volume[slice_number] += np.int16(intercept)

    return np.array(volume, dtype=np.int16)

# # Resampling

# A scan may have a pixel spacing of `[2.5, 0.5, 0.5]`, which means that the distance between slices is `2.5` millimeters. 

# For a different scan this may be `[1.5, 0.725, 0.725]`, 

# this can be problematic for automatic analysis (e.g. using ConvNets).

# A common method of dealing with this is resampling the full dataset to a certain isotropic resolution. 

# If we choose to resample everything to 1.5mm*1.5mm*1.5mm pixels we can use 3D convnets without worrying about learning zoom/slice thickness invariance. 

# Whilst this may seem like a very simple step, it has quite some edge cases due to rounding. Also, it takes quite a while.

def resample(scan_hu, scan_file, scan, new_spacing, verbose=False):

    # Determine current pixel spacing

    spacing = np.array([scan_file[0].SliceThickness] + list(scan_file[0].PixelSpacing), dtype=np.float32)

    if verbose:

        print('Spacing:', spacing)

    resize_factor = spacing / new_spacing

    new_real_shape = scan_hu.shape * resize_factor

    new_shape = np.round(new_real_shape)

    real_resize_factor = new_shape / scan_hu.shape

    new_spacing = spacing / real_resize_factor

    scan_hu = ndimage.interpolation.zoom(scan_hu, real_resize_factor, mode='nearest')

    return scan_hu, new_spacing 

def largest_label_volume(im, bg=-1):

    vals, counts = np.unique(im, return_counts=True)

    counts = counts[vals != bg]

    vals = vals[vals != bg]

    if len(counts) > 0:

        return vals[np.argmax(counts)]

    else:

        return None

# # Lung segmentation

# In order to reduce the problem space, we segment the lungs (and usually some tissue around it).

# It consists of a series of applications of region growing and morphological operations. In this case, 

# we will use only connected component analysis.

# 

# The steps:  

# * Threshold the volume (-320 HU is a good threshold, but it doesn't matter much for this approach)

# * Do connected components, determine label of air around person, fill this with 1s in the binary volume

# * Optionally: For every axial slice in the scan, determine the largest solid connected component 

# (the body+air around the person), and set others to 0. This fills the structures in the lungs in the mask.

# * Keep only the largest air pocket (the human body has other pockets of air here and there).

def segment_lung_mask(volume, fill_lung_structures=True):

    # not actually binary, but 1 and 2. 

    # 0 is treated as background, which we do not want

    binary_volume = np.array(volume > -320, dtype=np.int8)+1

    labels = measure.label(binary_volume)

    # Pick the pixel in the very corner to determine which label is air.

    #   Improvement: Pick multiple background labels from around the patient

    #   More resistant to "trays" on which the patient lays cutting the air 

    #   around the person in half

    background_label = labels[0,0,0]

    #Fill the air around the person

    binary_volume[background_label == labels] = 2

    # Method of filling the lung structures (that is superior to something like 

    # morphological closing)

    if fill_lung_structures:

        # For every slice we determine the largest solid structure

        for i, axial_slice in enumerate(binary_volume):

            axial_slice = axial_slice - 1

            labeling = measure.label(axial_slice)

            l_max = largest_label_volume(labeling, bg=0)

            if l_max is not None: # This slice contains some lung

                binary_volume[i][labeling != l_max] = 1

    binary_volume -= 1 # Make the volume actual binary

    binary_volume = 1-binary_volume # Invert it, lungs are now 1

    # Remove other air pockets insided body

    labels = measure.label(binary_volume, background=0)

    l_max = largest_label_volume(labels, bg=0)

    if l_max is not None: # There are air pockets

        binary_volume[labels != l_max] = 0

    return binary_volume

def bbox2_3D(volume):

    r = np.any(volume, axis=(1, 2))

    c = np.any(volume, axis=(0, 2))

    z = np.any(volume, axis=(0, 1))

    rmin, rmax = np.where(r)[0][[0, -1]]

    cmin, cmax = np.where(c)[0][[0, -1]]

    zmin, zmax = np.where(z)[0][[0, -1]]

    return rmin, rmax, cmin, cmax, zmin, zmax

def preprocess(scan, errors_map, num_slices=224, crop_size=224, voxel_size=1.5, windowing=False, sample_volume=True, verbose=True):

    orig_scan = load_scan(scan)

    num_orig_slices = len(orig_scan)

    if num_orig_slices < 50:

        errors_map['insufficient_slices'] += 1

        raise ValueError(scan[-4:] + ': Insufficient muber of slices (<50).')

    orig_scan_np = np.stack([s.pixel_array for s in orig_scan]).astype(np.int16)

    scan_hu = get_pixels_hu(orig_scan)

    # Let's resample our patient's pixels to an isomorphic resolution

    resampled_scan, _ = resample(scan_hu, orig_scan, orig_scan_np, [voxel_size, voxel_size, voxel_size], verbose=verbose)

    if verbose:

        print("Shape before resampling:", scan_hu.shape)

        print("Shape after resampling:", resampled_scan.shape)

    if resampled_scan.shape[0] < 180:

        errors_map['small_z'] += 1

        raise ValueError(scan[-4:] + ': Insufficient number of resampled slices (<200).')

    lung_mask = segment_lung_mask(resampled_scan, True)

    z_min, z_max, x_min, x_max, y_min, y_max = bbox2_3D(lung_mask)

    box_size = (z_max - z_min, x_max - x_min, y_max - y_min)

    if verbose:

        print('Lung bounding box (min, max):', (z_min, z_max), (x_min, x_max), (y_min, y_max))

        print('Bounding box size:', box_size)

    for dim in box_size:

        if dim < 100:

            errors_map['seg_error'] += 1

            raise ValueError(scan[-4:] + ': Segmentation error.')   

    lung_center = np.array([z_min + z_max, x_min + x_max, y_min + y_max]) // 2

    context = np.array([num_slices, crop_size, crop_size])

    volume_starts = np.array([max(0, lung_center[i] - context[i] // 2) for i in range(3)])

    volume_ends = np.array([min(resampled_scan.shape[i], lung_center[i] + context[i] // 2) for i in range(3)])

    volume_size = volume_ends - volume_starts

    starts = context // 2 - volume_size // 2

    ends = starts + volume_size

    lungs_padded = np.zeros((num_slices, crop_size, crop_size))

    lungs_padded[starts[0]: ends[0], starts[1]: ends[1], starts[2]: ends[2]] = \

            resampled_scan[volume_starts[0]: volume_ends[0], volume_starts[1]: volume_ends[1], volume_starts[2]: volume_ends[2]]

    if verbose:

        print("Final shape", lungs_padded.shape)

    if sample_volume:

        # Generate an RGB slice for display

        lungs_rgb = np.stack((lungs_padded, lungs_padded, lungs_padded), axis=3)

        lungs_sample_slice = lungs_rgb[lungs_rgb.shape[0] // 2]

    else:

        lungs_sample_slice = None

    return lungs_padded, lungs_sample_slice

def walk_dicom_dirs(base_in, base_out=None, print_dirs=True):

    for root, _, files in os.walk(base_in):

        path = root.split(os.sep)

        if print_dirs:

            print((len(path) - 1) * '---', os.path.basename(root))

        # sample_filename = os.path.splitext(files[0])

        if len(files) >= 50: # and (sample_filename[0].isdigit() or 'dcm' in sample_filename[1]):

            if base_out:

                yield root, base_out + os.path.relpath(root, base_in)

            else:

                yield root

def walk_np_files(base_in, print_dirs=True):

    pathlist = Path(base_in).glob('**/*.np*')

    for path in pathlist:

        np_path = str(path)

        print(np_path)

        yield np_path

def preprocess_all(input_dir, overwrite=False, num_slices=224, crop_size=224, voxel_size=1.5):

    start = time.time()

    scans = os.listdir(input_dir)

    scans.sort()

    errors_map = defaultdict(int)

    base_out = input_dir.rstrip('/') + '_preprocessed/'

    valid_scans = 0

    scans_num = len(list(walk_dicom_dirs(input_dir, base_out, False)))

    for scan_dir_path, out_path in tqdm(walk_dicom_dirs(input_dir, base_out), total=scans_num):

        try:

            out_dir = os.path.dirname(out_path)

            os.makedirs(out_dir, exist_ok=True)

            if overwrite or not os.path.exists(out_dir) or not os.listdir(out_dir):

                preprocessed_scan, scan_rgb_sample = \

                    preprocess(scan_dir_path, errors_map, num_slices, crop_size, voxel_size)

                plt.imshow(scan_rgb_sample)

                plt.savefig(out_path + '.png', bbox_inches='tight')

                np.savez_compressed(out_path + '.npz', data=preprocessed_scan)

            valid_scans += 1

            print('\n++++++++++++++++++++++++\nDiagnostics:')

            print(errors_map.items())

        except FileExistsError as e:

            valid_scans += 1

            print('Exists:', out_path)

        except ValueError as e:

            print('\nERROR!')

            print(e)

    print('Total scans: {}'.format(scans_num))

    print('Valid scans: {}'.format(valid_scans))

    print('Scans with insufficient slices: {}'.format(errors_map['insufficient_slices']))

    print('Scans with bad segmentation: {}'.format(errors_map['bad_seg']))

    print('Scans with small resampled z dimension: {}'.format(errors_map['small_z']))

    print((time.time() - start) / scans_num, 'sec/volume')

def split(positives, negatives, lists_dir, print_dirs=False, split_ratio=0.7):

    positive_paths = []

    negative_paths = []

    for preprocessed_dir, path_list, label in (positives, positive_paths, '1'), (negatives, negative_paths, '0'):

        for root, _, files in os.walk(preprocessed_dir):

            path = root.split(os.sep)

            if print_dirs:

                print((len(path) - 1) * '---', os.path.basename(root))

            for f in files:

                if '.np' in f:

                    path_list.append((root + '/' + f, label))

        print('\n INFO:', 'volumes with label', label, len(path_list))

    train_list = []

    val_list = []

    shuffle(positive_paths)

    split_pos = round(split_ratio * len(positive_paths))

    shuffle(negative_paths)

    split_neg = round(split_ratio * len(negative_paths))

    train_list = positive_paths[:split_pos] + negative_paths[:split_neg]

    val_list = positive_paths[split_pos:] + negative_paths[split_neg:]

    shuffle(train_list)

    shuffle(val_list)

    os.makedirs(lists_dir, exist_ok=True)

    with open(lists_dir + '/val.list', 'w') as val_f:

        for path, label in val_list:

            val_f.write(path + ' ' + label + '\n')

    with open(lists_dir + '/train.list', 'w') as train_f:

        for path, label in train_list:

            train_f.write(path + ' ' + label + '\n')