"""
Downsample the tiffs for easier postprocessing
"""

import os
import logging
import math
from argparse import ArgumentParser
from functools import partial
import multiprocessing as mp
import skimage.external.tifffile
import skimage.transform
from skimage import exposure
import utility

def do_downsample(input_file, working_directory, downsample_ratios):
    
    im = skimage.external.tifffile.imread(input_file)
    downsampled = skimage.transform.downscale_local_mean(im, downsample_ratios)
    out_path = working_directory + "/" + os.path.splitext(os.path.basename(input_file))[0] + "_downsampled.tiff"
    scaled_image = exposure.rescale_intensity(downsampled, out_range=(0.0, 1.0))
    skimage.external.tifffile.imsave(out_path, skimage.img_as_ubyte(scaled_image))
    
def downsample_tiffs(input_folder, downsample_ratios, channel = 0, series = 0):
    
    file_names = []
    dir_name = os.path.dirname(input_folder)
    for eachFile in os.listdir(input_folder):
        if eachFile.endswith(".tiff"):
            file_names.append(dir_name + "/" + eachFile)
    
    # Make a partial function for multiple arg mapping
    working_directory = os.getcwd() + "/Downsampled_X"+str(downsample_ratios[2]) + "_Y"+str(downsample_ratios[1])
    working_directory += "_Z"+str(downsample_ratios[0]) + "/"
    if not os.path.exists(working_directory):
        os.makedirs(working_directory)
    partial_extract_pixel_values_only = partial(do_downsample, 
                                                working_directory=working_directory, 
                                                downsample_ratios=downsample_ratios)
     
    # Limit to max 6 cpus
    num_available_cpus = mp.cpu_count()
    if num_available_cpus > 6:
        num_available_cpus = 6
        
    # Distribute across processes - Be conservative with available CPUs
    run_cpus = int(math.floor(num_available_cpus/2))
    if run_cpus == 0:
        run_cpus = 1
    #pool = mp.Pool(processes = run_cpus)
    pool = mp.Pool(processes = 3)
    pool.map(partial_extract_pixel_values_only, file_names)

if __name__ == "__main__":
    
    # Do setup
    tool_name = "zeiss_to_tiff"
    utility.do_setup(tool_name)
    logger1 = logging.getLogger('downsample_tiffs.'+tool_name)
    
    # Suppress XML Parse warnings
    pyxb_logger = logging.getLogger('pyxb')
    pyxb_logger.setLevel(logging.CRITICAL)
    
    parser = ArgumentParser()
    parser.add_argument("-i", "--input_folder", type=str, help='Folder with tiffs to be downsampled.')
    parser.add_argument("-x", "--x_reduction", type=int, help='Factor reduction in x.', default=1)
    parser.add_argument("-y", "--y_reduction", type=int, help='Factor reduction in y.', default=1)
    parser.add_argument("-z", "--z_reduction", type=int, help='Factor reduction in z.', default=1)
    args = parser.parse_args()
    
    logger1.info('Start Downsampling Images at: ' + args.input_folder)
    downsample_ratios = (args.z_reduction, args.y_reduction, args.x_reduction)
    downsample_tiffs(args.input_folder, downsample_ratios)
    logger1.info('Completed Downsampling')