##! /usr/bin/python

import os
import sys
from rtplanchanger import *

d_ref = '../data/ref_plans/'
d_RS = '../data/data_RS/'
d_meas = '../data/auto_measures/'
d_plan = '../data/autoPlans/'

#process all the measurements files generated by matlab code, translate the coordinates to CT system


if not os.path.exists(d_plan):
	os.makedirs(d_plan)

def printpage():
    print "Script for auto plan emulation! Mode: FILE or DIR. Two types of usage:"
    print "./auto_gen.py FILE refplan surID (optional)"
    print "./auto_gen.py DIR (optional)"
    print "..."
    print "default: dir_ref: ../ref_plans/; dir_RS: ../data_RS/; dir_meas = ../measures/ dir_plan = ../autoPlans/ "
    print "optional: -dir_ref PATH(end with /) -dir_RS PATH -dir_m PATH -dir_plan PATH"
    quit()


if len(sys.argv) == 1: #if no parameter
    printpage()
else:
    if ( len(sys.argv) == 2 and sys.argv[1] == '-help' ): #if user ask for -help
        printpage()

# set custom options if present
for i in range(2, len(sys.argv)):
    if sys.argv[i] == '-dir_ref':
        d_ref = sys.argv[i+1]
    elif sys.argv[i] == '-dir_RS':
        d_RS = sys.argv[i+1]
    elif sys.argv[i] == '-dir_m':
        d_meas = sys.argv[i+1]
    elif sys.argv[i] == '-dir_plan':
        d_plan = sys.argv[i+1]


def performPlanEmulation(refplan, surID, block =''):
	if "sampled" in refplan:
		refID = refplan[0:refplan.find('_sampled')]
	else:
		refID = refplan

	refRP = d_ref + 'RP.ref' + refplan +'.'+block + refID+ '.dcm'
	refRS = d_RS + 'RS.' + refID + '.dcm'
	surRS = d_RS + 'RS.' + surID + '.dcm'
	measure_ref = d_meas + 'measurement_' + refID + '.txt'
	measure_sur = d_meas + 'measurement_' + surID + '.txt'
	outRP = d_plan + 'RP.ref' + refplan + '.' + block + surID + '.dcm'
	#print refRP, refRS, surRS, measure_ref, measure_sur
	print 'Emulate plan ' + refplan + ' on patient ' + surID
	if (not os.path.exists(refRP)) or (not os.path.exists(refRS)) or (not os.path.exists(surRS)) or (not os.path.exists(measure_ref)) or (not os.path.exists(measure_sur)):
		print 'Error! Some files do not exist!'
		quit()
	rtplangenerator(refRP, refRS, measure_ref, surRS, measure_sur, outRP)




def main():
    # FILE mode
    # reconstruct the specific reference plan (ID) and surrogate patient (ID)
    if sys.argv[1] == 'FILE':
        
        refplan = sys.argv[2]
        surID = sys.argv[3]
        if "block" in refplan:
            plan = refplan[0:refplan.find('.')]
            block = refplan[refplan.find('.')+1:len(refplan)]+'.'
        else:
            plan = refplan
            block = ''
                    
        #print plan, block
        performPlanEmulation(plan, surID, block)


    # DIR mode
    # reconstruct the reference plans in dir_plan on all the surrogate patients in dir_RS
    elif sys.argv[1] == 'DIR':

        for refRP in os.listdir(d_ref):
            
            # skip non-.dcm files
            if not refRP.endswith('.dcm'):
                continue

            ID = refRP[ refRP.rfind('ref') + 3 : refRP.rfind('.dcm') ]
            refID = ID[0:ID.find('.')]

            filename = "AP_diameter_isoc_ref"+refID+".txt"
            block = ID[ID.find('.')+1:ID.rfind('.')]
            if len(block) > 0: block = block + '.'
            # generate automatic plan for all surrogates
            for surRS in os.listdir(d_RS):
                # skip non-.dcm files
                if not surRS.endswith('.dcm'):
                    continue
                surID = surRS[ surRS.rfind('RS.') + 3 : surRS.rfind('.dcm') ]
                # skip reference patient with itself as surrogate
                if surID == refID:
                    continue

                performPlanEmulation(refID, surID, block)

    print 'Finished processing! Automaic emulated plans saved at', d_plan

if __name__ == "__main__":
    main()
