#
%load_ext autoreload
%autoreload 2
%matplotlib inline
#
import numpy
import pylab as plt
import scipy
import math
#import PIL
import os
import time
import h5py
import glob
#
import sys
import os
#
# PyDICOM dependencies:
# Ran into some problems getting some of these libraries to load in Jupyter. Dunno... just working around it.
if not '/Users/myoder96/anaconda3/envs/dicom/lib/python3.13/site-packages' in sys.path:
    sys.path += ['/Users/myoder96/anaconda3/envs/dicom/lib/python3.13/site-packages']
import gdcm
#import Pillow
import PIL
import pylibjpeg

import pydicom

# Write a simple images viewer:
#
def view_images(data_file, dataset=None, n_cols=4, img_size=2., palette=None, cmap=None):
    #
    ext = os.path.splitext(data_file)[1]
    #
    #print('*** ext: ', ext)
    if ext.lower().endswith('h5'):
        with h5py.File(data_file, 'r') as fin:
            data = fin[dataset][:]
    else:
        with pydicom.dcmread(data_file) as fin:
            data = fin.pixel_array[:]
    #
    #
    Nz, Ny, Nx = data.shape
    n_rows = int(numpy.ceil(Nz/n_cols))
    #
    while (n_cols * n_rows) < N:
        n_rows += 1
    #
    ax_side = img_size
    fg = plt.figure(figsize=(ax_side*n_cols, ax_side*n_rows))
    fg.clear()
    #
    for k,X in enumerate(data):
        # if k>N:
        #     break
        #
        #print(f'** ary[{k}]')
        ax = fg.add_subplot(n_rows, n_cols, k+1)
        if not palette is None:
            #print('*** executing with pallette...')
            ax.imshow(pydicom.pixels.apply_color_lut(X, palette=palette), cmap=cmap)
        else:
            #print(f'** executing as if None... [{pallette}], { (pallette is None) }')
            ax.imshow(X, cmap=cmap)
        #
        ax.set_title(f'Image: {k}')
        #
    plt.show()       
#
# Use is like:
# test_file = '0003.DCM'
# view_images(os.path.join(dicom_data_path, test_file), palette=None, cmap=plt.cm.afmhot)

#ls -lh dicom_data
# Also see built in examples:
examples_path = pydicom.examples.get_path("mr")  # The path to the examples.mr dataset
print(f'** examples_path: {examples_path}')
#
# and let's have a look at the array shapes, to see if they might be good examples for ML workflows
#  or other consolidation applications:
for k,fl in enumerate(glob.glob(f'{os.path.split(examples_path)[0]}/*.dcm')):
    with pydicom.dcmread(fl, force=True) as dcf:
        #print(f'** {fl}') 
        #
        try:
            ary = dcf.pixel_array
            print(f'** fl[{k}] {os.path.split(fl)[1]}: {numpy.shape(ary)}')
        except:
            ary = numpy.empty(shape=(0,0,0))
            print('** broken. skipping.')
        #
        
    #
    

dicom_data_path = './dicom_data'
#dicom_files = list(glob.glob(os.path.join(dicom_data_path, '*.DCM')))
dicom_files = [s  for ex in ['dcm', 'DCM'] for s in glob.glob(os.path.join(dicom_data_path, f'*.{ex}'))]
#
for fl in dicom_files:
    print(f'** {fl}')

#dc_1 = dicom.dcmread(dicom_files[0])
#
for k, fl in enumerate(dicom_files):
    with pydicom.dcmread(fl) as dcf:
        ary = dcf.pixel_array
        print(f'** shape[k] ({os.path.split(fl)[1]}): {numpy.shape(ary)}')

print(f'** shape: {numpy.shape(ary)}')

test_file = '0003.DCM'
zz = view_images(os.path.join(dicom_data_path, test_file), palette=None, cmap=plt.cm.afmhot)

working_data_h5 = os.path.join(dicom_data_path, 'working_data.h5')
#
# NOTE: we double up some of the data sets to illustrate appending data sets.
training_data_files  = [os.path.join(dicom_data_path, s) for s in ['0002.DCM', '0002.DCM']]
testing_data_files   = [os.path.join(dicom_data_path, s) for s in ['0003.DCM', '0004.DCM', '0003.DCM']]
inference_data_files = [os.path.join(dicom_data_path, s) for s in ['0009.DCM', '0009.DCM']]
#
# NOTE: having some issues making numpy.append() work well, so getting correct array shape for multiple inputs is
# messy, so do not assume this is the best way to do this part.
#
with h5py.File(working_data_h5, 'w') as h5f:
    #
    # There are lots of good and bad ways to import and write the data. This step -- consolidating, orgainizing,
    # and pre-processing data is not always pretty. Often, it is more important that this process be flexible
    # and transparent than highly optimized, so these steps might not always be your proudest code.
    #
    ########################
    # training data:
    # there are lots of ways to work through the data. This approach makes two passes through the data,
    # so that the array size and shape can be determined in the first pass, then the array can be created
    # to it's full size in one step, which can save memory and time.
    #
    Nz = 0
    for n,fl in enumerate(training_data_files):
        with pydicom.dcmread(fl) as dcf:
            Nz += dcf.pixel_array.shape[0]
            if n == 0:
                # NOTE: you can also use this step to confirm that all the sub-arrays have the same shape.
                ny, nx = dcf.pixel_array.shape[1:]
    # Now, spin through it again and assign data:
    k=0
    h5f.create_dataset('training', (Nz, ny, nx) )
    for fl in training_data_files:
        with pydicom.dcmread(fl) as dcf:
            n_arrays = dcf.pixel_array.shape[0]
            h5f['training'][k:k+n_arrays] = dcf.pixel_array[:]
            k += n_arrays
            #
    #
    ####################
    # This approach uses numpy.append(). numpy.append() appends values to a copy of the
    #  array, which can be expensive in terms of memory write time and also can result in 
    #  memory use spikes.
    #
    # testing:
    data = None
    for fl in testing_data_files:
        with pydicom.dcmread(fl) as dcf:
            if data is None:
                data = dcf.pixel_array[:]
                sh = data.shape
            else:
                # NOTE: Specify the append axis= to preserve the shape. If this is excluded,
                #  numpy.append() will return a shape=(1,) array (ie, numpy.ravel(data) ).
                data = numpy.append(data, dcf.pixel_array, axis=0)
            #
    h5f.create_dataset('testing', data=data)
    #
    ########################
    # inference:
    data = None
    #
    # here's another way to assign data. Also, having trouble with append...
    for fl in inference_data_files:
        with pydicom.dcmread(fl) as dcf:
            if data is None:
                data = dcf.pixel_array[:]
                sh = data.shape
            else:
                data = numpy.append(data, dcf.pixel_array[:], axis=0)
                #data.shape=numpy.array([sh[0] + dcf.pixel_array.shape[0], sh[1:]])
                #
                #pass
            #
    h5f.create_dataset('inference', data=data)
    #
    # create a group out of principle, to show we can (and some other stuff later)
    h5f.create_group('some_group')
    
    
#########
# have a look at what we created:
with h5py.File(working_data_h5, 'r') as h5f:
    print('** ', h5f.keys())
    #
    for ds in h5f.keys():
        #print(f'** {type(h5f[ds])}')
        is_DS = isinstance(h5f[ds], h5py.Dataset)
        print(f'** [{k}] {ds}: {(f"(DS): {h5f[ds].shape}" if is_DS else "(GR)")}')
        #
        # if is_DS:
        #     print(f'** SH: {h5f[ds].shape}')

# Plot some images:
view_images('dicom_data/working_data.h5', dataset='training', n_cols=6)

from pydicom import examples
from pydicom.pixels import apply_color_lut

ds = examples.palette_color
arr = ds.pixel_array
# You could also use the corresponding well-known SOP Instance UID
rgb = apply_color_lut(arr, palette='PET')
#
fg = plt.figure(figsize=(14,6))
ax1 = fg.add_subplot(1,2,1)
ax2 = fg.add_subplot(1,2,2)
ax1.imshow(arr)
ax2.imshow(rgb)
plt.show()