"""
Conversion of IMS/LSM/CZI files to OME and then Generic TIFFs
"""

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

def extract_pixel_values_only(input_file, working_directory, channel=0, significant_bits=8):

    im = skimage.external.tifffile.imread(input_file)
    pixel_vals_only = im[0, :, 0, channel, :, :] # Strip excess data
    out_path = working_directory + "/" + os.path.splitext(os.path.basename(input_file))[0] + "_pixels_only.tiff"
    
    if significant_bits==8:
        skimage.external.tifffile.imsave(out_path, skimage.img_as_ubyte(pixel_vals_only))
    elif significant_bits==16:
        skimage.external.tifffile.imsave(out_path, skimage.img_as_uint(pixel_vals_only))
    else:
        # Let skiimage decide
        skimage.external.tifffile.imsave(out_path, pixel_vals_only)
        
def do_conversion(input_file, output_file, channel = 0, series = 0):
#     
#     # Get the image metadata
    extract_zeiss_metadata.extract_metadata(input_file, os.getcwd() + "temp_metadata.xml")
    image_metadata = extract_zeiss_metadata.get_metadata_as_class(os.getcwd() + "temp_metadata.xml")
#     
#     # Choose the tile size
    tile_x = 1024
    tile_y = 1024
    
    # Black box JVM conversion via CLI. Be skeptical of outputs, 
    # at least the Z-physical size is silently lost!!!
    bftools_wrapper.convert(input_file, output_file, tile_x, tile_y, channel, series)
    
    # Make versions with just XYZ pixel values for easier use in scikit-image/VTK etc. 
    # DANGER, now all physical size is lost. Only trust physical sizes
    # in the standalone xml metadata file, obtained (and validated) using `extract_zeiss_metadata()`.
    
    # Collect the file names
    file_names = []
    for eachFile in os.listdir(output_file):
        if eachFile.endswith(".tiff"):
            file_names.append(output_file + "/" + eachFile)

    # Make a partial function for multiple arg mapping
    working_directory = output_file + "/Pixel_Values_Only/"
    if not os.path.exists(working_directory):
        os.makedirs(working_directory)
        
    significant_bits = int(image_metadata.Image[0].Pixels.SignificantBits)
    partial_extract_pixel_values_only = partial(extract_pixel_values_only, 
                                                working_directory=working_directory, 
                                                channel=channel,
                                                significant_bits=significant_bits)
     
    # 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
    run_cpus = 1
    #pool = mp.Pool(processes = run_cpus)
    pool = mp.Pool(processes = 1)
    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('format_conversion.'+tool_name)
    
    # Suppress XML Parse warnings
    pyxb_logger = logging.getLogger('pyxb')
    pyxb_logger.setLevel(logging.CRITICAL)
    
    parser = ArgumentParser()
    parser.add_argument("-i", "--input_file", type=str, help='Input file in a ZEISS format.')
    parser.add_argument("-o", "--output_folder", type=str, help='Output Tiff Location.')
    args = parser.parse_args()
    
    logger1.info('Start Converting Image at: ' + args.input_file)
    do_conversion(args.input_file, args.output_folder, channel = 1, series = 0)
    logger1.info('Completed Processing')