#********************************************************************************
# Title:   extractDoseProfiles.py
# Author:  John Eley
# Date:    14 June 2015
#
# Purpose: Extract 1D dose profile from dose grid and write to text file
#
# Input:   pcDirectoryAndFileNameOut 
#          sProfile
#              iFlagProfileDirection (1=x, 2=y, 3=z, -1=-x, -2=-y, -3=-z, directions are in DICOM image coordinate system) 
#              dXOrigin (cm, 1D profiles will originate from these specified coordinates, coordinates must given in the DICOM image coordinate frame, hint: when you hover the cursor over the image in RayStation, it will show you the DICOM coordinates as [x z -y])
#              dYOrigin 
#              dZOrigin 
#
# Output: pdPathlength (1D array of pathlength distances, distance from origin)
#         pdDose (1D array of dose corresponding to the above pathlength coordinates)
#       
# Notes: Currently only supports profiles along cardinal directions
#        That is why I am using the terminology "pathlength" as output instead of just e.g., z-coordinates because eventually we want to support diagonal profiles etc
#        The current code has been tested by JE on 6/16/15 and profiles appear nearly identical to those achieved by the RayStation profile tool, 
#            however, the absolute position has not been closely verified and for example the profiles might be shifted by +/-1 voxel due to exact definition of voxel coordinate (is the coordinate at the lower bound, the middle, or the upper borde of a voxel)
#            thus, perfect agreement in position is not expected until this has been closely looked at, shifts of +/- 1 voxel are reasonable in the meantime to find agreement w measured data
#
#********************************************************************************


# Initialize empty list (structure) to hold requested from profiles
sProfile = [] # This must occur before the user input section below


#
# USER INPUT
#


pcDirectoryAndFileNameOut = 'J:/raystationScripts/testExtractDoseProfiles.txt'

# Specify requested profile(s) (can add as few or many profiles as you wish using the append command, just repeat the following line with new parameters), see notes above about directions and coordinates
sProfile.append({'iFlagProfileDirection': +1, 'dXProfileOrigin': 0.37, 'dYProfileOrigin': 13.3, 'dZProfileOrigin': 10.6})
sProfile.append({'iFlagProfileDirection': +2, 'dXProfileOrigin': 0.37, 'dYProfileOrigin': 13.3, 'dZProfileOrigin': 10.6})
sProfile.append({'iFlagProfileDirection': +3, 'dXProfileOrigin': 0.37, 'dYProfileOrigin': 13.3, 'dZProfileOrigin': 10.6})


#
# END USER INPUT
#




# Load necessary objects into memory and import libraries
from connect import *
import math
import string
import json
import time
patient = get_current("Patient") 
plan = get_current("Plan")
beamSet = get_current("BeamSet") 


# Specify or retrieve the plan and beamset for which we will extract dose, for now this is the current active plan, but later may wish to make this variable to for example auto create new plans and keep extracting profiles
pcPlanName = plan.Name 
pcBeamSetName = plan.Name # For now, I assume that beamset name is the plan name (I don't know how else to get it, and this is usually the case at least when you only have 1 beamset)


nProfiles = len(sProfile) # determine how many profiles the user has requested


# Determine the dimensions and coordinates of the existing dose grid for this plan
dXMin = patient.TreatmentPlans[pcPlanName].BeamSets[pcBeamSetName].FractionDose.InDoseGrid.Corner.x
dYMin = patient.TreatmentPlans[pcPlanName].BeamSets[pcBeamSetName].FractionDose.InDoseGrid.Corner.y
dZMin = patient.TreatmentPlans[pcPlanName].BeamSets[pcBeamSetName].FractionDose.InDoseGrid.Corner.z
dDeltaX = patient.TreatmentPlans[pcPlanName].BeamSets[pcBeamSetName].FractionDose.InDoseGrid.VoxelSize.x
dDeltaY = patient.TreatmentPlans[pcPlanName].BeamSets[pcBeamSetName].FractionDose.InDoseGrid.VoxelSize.y
dDeltaZ = patient.TreatmentPlans[pcPlanName].BeamSets[pcBeamSetName].FractionDose.InDoseGrid.VoxelSize.z
nX = patient.TreatmentPlans[pcPlanName].BeamSets[pcBeamSetName].FractionDose.InDoseGrid.NrVoxels.x
nY = patient.TreatmentPlans[pcPlanName].BeamSets[pcBeamSetName].FractionDose.InDoseGrid.NrVoxels.y
nZ = patient.TreatmentPlans[pcPlanName].BeamSets[pcBeamSetName].FractionDose.InDoseGrid.NrVoxels.z
dXMax = dXMin + (nX - 1) * dDeltaX
dYMax = dYMin + (nY - 1) * dDeltaY
dZMax = dZMin + (nZ - 1) * dDeltaZ


# Testing
#print("nX: %d\n" % nX)
#print("nY: %d\n" % nY)
#print("nZ: %d\n" % nZ)



# Initialize lists to hold the profile data, can be indexed by iProfile
sProfilePathlengthData = []
sProfileDoseData = []


for iProfile in range(0, nProfiles):


    # Check that requested origin is withing the range of the dose grid
    if (sProfile[iProfile]['dXProfileOrigin'] < dXMin) or (sProfile[iProfile]['dXProfileOrigin'] > dXMax): 
        print("<E> Requested x-coordinate outside bounds of dose grid. iProfile = %d\n" % iProfile)
    elif (sProfile[iProfile]['dYProfileOrigin'] < dYMin) or (sProfile[iProfile]['dYProfileOrigin'] > dYMax): 
        print("<E> Requested y-coordinate outside bounds of dose grid. iProfile = %d\n" % iProfile)
    elif (sProfile[iProfile]['dZProfileOrigin'] < dZMin) or (sProfile[iProfile]['dZProfileOrigin'] > dZMax): 
        print("<E> Requested z-coordinate outside bounds of dose grid. iProfile = %d\n" % iProfile)


    # Determine index of origin in x, y, z
    iXOrigin = int(math.floor((sProfile[iProfile]['dXProfileOrigin'] - dXMin)/dDeltaX))
    iYOrigin = int(math.floor((sProfile[iProfile]['dYProfileOrigin'] - dYMin)/dDeltaY))
    iZOrigin = int(math.floor((sProfile[iProfile]['dZProfileOrigin'] - dZMin)/dDeltaZ))


    # Initialize some things
    iL = 0 # counter for path (profile)
    iX = 0
    iY = 0
    iZ = 0
    iXStart = 0
    iYStart = 0
    iZStart = 0




    #
    # Extract profiles and store in 1D path arrays
    #


    if sProfile[iProfile]['iFlagProfileDirection'] == 1:


        sProfile[iProfile]['pcPathDirection'] = '+x'

        # Find starting index for x
        iXStart = int(math.floor((sProfile[iProfile]['dXProfileOrigin'] - dXMin)/dDeltaX))

        # Initialize 1D arrays to store dose profile
        pdPathLength = [0] * (nX - iXStart) 
        pdDose = [0] * (nX - iXStart)
    
        for iX in range(iXStart,nX):
            pdPathLength[iL] = iL * dDeltaX 
            pdDose[iL] = patient.TreatmentPlans[pcPlanName].BeamSets[pcBeamSetName].FractionDose.DoseValues.DoseData[iZOrigin, iYOrigin, iX] 
            iL = iL + 1
        sProfile[iProfile]['nL'] = iL - 1 # grab the highest index after the loop ends

        # Store profile in multi-profile structures
        sProfilePathlengthData.append(pdPathLength)
        sProfileDoseData.append(pdDose)



    elif sProfile[iProfile]['iFlagProfileDirection'] == 2:


        sProfile[iProfile]['pcPathDirection'] = '+y'

        # Find starting index for y
        iYStart = int(math.floor((sProfile[iProfile]['dYProfileOrigin'] - dYMin)/dDeltaY))

        # Initialize 1D arrays to store dose profile
        pdPathLength = [0] * (nY - iYStart) 
        pdDose = [0] * (nY - iYStart)
    
        for iY in range(iYStart,nY):
            pdPathLength[iL] = iL * dDeltaY 
            pdDose[iL] = patient.TreatmentPlans[pcPlanName].BeamSets[pcBeamSetName].FractionDose.DoseValues.DoseData[iZOrigin, iY, iXOrigin] 
            iL = iL + 1
        sProfile[iProfile]['nL'] = iL - 1 # grab the highest index after the loop ends

        # Store profile in multi-profile structures
        sProfilePathlengthData.append(pdPathLength)
        sProfileDoseData.append(pdDose)


    elif sProfile[iProfile]['iFlagProfileDirection'] == 3:


        sProfile[iProfile]['pcPathDirection'] = '+z'

        # Find starting index for z
        iZStart = int(math.floor((sProfile[iProfile]['dZProfileOrigin'] - dZMin)/dDeltaZ))

        # Initialize 1D arrays to store dose profile
        pdPathLength = [0] * (nZ - iZStart) 
        pdDose = [0] * (nZ - iZStart)
    
        for iZ in range(iZStart,nZ):
            pdPathLength[iL] = iL * dDeltaZ 
            pdDose[iL] = patient.TreatmentPlans[pcPlanName].BeamSets[pcBeamSetName].FractionDose.DoseValues.DoseData[iZ, iYOrigin, iXOrigin] 
            iL = iL + 1
        sProfile[iProfile]['nL'] = iL - 1 # grab the highest index after the loop ends

        # Store profile in multi-profile structures
        sProfilePathlengthData.append(pdPathLength)
        sProfileDoseData.append(pdDose)


    elif sProfile[iProfile]['iFlagProfileDirection'] == -1:


        sProfile[iProfile]['pcPathDirection'] = '-x'

        # Find starting index for x
        iXStart = int(math.floor((sProfile[iProfile]['dXProfileOrigin'] - dXMin)/dDeltaX))

        # Initialize 1D arrays to store dose profile
        pdPathLength = [0] * (iXStart) # Since we reverse the profile, the length will be exactly iXStart
        pdDose = [0] * (iXStart)
    
        for iX in range(0,(iXStart)):
            iXReverseLookup = iXStart - iX
            pdPathLength[iL] = iL * dDeltaX 
            pdDose[iL] = patient.TreatmentPlans[pcPlanName].BeamSets[pcBeamSetName].FractionDose.DoseValues.DoseData[iZOrigin, iYOrigin, iXReverseLookup] 
            iL = iL + 1
        sProfile[iProfile]['nL'] = iL - 1 # grab the highest index after the loop ends

        # Store profile in multi-profile structures
        sProfilePathlengthData.append(pdPathLength)
        sProfileDoseData.append(pdDose)


    elif sProfile[iProfile]['iFlagProfileDirection'] == -2:


        sProfile[iProfile]['pcPathDirection'] = '-y'

        # Find starting index for y
        iYStart = int(math.floor((sProfile[iProfile]['dYProfileOrigin'] - dYMin)/dDeltaY))

        # Initialize 1D arrays to store dose profile
        pdPathLength = [0] * (iYStart) # Since we reverse the profile, the length will be exactly iYStart
        pdDose = [0] * (iYStart)
    
        for iY in range(0,(iYStart)):
            iYReverseLookup = iYStart - iY
            pdPathLength[iL] = iL * dDeltaY 
            pdDose[iL] = patient.TreatmentPlans[pcPlanName].BeamSets[pcBeamSetName].FractionDose.DoseValues.DoseData[iZOrigin, iYReverseLookup, iXOrigin] 
            iL = iL + 1
        sProfile[iProfile]['nL'] = iL - 1 # grab the highest index after the loop ends

        # Store profile in multi-profile structures
        sProfilePathlengthData.append(pdPathLength)
        sProfileDoseData.append(pdDose)


    elif sProfile[iProfile]['iFlagProfileDirection'] == -3:


        sProfile[iProfile]['pcPathDirection'] = '-z'

        # Find starting index for z
        iZStart = int(math.floor((sProfile[iProfile]['dZProfileOrigin'] - dZMin)/dDeltaZ))

        # Initialize 1D arrays to store dose profile
        pdPathLength = [0] * (iZStart) # Since we reverse the profile, the length will be exactly iZStart
        pdDose = [0] * (iZStart)
    
        for iZ in range(0,(iZStart)):
            iZReverseLookup = iZStart - iZ
            pdPathLength[iL] = iL * dDeltaZ 
            pdDose[iL] = patient.TreatmentPlans[pcPlanName].BeamSets[pcBeamSetName].FractionDose.DoseValues.DoseData[iZReverseLookup, iYOrigin, iXOrigin] 
            iL = iL + 1
        sProfile[iProfile]['nL'] = iL - 1 # grab the highest index after the loop ends

        # Store profile in multi-profile structures
        sProfilePathlengthData.append(pdPathLength)
        sProfileDoseData.append(pdDose)


    else:
        print("<E> iFlagProfileDirection outside valid range.")




# Write data to file
fileOut = open(pcDirectoryAndFileNameOut,'w')
fileOut.write('1D Dose Profile(s)\n')
fileOut.write('Created by - extractDoseProfiles.py - Eley\n')
fileOut.write('%s\n' % time.strftime("%c"))

for iProfile in range(0, nProfiles):
    fileOut.write(' \n')
    fileOut.write('iProfile: %d\n' % iProfile)
    fileOut.write('Path Direction: %s\n' % sProfile[iProfile]['pcPathDirection'])
    fileOut.write('Path Origin [x y z] (cm): %.3f %.3f %.3f\n' % (sProfile[iProfile]['dXProfileOrigin'], sProfile[iProfile]['dYProfileOrigin'], sProfile[iProfile]['dZProfileOrigin']))
    fileOut.write('Pathlength (cm) - Dose (cGy) \n')
    for iL in range(0, sProfile[iProfile]['nL']):
        fileOut.write('%.3f %.3f\n' % (sProfilePathlengthData[iProfile][iL], sProfileDoseData[iProfile][iL]))


fileOut.close()