import matplotlib.pyplot as plt
import numpy as np
import scipy
import scipy.ndimage
import data_transforms
import pathfinder
import utils
import utils_lung
import logger
import sys
import collections


def resample(image, spacing, new_spacing=[1, 1, 1]):
    # Determine current pixel spacing
    spacing = np.array(spacing)
    resize_factor = spacing / new_spacing
    new_real_shape = image.shape * resize_factor
    new_shape = np.round(new_real_shape)
    real_resize_factor = new_shape / image.shape
    new_spacing = spacing / real_resize_factor
    image = scipy.ndimage.interpolation.zoom(image, real_resize_factor)
    return image, new_spacing


def plot_2d(image3d, axis, pid, img_dir):
    fig = plt.figure()
    fig.canvas.set_window_title(pid)
    ax = fig.add_subplot(111)
    idx = image3d.shape[axis] / 2
    if axis == 0:  # sax
        ax.imshow(image3d[idx, :, :], cmap=plt.cm.gray)
    if axis == 1:  # 2 lungs
        ax.imshow(image3d[:, idx, :], cmap=plt.cm.gray)
    if axis == 2:  # side view
        ax.imshow(image3d[:, :, idx], cmap=plt.cm.gray)
    plt.show()
    fig.savefig(img_dir + '/%s.png' % pid, bbox_inches='tight')
    fig.clf()
    plt.close('all')


def test1():
    image_dir = utils.get_dir_path('analysis', pathfinder.METADATA_PATH)
    image_dir = image_dir + '/test_1/'
    utils.auto_make_dir(image_dir)

    sys.stdout = logger.Logger(image_dir + '/%s.log' % 'test1_log')
    sys.stderr = sys.stdout

    patient_data_paths = utils_lung.get_patient_data_paths(pathfinder.DATA_PATH)
    print len(patient_data_paths)

    for k, p in enumerate(patient_data_paths):
        pid = utils_lung.extract_pid_dir(p)
        try:
            sid2data, sid2metadata = utils_lung.get_patient_data(p)
            sids_sorted = utils_lung.sort_sids_by_position(sid2metadata)
            sids_sorted_jonas = utils_lung.sort_slices_jonas(sid2metadata)
            sid2position = utils_lung.slice_location_finder(sid2metadata)

            try:
                slice_thickness_pos = np.abs(sid2metadata[sids_sorted[0]]['ImagePositionPatient'][2] -
                                             sid2metadata[sids_sorted[1]]['ImagePositionPatient'][2])
            except:
                print 'This patient has no ImagePosition!'
                slice_thickness_pos = 0.
            try:
                slice_thickness_loc = np.abs(
                    sid2metadata[sids_sorted[0]]['SliceLocation'] - sid2metadata[sids_sorted[1]]['SliceLocation'])
            except:
                print 'This patient has no SliceLocation!'
                slice_thickness_loc = 0.

            jonas_slicethick = []
            for i in xrange(len(sids_sorted_jonas) - 1):
                s = np.abs(sid2position[sids_sorted_jonas[i + 1]] - sid2position[sids_sorted_jonas[i]])
                jonas_slicethick.append(s)

            full_img = np.stack([data_transforms.ct2normHU(sid2data[sid], sid2metadata[sid]) for sid in sids_sorted])
            del sid2data, sid2metadata
            print np.min(full_img), np.max(full_img)
            # spacing = sid2metadata[sids_sorted[0]]['PixelSpacing']
            # spacing = [slice_thickness, spacing[0], spacing[1]]
            # resampled_image, _ = resample(full_img, spacing)
            plot_2d(full_img, axis=0, pid=pid + 'ax0', img_dir=image_dir)
            plot_2d(full_img, axis=1, pid=pid + 'ax1', img_dir=image_dir)
            plot_2d(full_img, axis=2, pid=pid + 'ax2', img_dir=image_dir)
            print k, pid, full_img.shape, slice_thickness_pos, slice_thickness_loc, set(jonas_slicethick)
            del full_img
        except:
            print 'exception!!!', pid


def test2():
    patient_data_paths = utils_lung.get_patient_data_paths(pathfinder.DATA_PATH)
    print len(patient_data_paths)
    pixel_spacings_xy = []
    n_slices = []

    for k, p in enumerate(patient_data_paths):
        pid = utils_lung.extract_pid_dir(p)
        sid2data, sid2metadata = utils_lung.get_patient_data(p)
        mtd = sid2metadata.itervalues().next()

        assert mtd['PixelSpacing'][0] == mtd['PixelSpacing'][1]
        pixel_spacings_xy.append(mtd['PixelSpacing'][0])
        n_slices.append(len(sid2metadata))
        print pid, pixel_spacings_xy[-1], n_slices[-1]

    print 'nslices', np.max(n_slices), np.min(n_slices), np.mean(n_slices)
    counts = collections.Counter(pixel_spacings_xy)
    new_list = sorted(pixel_spacings_xy, key=counts.get, reverse=True)
    print 'spacing', new_list


if __name__ == '__main__':
    test1()