#!/usr/bin/env python
# coding: utf-8

# In[1]:


import dicom
import numpy as np
import matplotlib.pyplot as plt
get_ipython().run_line_magic('matplotlib', 'inline')

from npgamma import calc_gamma


# In[2]:


dcm_ref = dicom.read_file("data_reference.dcm")
dcm_evl = dicom.read_file("data_evaluation.dcm")


# In[3]:


# The x, y, and z defined here have not been sufficiently verified
# They do not necessarily match either what is within Dicom nor what is within your
# TPS. Please verify these and see if they are what you expect them to be.

# If these functions are incorrect or there is a better choice of dimension definitions 
# please contact me by creating an issue within the github repository:
#   https://github.com/SimonBiggs/npgamma/issues

# If you are able to validate these functions please contact me in the same way.


def load_dose_from_dicom(dcm):
    """Imports the dose in matplotlib format, with the following index mapping:
        i = y
        j = x
        k = z
    
    Therefore when using this function to have the coords match the same order,
    ie. coords_reference = (y, x, z)
    """
    pixels = np.transpose(
        dcm.pixel_array, (1, 2, 0))
    dose = pixels * dcm.DoseGridScaling

    return dose


def load_xyz_from_dicom(dcm):
    """Although this coordinate pull from Dicom works in the scenarios tested
    this is not an official x, y, z pull. It needs further confirmation.
    """
    resolution = np.array(
        dcm.PixelSpacing).astype(float)
    # Does the first index match x? 
    # Haven't tested with differing grid sizes in x and y directions.
    dx = resolution[0]

    # The use of dcm.Columns here is under question
    x = (
        dcm.ImagePositionPatient[0] +
        np.arange(0, dcm.Columns * dx, dx))

    # Does the second index match y? 
    # Haven't tested with differing grid sizes in x and y directions.
    dy = resolution[1]
    
    # The use of dcm.Rows here is under question
    y = (
        dcm.ImagePositionPatient[1] +
        np.arange(0, dcm.Rows * dy, dy))
    
    # Is this correct?
    z = (
        np.array(dcm.GridFrameOffsetVector) +
        dcm.ImagePositionPatient[2])

    return x, y, z


dose_reference = load_dose_from_dicom(dcm_ref)
dose_evaluation = load_dose_from_dicom(dcm_evl)

x_reference, y_reference, z_reference = load_xyz_from_dicom(dcm_ref)
x_evaluation, y_evaluation, z_evaluation = load_xyz_from_dicom(dcm_evl)


# In[4]:


# Input coordinates need to match the same order as the dose grid in 
# index reference order.

coords_reference = (
    y_reference, x_reference, z_reference)

coords_evaluation = (
    y_evaluation, x_evaluation, z_evaluation)


# In[5]:


distance_threshold = 3
distance_step_size = distance_threshold / 10

dose_threshold = 0.03 * np.max(dose_reference)
lower_dose_cutoff = np.max(dose_reference) * 0.2
maximum_test_distance = distance_threshold * 2
max_concurrent_calc_points = 30000000
num_threads = 4


# # Here starts internal function
# 
# With inputs:
# 
#  * `coords_reference`
#  * `dose_reference`
#  * `coords_evaluation`
#  * `dose_evaluation`
#  * `distance_threshold`
#  * `dose_threshold`
#  * `lower_dose_cutoff=0`
#  * `distance_step_size=None` (default is 1/10th of `distance_threshold`)
#  * `maximum_test_distance=np.inf` (default of `np.inf` implies no maximum. Recommend using 2 * `distance_threshold`
#  * `max_concurrent_calc_points=np.inf` (Use this parameter to limit the number of points per calculation set. A value of 30000000 should keep function usage under 2 GB of RAM for a 32 bit python.))
#  * `num_threads=1` (Use this parameter to enable multithreading)

# In[6]:


gamma = calc_gamma(
    coords_reference, dose_reference,
    coords_evaluation, dose_evaluation,
    distance_threshold, dose_threshold,
    lower_dose_cutoff=lower_dose_cutoff, 
    distance_step_size=distance_step_size,
    maximum_test_distance=maximum_test_distance,
    max_concurrent_calc_points=max_concurrent_calc_points,
    num_threads=num_threads)


# In[7]:


valid_gamma = gamma[~np.isnan(gamma)]
valid_gamma[valid_gamma > 2] = 2

plt.hist(valid_gamma, 30);
plt.xlim([0,2])


# In[8]:


np.sum(valid_gamma <= 1) / len(valid_gamma)


# In[9]:


relevant_slice = (
    np.max(dose_evaluation, axis=(0, 1)) > 
    lower_dose_cutoff)
slice_start = np.max([
        np.where(relevant_slice)[0][0], 
        0])
slice_end = np.min([
        np.where(relevant_slice)[0][-1], 
        len(z_evaluation)])


# In[10]:


max_ref_dose = np.max(dose_reference)

cut_off_gamma = gamma.copy()
greater_than_2_ref = (cut_off_gamma > 2) & ~np.isnan(cut_off_gamma)
cut_off_gamma[greater_than_2_ref] = 2

for z_i in z_evaluation[slice_start:slice_end:5]:
    i = np.where(z_i == z_evaluation)[0][0]
    j = np.where(z_i == z_reference)[0][0]
    print("======================================================================")
    print("Slice = {0}".format(z_i))  
   
    plt.contourf(
        x_evaluation, y_evaluation, dose_evaluation[:, :, j], 30, 
        vmin=0, vmax=max_ref_dose, cmap=plt.get_cmap('gist_heat'))
    plt.title("Evaluation")
    plt.colorbar()
    plt.show()
    
    plt.contourf(
        x_reference, y_reference, dose_reference[:, :, j], 30, 
        vmin=0, vmax=max_ref_dose, cmap=plt.get_cmap('gist_heat'))
    plt.title("Reference")  
    plt.colorbar()
    plt.show()
    
    plt.contourf(
        x_evaluation, y_evaluation, cut_off_gamma[:, :, i], 30, 
        vmin=0, vmax=2, cmap=plt.get_cmap('bwr'))
    plt.title("Gamma")    
    plt.colorbar()  
    plt.show()
    
    print("\n")