#!/usr/bin/env python

import sys
import os

from efidir.config import *
from efidir.param import *
from efidir.image import *
from efidir.sws import *
from efidir.utils import listimage

define_param("masterFileName", STRING, "Master image file name", None)
define_optional_param("maskFileName", STRING, "Mask image file name", None)
define_param("slaveFileName", STRING, "Slave image file name", None)
define_optional_param("first_line", INT, "Index of first line of the coorelated lines band (from 0: default value)", None)
define_optional_param("last_line", INT, "Index of last line of the coorelated lines band (if unset last_line = master image lines)", None)
define_param("displacementFileName", STRING, "Displacement image file name", None)
define_optional_param("displacementDescription", STRING, "Displacement image description", None)
define_optional_param("correlationPeakFileName", STRING, "CorrelationPeak image file name", None)
define_optional_param("correlationPeakDescription", STRING, "CorrelationPeak image description", None)
define_optional_param("correlationFWHMFileName", STRING, "CorrelationFWHM image file name", None)
define_optional_param("correlationFWHMDescription", STRING, "CorrelationFWHM image description", None)
define_optional_param("correlationHPeakFileName", STRING, "CorrelationHPeak image file name (Use NAN in some results for Parabolic 3D error)", None)
define_optional_param("correlationHPeakDescription", STRING, "CorrelationHPeak image description", None)
define_optional_param("correlationSubPixColorMapFileName", STRING, "CorrelationSubPixColorMap image file name (Use NAN in some results for Parabolic 3D error)", None)
define_optional_param("correlationSubPixColorMapDescription", STRING, "CorrelationSubPixColorMap image description", None)
define_optional_param("optimization", INT, "Optimization method (0: none, 1: opt1, 2: opt2, 3: opt3 (default), 4: para opt1, 5: para opt3)", None)
define_optional_param("requiredCore", INT, "Number of Core requires for computation (0: automatic detection (default))", None)

# dist corr op param
define_param("window_nb_rows", INT, "Number of lines for the process window", None)
define_optional_param("window_nb_columns", INT, "Number of comuns for the process window (if unset window_nb_columns = window_nb_rows)", None)

define_param("window_search_nb_rows", INT, "Number of lines for the search window", None)
define_optional_param("window_search_nb_columns", INT, "Number of columns for the search window (if unset window_search_nb_columns = window_search_nb_rows)", None)

define_param("sub_pixel", INT, "Sub-pixel correlation (0: None, 1:Parabolic 3D, 2: Parabolic 2D, 3: Mix parabolic 3D and 2D (2D instead of 3D where 3D does not work))", None)
define_optional_param("step", INT, "step between computation points (if unset step = 1, if set optimisations cannot be used )", None)
#
#   run parameter manager
#    
init_param(sys.argv)

#
# get parameters
#

outputFiles = []

masterFileName = get_string_value_param("masterFileName")


if(is_set_param("maskFileName")):
    maskFileName = get_string_value_param("maskFileName")
else:
    maskFileName = None
    
slaveFileName = get_string_value_param("slaveFileName")

if(is_set_param("first_line")):
    first_line = get_int_value_param("first_line")
else:
    first_line = 0

if(is_set_param("last_line")):
    last_line = get_int_value_param("last_line")
else:
    last_line = 0
    
displacementFileName = get_string_value_param("displacementFileName")
outputFiles.append(displacementFileName)
displacementDescription = get_string_value_param("displacementDescription")

correlationPeakFileName = None
if(is_set_param("correlationPeakFileName")):
    correlationPeakFileName = get_string_value_param("correlationPeakFileName")
    outputFiles.append(correlationPeakFileName)

correlationPeakDescription = None
if(is_set_param("correlationPeakDescription")):
    correlationPeakDescription = get_string_value_param("correlationPeakDescription")

correlationFWHMFileName = None
if(is_set_param("correlationFWHMFileName")):
    correlationFWHMFileName = get_string_value_param("correlationFWHMFileName")
    outputFiles.append(correlationFWHMFileName)

correlationFWHMDescription = None
if(is_set_param("correlationFWHMDescription")):
    correlationFWHMDescription = get_string_value_param("correlationFWHMDescription")
    
optimization = 3
if(is_set_param("optimization")):
    optimization = get_int_value_param("optimization")

correlationHPeakFileName = None
if(is_set_param("correlationHPeakFileName")):
    correlationHPeakFileName = get_string_value_param("correlationHPeakFileName")
    outputFiles.append(correlationHPeakFileName)

correlationHPeakDescription = None
if(is_set_param("correlationHPeakDescription")):
    correlationHPeakDescription = get_string_value_param("correlationHPeakDescription")

    
correlationSubPixColorMapFileName = None
if(is_set_param("correlationSubPixColorMapFileName")):
    correlationSubPixColorMapFileName = get_string_value_param("correlationSubPixColorMapFileName")
    outputFiles.append(correlationSubPixColorMapFileName)

correlationSubPixColorMapDescription = None
if(is_set_param("correlationSubPixColorMapDescription")):
    correlationSubPixColorMapDescription = get_string_value_param("correlationSubPixColorMapDescription")
    
requiredCore = 0
if(is_set_param("requiredCore")):
    requiredCore = get_int_value_param("requiredCore")

    
coreAvailable_predefined = get_njobs() 
coreAvailable = -1
if (coreAvailable_predefined == 0):
    coreAvailable = set_njobs(requiredCore) # return the real nuber of core available
elif (coreAvailable_predefined < requiredCore):
    coreAvailable = coreAvailable_predefined
else:
    coreAvailable = requiredCore
    


# dist corr op param
window_nb_r = get_int_value_param("window_nb_rows")
if(is_set_param("window_nb_columns")):
    window_nb_c = get_int_value_param("window_nb_columns")
else:
    window_nb_c = window_nb_r

window_search_nb_r = get_int_value_param("window_search_nb_rows")
if(is_set_param("window_search_nb_columns")):
    window_search_nb_c = get_int_value_param("window_search_nb_columns")
else:
    window_search_nb_c = window_search_nb_r
    
sub_pixel = get_int_value_param("sub_pixel")

if(is_set_param("step")):
    step = get_int_value_param("step")
else:
    step = 1

targets = []

image = read_image_open(masterFileName)
number_lines = int(image["lines"])
number_samples = int(image["samples"])

window_rows_border = (window_search_nb_r - 1)/2

subimage_size = (number_lines - window_search_nb_r + 1)/coreAvailable
if(subimage_size < 1):
    subimage_size = 1
    coreAvailable = number_lines - window_search_nb_r + 1
    
tmp_files = {}
tmp_files[displacementFileName] = []
if correlationPeakFileName is not None:
    tmp_files[correlationPeakFileName] = []
if correlationFWHMFileName is not None:
    tmp_files[correlationFWHMFileName] = []
if correlationHPeakFileName is not None:
    tmp_files[correlationHPeakFileName] = []
if  correlationSubPixColorMapFileName is not None:
    tmp_files[correlationSubPixColorMapFileName] = []

sub_dir = os.path.join(os.path.dirname(displacementFileName),
                       "dist_corr_splitter_" + os.path.basename(displacementFileName))


inputs = []
outputs = []    
for i in range(coreAvailable):
    tmp_images = []
    # compute size of sub-image
    if(i == 0):
        first_line = 0
        
    if(i > 0):
        first_line = (last_line + 1) - (window_search_nb_r - 1)
        
    if(i < coreAvailable-1):
        last_line = first_line + subimage_size + 2*window_rows_border - 1
        
    if(i == coreAvailable-1):
        last_line = number_lines - 1

    inputs = []
    outputs = []      

    command = "dist_corr "

    command = command + " --masterFileName " + masterFileName
    inputs.append(masterFileName)
    inputs.append(masterFileName+".hdr")

    if maskFileName is not None:
        command = command + " --maskFileName " + maskFileName
        inputs.append(maskFileName)
        inputs.append(maskFileName+".hdr")

    command = command + " --slaveFileName " + slaveFileName
    inputs.append(slaveFileName)
    inputs.append(slaveFileName+".hdr")

    command = command + " --first_line " + str(first_line)
    command = command + " --last_line " + str(last_line)

  
    outputs.append(os.path.join(sub_dir, os.path.basename(displacementFileName) + "_" + str(i)))
    outputs.append(os.path.join(sub_dir, os.path.basename(displacementFileName) + "_" + str(i)+".hdr"))
    command = command + " --displacementFileName " + outputs[-2] # -2 => second last added into outputs
    tmp_files[displacementFileName].append(os.path.basename(outputs[-2])) 

    if correlationPeakFileName is not None:
        outputs.append(os.path.join(sub_dir, os.path.basename(correlationPeakFileName)  + "_" + str(i)))
        outputs.append(os.path.join(sub_dir, os.path.basename(correlationPeakFileName)  + "_" + str(i)+".hdr"))
        command = command + " --correlationPeakFileName " + outputs[-2] # -2 => second last added into outputs
        tmp_files[correlationPeakFileName].append(os.path.basename(outputs[-2]))

    if correlationFWHMFileName is not None:
        outputs.append(os.path.join(sub_dir, os.path.basename(correlationFWHMFileName) + "_" + str(i)))
        outputs.append(os.path.join(sub_dir, os.path.basename(correlationFWHMFileName) + "_" + str(i)+".hdr"))
        command = command + " --correlationFWHMFileName " + outputs[-2] # -2 => second last added into outputs
        tmp_files[correlationFWHMFileName].append(os.path.basename(outputs[-2]))

    if correlationHPeakFileName is not None:
        outputs.append(os.path.join(sub_dir, os.path.basename(correlationHPeakFileName) + "_" + str(i)))
        outputs.append(os.path.join(sub_dir, os.path.basename(correlationHPeakFileName) + "_" + str(i)+".hdr"))
        command = command + " --correlationHPeakFileName " + outputs[-2] # -2 => second last added into outputs
        tmp_files[correlationHPeakFileName].append(os.path.basename(outputs[-2]))

    if correlationSubPixColorMapFileName is not None:
        outputs.append(os.path.join(sub_dir, os.path.basename(correlationSubPixColorMapFileName) + "_" + str(i)))
        outputs.append(os.path.join(sub_dir, os.path.basename(correlationSubPixColorMapFileName) + "_" + str(i)+".hdr"))
        command = command + " --correlationSubPixColorMapFileName " + outputs[-2] # -2 => second last added into outputs
        tmp_files[correlationSubPixColorMapFileName].append(os.path.basename(outputs[-2]))

    command = command + " --optimization " + str(optimization)
    command = command + " --window_nb_rows " + str(window_nb_r)
    command = command + " --window_nb_columns " + str(window_nb_c)
    command = command + " --window_search_nb_rows " + str(window_search_nb_r)
    command = command + " --window_search_nb_columns " + str(window_search_nb_c)
    command = command + " --sub_pixel " + str(sub_pixel)

    #print("command", command)
    #print("in", inputs)
    #print("out", outputs)
    
    targets.append(target_create([command], inputs, outputs))

chain = chain_create("dist_corr_splitter.mk", targets)
chain_save(chain, "dist_corr_splitter.mk")
chain_run(chain)

for list_file_name in tmp_files.keys():

    #generate list file
    #listimage.genListImage(sub_dir, ".hdr", os.path.basename(list_file_name) + "_list")
    f = open(sub_dir + "/" + os.path.basename(list_file_name) + "_list", 'w')
    f.write('\n'.join(tmp_files[list_file_name]))
    f.write('\n')
    f.close()

    inputs =[]
    for f in tmp_files[list_file_name]:
        inputs = inputs +[os.path.join(sub_dir, f)]
        
    inputs = inputs + [os.path.join(sub_dir, os.path.basename(list_file_name) + "_list")]
    outputs = [os.path.join(sub_dir, os.path.basename(list_file_name) + ".nopad"), 
               os.path.join(sub_dir, os.path.basename(list_file_name) + ".nopad.hdr")]  
    command = "merge_series "
    command = command + " --save " + os.path.join(sub_dir, "merge_"+os.path.basename(list_file_name)+".save")
    command = command + " --input " + inputs[-1] # -1 for last added
    command = command + " --output " + outputs[0]

    targets.append(target_create([command], inputs, outputs))

    #shrink if step != 1
    #...

    inputs = [os.path.join(sub_dir, os.path.basename(list_file_name) + ".nopad"),
              os.path.join(sub_dir, os.path.basename(list_file_name) + ".nopad.hdr")]
    outputs = [list_file_name, 
               list_file_name + ".hdr"]  
    command = "zeropad "
    command = command + " --input " + inputs[0]
    command = command + " --output " + outputs[0]
    command = command + " --ncols " + str(number_samples)
    command = command + " --nrows " + str(number_lines)
    command = command + " --xoff 0 "
    command = command + " --yoff "  + str((window_search_nb_r -1)/2)

    targets.append(target_create([command], inputs, outputs))

chain = chain_create("dist_corr_splitter.mk", targets)
chain_save(chain, "dist_corr_splitter.mk")
chain_run(chain)


