from parseXML import load_xml_file, get_function
import numpy as np
import time
conf = SparkConf().setAppName('test').setMaster('local[1]').set('spark.executor.memory','2g').set('spark.driver.maxResultSize','6g').set('spark.driver.memory','8g').set('spark.local.dir','/dev/shm').set('spark.storage.memoryFraction','0.2').set('spark.default.parallelism','10')
tsc=lambdaimageContext.start(conf=conf)
result = load_xml_file("./lambdaimage.xml")
count = 0
log('info')('load tiff ...')
rddA = tsc.loadImages('/home/wb/data/1-L/*.tif', inputFormat='tif-stack')
rddB = tsc.loadImages('/home/wb/data/1-R/*.tif', inputFormat='tif-stack')
log('info')('preprocess ...')
fun, para = get_function(count, result)
_rddA = eval(fun)(rddA,int(para[0]))
print fun
_rddB = eval(fun)(rddB,int(para[0]))
print fun
count += 1
fun, para = get_function(count, result)
_rddB = eval(fun)(_rddB)
print fun
rddB = eval(fun)(rddB)
print fun
count += 1
fun, para = get_function(count, result)
log('info')('registration ...')
rddB = eval(fun)(rddB)(_rddA.get(int(para[0])), _rddB.get(int(para[0])))
rddA = tsc.loadImages("/home/wb/data/1-L/*.tif", inputFormat="tif-stack")
rddB = tsc.loadImages("/home/wb/data/1-R/*.tif", inputFormat="tif-stack")
log("info")("tiff load over...")
log("info")("intensity normalization start ...")
rddA = prep.intensity_normalization(rddA)
rddB = prep.intensity_normalization(rddB)
rddB = prep.flip(rddB)
_rddA = prep.intensity_normalization(rddA, 8)
_rddB = prep.intensity_normalization(rddB, 8)
log("info")("intensity normalization over ...")
log("info")("registration start ...")
vec0 = [0, 0, 0, 1, 1, 0, 0]
# vec = reg.c_powell(_rddA.get(4), _rddB.get(4), vec0)
vec = eval(get_function("reg", result))(_rddA.get(4), _rddB.get(4), vec0)
rddB = reg.execute(rddB, vec)
log("info")("registration over ...")
log("info")("fusion start ...")
L_img_stack = rddA.collectValuesAsArray()
R_img_stack = rddB.collectValuesAsArray()
img_stack = zip(L_img_stack, R_img_stack)
rdd = tsc.loadImagesFromArray(img_stack)
# fused_img = fus.wavelet_fusion(rdd)
fused_img = eval(get_function("fus", result))(rdd)
fused_img = tsc.loadImagesFromArray(fused_img)
log("info")("fusion over ...")
log("info")("saving ...")
fused_img.exportAsTiffs("/home/wb/data/lambdaimage/fusion", overwrite=True)