%%
% This file has been modified by Surendra Prajapati especially to run
% WiscPlan for KV beams. Works for running locally as well as in Server.
%
% RTF 11/4/05
% Sets up the run files needed for running a Condor convolution
% superposition dose calculation with photon beams.

% Versions of the convolution code that run in both Windows XP, Condor, and
% directly through Matlab can all be created here.  The only real
% differences between the files are in the access conventions, which just%
% This has to do with switching forward slashes with backslashes.

% Surendra edits: num_batches, SAD, pitch, xpmin/max, ypmin/max, Mxp, Nphi 
% Also edit: 
% Kernel should always be named Kernels.mat 
% SAD < 20 was edited from SAD < 50 for clinical system (error message)
% 

% specify where kernel and geometry files are located
patient_dir = 'C:\localGit\wiscPlanv2\WiscPlanPhotonkV125\PatientData';
num_batches = 2; % number of cores you want to run the beam calculation

iso = [0 0 0]; % Point about which gantry rotations begin
SAD = 35;  % source-axis distance for the x-ray source
pitch = 0.86; % fraction of beam with couch translates per rotation

% define the overall beam field size for each beam angle
xpmin = -0.5;     % -field width / 2
xpmax = 0.5;      % +field width / 2
ypmin = -0.5;   % -jaw width / 2
ypmax = 0.5;    % +jaw width / 2
% y-prime points in the z-direction in the CT coordinate system

% Number of beamlets in the BEV for each direction
Mxp = 20;        % number of leaves; leaf size is 1/20cm= 0.5mm
Nyp = 1;        % always 1 for Tomo due to binary mlc

% ============================================= End of user-supplied inputs
executable_path = 'C:\localGit\wiscPlanv2\WiscPlanPhotonkV125\WiscPlanEXE\RyanCsphoton.x86.exe';
kernel_file = 'Kernels.mat';
geometry_file = fullfile(patient_dir, 'matlab_files\Geometry.mat');
% tumor bow and stern locations
load(geometry_file);
ROI_names = cellfun(@(c)c.name, Geometry.ROIS, 'UniformOutput', false);
[target_idx okay] = listdlg('ListString', ROI_names, ...
    'SelectionMode', 'single', 'Name', 'Target Selection', ...
    'PromptString', 'Please select the target ROI. ');
if okay ~= 1
    msgbox('Plan creation aborted');
    return;
end
targetMask = zeros(size(Geometry.data));
targetMask(Geometry.ROIS{target_idx}.ind) = 1;
targetMaskZ = sum(sum(targetMask,1),2);
zBow = find(targetMaskZ>0, 1, 'first')*Geometry.voxel_size(3) + Geometry.start(3) + ypmin;
zStern = find(targetMaskZ>0, 1, 'last')*Geometry.voxel_size(3) + Geometry.start(3) + ypmax;
[subi, subj, subk] = ind2sub(size(Geometry.data), Geometry.ROIS{target_idx}.ind);
iso = [Geometry.start(1)+Geometry.voxel_size(1)*mean(subi) ...
    Geometry.start(2)+Geometry.voxel_size(2)*mean(subj) 0];

% flags used to select which calculations will be set up
Condor_flag = 1;

ptvInd = target_idx;  % PTV index in Geometry.ROIS

fieldWidth = ypmax - ypmin;

% total number of rotations required for treatment
Nrot = ceil(abs(zBow - zStern)/(pitch*fieldWidth));

Nphi = Nrot*51;  % number of angles used in the calculation

load(geometry_file);

% define the limits of the angles that will be used for the calculation
phimin = 0;  % starting angle in radians
phimax = 2*pi*Nphi;

phi = [0:Nphi-1]/Nphi*2*pi*Nrot;

condor_folder = patient_dir;
winxp_folder = 'winxp';

% create names for condor input and output folders
input_folder = '.';
output_folder = '.';

% name of the convolution/superposition executable, which will be in the
% 'code' folder of each respective run type folder
condor_exectuable_name = 'convolutionCondor';  % relative path on the cluster where code will be
winxp_executable_name = 'convolution.exe';
matlab_executable_name = 'convolution_mex';  % name of the Matlab version of the dose calculation code

% set the beam parameters, assuming a helical beam trajectory
% folders that will be inside the 'input' folder
beamspec_folder = 'beamspecfiles';    % directory where beam files will be stored
beamspec_batches_folder = 'beamspecbatches';
beamspec_batch_base_name = 'beamspecbatch';  % base name for a beamlet batch file
kernel_folder = 'kernelfiles';  % folder where kernel information will be saved
kernel_filenames_condor = 'kernelFilenamesCondor.txt';
kernel_filenames_winxp = 'kernelFilenamesWinXP.txt';

% output folders
beamlet_batch_base_name = 'beamletbatch';  % base name for a dose batch file

geometry_header_filename = 'geometryHeader.txt';
geometry_density_filename = 'density.bin';  % save the density, not the Hounsfield units!

% end of user-defined parameters

% check the validity of the user-defined variables
if xpmin >= xpmax
    error('xpmin must be less than xpmax.');
end

if ypmin >= ypmax
    error('ypmin must be less than ypmax.');b
end

if phimin > phimax
    error('phimin must be less than or equal to phimax.');
end

if Mxp <= 0 || Nyp <= 0 || Nphi <= 0
    error('Mxp, Nyp, and Nphi must be greater than zero.');
end

if SAD < 20
    error('It is recommended that the SAD be greater than 20 cm.');
end

% the xy plane is perpendicular to the isocenter axis of the linac gantry

% size of each beam aperture, making them vectors so extension to
% non-uniform aperture sizes becomes obvious
del_xp = (xpmax - xpmin)/Mxp;
del_yp = (ypmax - ypmin)/Nyp;

% Calculate the xp and yp offsets, which lie at the centers of the
% apertures.
xp = [xpmin:del_xp:xpmax-del_xp] + del_xp/2;
yp = [ypmin:del_yp:ypmax-del_yp] + del_yp/2;

[M,N,Q] = size(Geometry.rhomw);

START = single(Geometry.start - iso);
INC = single(Geometry.voxel_size);

% define the tumor mask
tumorMask = zeros(size(Geometry.rhomw),'single');
tumorMask(Geometry.ROIS{ptvInd}.ind) = 1;

BW = bwdist(tumorMask);
tumorMaskExp = tumorMask;
tumorMaskExp(BW <= 4) = 1;

P = zeros(Mxp,Nphi);

fprintf('Checking beam''s eye view ...\n');
for p=1:Nphi
    % ir and jr form the beam's eye view (BEV)
    ir = [-sin(phi(p)); cos(phi(p)); 0];
    jr = [0 0 1]';
    % kr denotes the beam direction
    kr = [cos(phi(p)); sin(phi(p)); 0];
    
    for m=1:Mxp
        point1 = single(-kr*SAD + [0 0 zBow + pitch*fieldWidth*phi(p)/(2*pi)]');  % source point
        point2 = single(point1 + (SAD*kr + ir*xp(m))*10);
        [indVisited,deffVisited] = singleRaytraceClean(tumorMaskExp,START,INC,point1,point2);
        if ~isempty(indVisited)
            P(m,p) = max(deffVisited);
        end
    end
end
fprintf('Finished checking BEV\n');

% load data required for the dose calculator
load(kernel_file);

Geometry.rhomw(Geometry.rhomw < 0) = 0;
Geometry.rhomw(Geometry.rhomw < 0.0013) = 0.0013;  % fill blank voxels with air

% convert Geometry and kernels to single
f = fieldnames(Kernels);
for k=1:length(f)
    if isnumeric(getfield(Kernels,f{k}))
        Kernels = setfield(Kernels,f{k},single(getfield(Kernels,f{k})));    
    end
end

f = fieldnames(Geometry);
for k=1:length(f)
    if isnumeric(getfield(Geometry,f{k}))
        Geometry = setfield(Geometry,f{k},single(getfield(Geometry,f{k})));    
    end
end

% account for isocenter
Geometry.start = single(Geometry.start - iso);

% find the total number of beams
Nbeam = Nphi*Mxp*Nyp;

batch_num = 0;  % start the count for the number of total batches

% fill up a cell array of beam structures, grouped by batch
clear batches;
batch_num = 0;

batches = cell(1,Nrot);  % start the batches cell array (cell array of beam batches)
rotNum = 0;

% calculate beams for all source directions and apertures
for k=1:Nphi  % loop through all gantry angles
    % calculate the source location for a helical trajectory
    beam.SAD = single(SAD);
    
    % the kp vector is the beam direction, ip and jp span the beam's eye view
    beam.ip = single([-sin(phi(k)) cos(phi(k)) 0]);
    beam.jp = single([0 0 1]);
    beam.kp = single([cos(phi(k)) sin(phi(k)) 0]);
    beam.y_vec = single(-beam.kp*SAD + [0 0 zBow + pitch*fieldWidth*phi(k)/(2*pi)]);
    
    rotNumOld = rotNum;
    rotNum = floor(k/51) + 1;  % current rotation number
    if rotNum - rotNumOld > 0
        beam_num = 0; % if the rotation number has changed, start the beam count over
    end
    
    for m=1:Mxp  % loop through all apertures in the xp-direction
        % calculate the beam if the tomotherapy fluence value is non-zero
        if P(m,k) > 0
            num = m + (k-1)*Mxp - 1;         % beamlet number (overall)
            beam_num = beam_num + 1;

            % set the beam aperture parameters
            beam.del_xp = single(del_xp);
            beam.del_yp = single(del_yp);
            beam.xp = single(xp(m));
            beam.yp = single(0);
            beam.num = single(num);  % record the beam number to avoid any later ambiguity

            batches{rotNum}{beam_num} = beam;
        end
    end
end
% merge/split batches
all_beams = horzcat(batches{:});
num_beams_per_batch = ceil(numel(all_beams)/num_batches);
batches = cell(num_batches,1);
for k = 1:(num_batches-1)
    batches{k} = all_beams(1+num_beams_per_batch*(k-1):num_beams_per_batch*k);
end
batches{num_batches} = all_beams(1+num_beams_per_batch*(k):end);


% Everything else in this file is related to saving the batches in a
% useable form.
if Condor_flag == 1
    % delete the old submission file
    err = rmdir(fullfile(condor_folder,beamspec_batches_folder),'s');
    err = rmdir(fullfile(condor_folder,kernel_folder),'s');

    % create folders where batch information will be sent
    mkdir([condor_folder '/' input_folder '/' beamspec_batches_folder]);

    % save the kernels
    save_kernels(Kernels,[condor_folder '/' input_folder '/' kernel_folder]);
    fprintf(['Successfully saved Condor kernels to ' input_folder '/' kernel_folder '\n']);

    % create kernel filenames files
    kernel_filenames_CHTC = 'kernelFilenamesCHTC.txt';
    kernel_filenames_condor = 'kernelFilenamesCondor.txt';
    fid = fopen([condor_folder '/' input_folder '/' kernel_filenames_condor],'w');
    fid2 = fopen([condor_folder '/' input_folder '/' kernel_filenames_CHTC],'w');
    
    fprintf(fid,'kernel_header\n');
    % fprintf(fid,['./' input_folder '/' kernel_folder '/kernel_header.txt\n']);
    fprintf(fid,'%s\n',fullfile(patient_dir,input_folder,kernel_folder,'kernel_header.txt'));
    fprintf(fid2,'kernel_header\n');
    fprintf(fid2, '%s/%s\n', kernel_folder,'kernel_header.txt');
    
    fprintf(fid,'kernel_radii\n');
    % fprintf(fid,['./' input_folder '/' kernel_folder '/radii.bin\n']);
    fprintf(fid,'%s\n',fullfile(patient_dir,input_folder,kernel_folder,'radii.bin'));
    fprintf(fid2,'kernel_radii\n');
    fprintf(fid2, '%s/%s\n', kernel_folder,'radii.bin');
    
    fprintf(fid,'kernel_angles\n');
    % fprintf(fid,['./' input_folder '/' kernel_folder '/angles.bin\n']);
    fprintf(fid,'%s\n',fullfile(patient_dir,input_folder,kernel_folder,'angles.bin'));
    fprintf(fid2,'kernel_angles\n');
    fprintf(fid2, '%s/%s\n', kernel_folder,'angles.bin');
      
    fprintf(fid,'kernel_energies\n');
    % fprintf(fid,['./' input_folder '/' kernel_folder '/energies.bin\n']);
    fprintf(fid,'%s\n',fullfile(patient_dir,input_folder,kernel_folder,'energies.bin'));
    fprintf(fid2,'kernel_energies\n');
    fprintf(fid2, '%s/%s\n', kernel_folder,'energies.bin');
    
    fprintf(fid,'kernel_primary\n');
    % fprintf(fid,['./' input_folder '/' kernel_folder '/primary.bin\n']);
    fprintf(fid,'%s\n',fullfile(patient_dir,input_folder,kernel_folder,'primary.bin'));
    fprintf(fid2,'kernel_primary\n');
    fprintf(fid2, '%s/%s\n', kernel_folder,'primary.bin');  
    
    fprintf(fid,'kernel_first_scatter\n');
    % fprintf(fid,['./' input_folder '/' kernel_folder '/first_scatter.bin\n']);
    fprintf(fid,'%s\n',fullfile(patient_dir,input_folder,kernel_folder,'first_scatter.bin'));
    fprintf(fid2,'kernel_first_scatter\n');
    fprintf(fid2, '%s/%s\n', kernel_folder,'first_scatter.bin');
    
    fprintf(fid,'kernel_second_scatter\n');
    % fprintf(fid,['./' input_folder '/' kernel_folder '/second_scatter.bin\n']);
    fprintf(fid,'%s\n',fullfile(patient_dir,input_folder,kernel_folder,'second_scatter.bin'));
    fprintf(fid2,'kernel_second_scatter\n');
    fprintf(fid2, '%s/%s\n', kernel_folder,'second_scatter.bin');
    
    fprintf(fid,'kernel_multiple_scatter\n');
    % fprintf(fid,['./' input_folder '/' kernel_folder '/multiple_scatter.bin\n']);
    fprintf(fid,'%s\n',fullfile(patient_dir,input_folder,kernel_folder,'multiple_scatter.bin'));
    fprintf(fid2,'kernel_multiple_scatter\n');
    fprintf(fid2, '%s/%s\n', kernel_folder,'multiple_scatter.bin');
    
    fprintf(fid,'kernel_brem_annih\n');
    % fprintf(fid,['./' input_folder '/' kernel_folder '/brem_annih.bin\n']);
    fprintf(fid,'%s\n',fullfile(patient_dir,input_folder,kernel_folder,'brem_annih.bin'));
    fprintf(fid2,'kernel_brem_annih\n');
    fprintf(fid2, '%s/%s\n', kernel_folder,'brem_annih.bin');
    
    fprintf(fid,'kernel_total\n');
    % fprintf(fid,['./' input_folder '/' kernel_folder '/total.bin\n']);
    fprintf(fid,'%s\n',fullfile(patient_dir,input_folder,kernel_folder,'total.bin'));
    fprintf(fid2,'kernel_total\n');
    fprintf(fid2, '%s/%s\n', kernel_folder,'total.bin');
    
    fprintf(fid,'kernel_fluence\n');
    % fprintf(fid,['./' input_folder '/' kernel_folder '/fluence.bin\n']);
    fprintf(fid,'%s\n',fullfile(patient_dir,input_folder,kernel_folder,'fluence.bin'));
    fprintf(fid2,'kernel_fluence\n');
    fprintf(fid2, '%s/%s\n', kernel_folder,'fluence.bin');
    
    fprintf(fid,'kernel_mu\n');
    % fprintf(fid,['./' input_folder '/' kernel_folder '/mu.bin\n']);
    fprintf(fid,'%s\n',fullfile(patient_dir,input_folder,kernel_folder,'mu.bin'));
    fprintf(fid2,'kernel_mu\n');
    fprintf(fid2, '%s/%s\n', kernel_folder,'mu.bin');
    
    fprintf(fid,'kernel_mu_en\n');
    % fprintf(fid,['./' input_folder '/' kernel_folder '/mu_en.bin\n']);
    fprintf(fid,'%s\n',fullfile(patient_dir,input_folder,kernel_folder,'mu_en.bin'));
    fprintf(fid2,'kernel_mu_en\n');
    fprintf(fid2, '%s/%s\n', kernel_folder,'mu_en.bin');
    
    fclose(fid);
end

% name for the condor submit file that will be used
condor_submit_file = 'convolutionSubmit.txt';

geometry_filenames_condor = 'geometryFilenamesCondor.txt';
geometry_filenames_CHTC = 'geometryFilenamesCHTC.txt';

% check the geometry file to ensure that it's not in Hounsfield units
if length(find(Geometry.rhomw > 20)) || length(find(Geometry.rhomw < 0))
    error('Double check the Geometry structure, it may still be in Hounsfield units!');
end

geometry_folder = 'geometryfiles';
batch_output_folder = 'batchoutput';  % folder to which stdout will be printed
beamlet_batches_folder = 'beamletbatches';  % folder where resulting beamlet batches will be stored

if Condor_flag == 1
    mkdir([condor_folder '/' output_folder '/' beamlet_batches_folder]);
    mkdir([condor_folder '/' output_folder '/' batch_output_folder]);

    save_geometry(Geometry,[condor_folder '/' input_folder '/' geometry_folder],geometry_header_filename,geometry_density_filename);
    fprintf(['Successfully saved Condor geometry to ' input_folder '/' geometry_folder '\n']);

    % create geometry filenames files
    fid = fopen([condor_folder '/' input_folder '/' geometry_filenames_condor],'w');
    fid2 = fopen([condor_folder '/' input_folder '/' geometry_filenames_CHTC],'w');
    
    fprintf(fid,'geometry_header\n');
    % fprintf(fid,['./' input_folder '/' geometry_folder '/' geometry_header_filename '\n']);
    fprintf(fid,'%s\n',fullfile(patient_dir,input_folder,geometry_folder,geometry_header_filename));
    fprintf(fid2,'geometry_header\n');
    fprintf(fid2, '%s/%s\n', geometry_folder,'geometryHeader.txt');
    
    fprintf(fid,'geometry_density\n');
    % fprintf(fid,['./' input_folder '/' geometry_folder '/' geometry_density_filename '\n']);
    fprintf(fid,'%s\n',fullfile(patient_dir,input_folder,geometry_folder,geometry_density_filename));
    fprintf(fid2,'geometry_density\n');
    fprintf(fid2, '%s/%s\n', geometry_folder,'density.bin');
    fclose(fid);
    
    % write command file
    % TODO consistent naming throughout script
    for k = 1:numel(batches)
        fid = fopen(fullfile(patient_dir,sprintf('run%d.cmd',k-1)), 'w');
        fprintf(fid, '"%s" "%s" "%s" "%s" "%s"', executable_path,...
            fullfile(patient_dir, kernel_filenames_condor),...
            fullfile(patient_dir, geometry_filenames_condor),...
            fullfile(patient_dir, 'beamspecbatches', sprintf('beamspecbatch%d.txt',k-1)),...
            fullfile(patient_dir, sprintf('batch_dose%d.bin',k-1)));
        fclose(fid);
    end

    % write the condor submit file
%     beamspec_batch_filename = ['./' input_folder '/' beamspec_batches_folder '/' beamspec_batch_base_name '$(Process).txt'];
%     beamlet_batch_filename = ['./' output_folder '/' beamlet_batches_folder '/' beamlet_batch_base_name '$(Process).bin'];
%     fid = fopen([condor_folder '/' condor_submit_file],'w');
%     fprintf(fid,'###############################################################\n');
%     fprintf(fid,'# Condor submission script for convolution/superposition code\n');
%     fprintf(fid,'###############################################################\n\n');
%     fprintf(fid,'copy_to_spool = false\n');
%     fprintf(fid,['Executable   = ' code_folder '/' condor_exectuable_name '\n']);
%     fprintf(fid,['arguments    = ' input_folder '/' kernel_filenames_condor ' ' input_folder '/' geometry_filenames_condor ' ' beamspec_batch_filename ' ' beamlet_batch_filename '\n']);
%     fprintf(fid,['Output       = ./' output_folder '/' batch_output_folder '/batchout$(Process).txt\n']);
%     fprintf(fid,['Log          = ./' output_folder '/' batch_output_folder '/log.txt\n']);
%     fprintf(fid,['Queue  ' num2str(Nrot)]);
%     fclose(fid);


 
% % write the condor submit file
%      beamspec_batch_filename = ['./' input_folder '/' beamspec_batches_folder '/' beamspec_batch_base_name '$(Process).txt'];
%      beamlet_batch_filename = ['./' output_folder '/' beamlet_batches_folder '/' beamlet_batch_base_name '$(Process).bin'];
     fid = fopen([condor_folder '/' condor_submit_file],'w');
     fprintf(fid,'###############################################################\n');
     fprintf(fid,'# Condor submission script for convolution/superposition code\n');
     fprintf(fid,'###############################################################\n\n');
     fprintf(fid,'copy_to_spool = false\n');
     fprintf(fid,['Executable   = ' condor_exectuable_name '\n']);
     fprintf(fid,['Arguments    = ' kernel_filenames_CHTC ' ' geometry_filenames_CHTC ' ' beamspec_batch_base_name '$(Process).txt ' 'batch_dose$(Process).bin\n']);
     fprintf(fid,['Transfer_input_files = ' kernel_folder ',' geometry_folder ',' beamspec_batches_folder '/' beamspec_batch_base_name '$(Process).txt' ',' kernel_filenames_CHTC ',' geometry_filenames_CHTC '\n']);
     fprintf(fid,['Request_memory = 1000' '\n']);
     fprintf(fid,['Request_disk = 500000' '\n']);
     fprintf(fid,['Output       = $(Cluster).out' '\n']);
     fprintf(fid,['Log          = $(Cluster).log' '\n']);
     fprintf(fid,['Error          = $(Cluster).err' '\n']);
     fprintf(fid,['Queue  ' num2str(num_batches) '\n']);
%      fclose(fid);
end


% write the batches to files
for n=1:numel(batches)
    batch = batches{n};  % current batch

    if Condor_flag == 1
        save_beamspec_batch(batch,[condor_folder '/' input_folder '/' beamspec_batches_folder],[beamspec_batch_base_name num2str(n-1) '.txt']);
    end
end

% for k = 1:numel(batches)
%         system([fullfile(patient_dir,sprintf('run%d.cmd',k-1)) ' &']);
%     end 

% Ask for User option to run the dose calculation locally on the computer 
% or just to get necessary files for CHTC server
% 'y' means run locally, 'n' means not to run locally on the computer

strBeamlet = '';
while(1)
    if strcmpi('y',strBeamlet)
        break;
    elseif strcmpi('n',strBeamlet)
        break;
    end
    
    strBeamlet = input('Run beamlet batches dose calculation locally? y/n \n','s');
end

if(strcmpi('y',strBeamlet)) 
    for k = 1:numel(batches)
        system([fullfile(patient_dir,sprintf('run%d.cmd',k-1)) ' &']);
    end 
end