"""
Convert to VTK. Use pvti files for final assembly.
"""

import os
import logging
import math
from argparse import ArgumentParser
from functools import partial
import multiprocessing as mp
import vtk
import utility
import extract_zeiss_metadata


def do_conversion(file_attributes, working_directory):
    
    reader = vtk.vtkTIFFReader()
    reader.SetFileName(file_attributes["FileName"])
    reader.Update()
    
    image_data = reader.GetOutput()
    image_data.SetExtent(file_attributes["Extents"])
    image_data.SetSpacing(file_attributes["PhysicalSize"])
    
    out_path = working_directory + "/" + os.path.splitext(os.path.basename(file_attributes["FileName"]))[0] + ".vti"
    writer = vtk.vtkXMLImageDataWriter()
    writer.SetInputData(image_data)
    writer.SetFileName(out_path)
    writer.Write()

def convert_to_vti(input_folder, metadata_file, tile_dimensions, downsample_dimensions):
    
    extract_zeiss_metadata.extract_metadata(metadata_file, os.getcwd() + "temp_metadata.xml")
    image_metadata = extract_zeiss_metadata.get_metadata_as_class(os.getcwd() + "temp_metadata.xml")
    
    num_x = int(math.floor(image_metadata.Image[0].Pixels.SizeX/tile_dimensions[0]))+1
    num_y = int(math.floor(image_metadata.Image[0].Pixels.SizeY/tile_dimensions[1]))+1
    size_x = int(math.floor(tile_dimensions[0]/downsample_dimensions[0]))+1
    size_y = int(math.floor(tile_dimensions[1]/downsample_dimensions[1]))+1
    size_z = int(math.floor(image_metadata.Image[0].Pixels.SizeZ/downsample_dimensions[2]))
    physical_x = float(image_metadata.Image[0].Pixels.PhysicalSizeX)*float(downsample_dimensions[0])
    physical_y = float(image_metadata.Image[0].Pixels.PhysicalSizeY)*float(downsample_dimensions[1])
    physical_z = float(image_metadata.Image[0].Pixels.PhysicalSizeZ)*float(downsample_dimensions[2])
    
    # Grab the file base name. Assume the input folder only has the tiff files we want to convert
    path_base_name = ""
    for eachFile in os.listdir(input_folder):
        if eachFile.endswith(".tiff"):
            loc_start_index = eachFile.find("_tile")
            if loc_start_index>-1:
                path_base_name = input_folder + "/" + eachFile[:loc_start_index]
                break
            
    file_attribute_collection = []
    for jdx in range(num_y):
        for idx in range(num_x):
            global_index = idx + jdx*num_x
            file_name = path_base_name + "_tile_X" + str(idx) + "_Y" + str(jdx) + "_ID_" + str(global_index)+"_pixels_only_downsampled.tiff"
            min_ext_x = idx*size_x
            min_ext_y = jdx*size_y
            max_ext_x = (idx+1)*size_x-1
            max_ext_y = (jdx+1)*size_y-1
            min_ext_z = 0
            max_ext_z = size_z-1
            if idx == num_x - 1:
                max_ext_x = int(math.floor(image_metadata.Image[0].Pixels.SizeX/downsample_dimensions[0]))+2
            if jdx == num_y - 1:
                max_ext_y = int(math.floor(image_metadata.Image[0].Pixels.SizeY/downsample_dimensions[1]))+2
            extents = [min_ext_x, max_ext_x, min_ext_y, max_ext_y, min_ext_z, max_ext_z]
            file_attributes = {}
            file_attributes["PhysicalSize"]= (physical_x, physical_y, physical_z)
            file_attributes["Extents"]= extents
            file_attributes["FileName"]= file_name
            file_attribute_collection.append(file_attributes)
            
    # Make a partial function for multiple arg mapping
    working_directory = os.getcwd() + "/VTK_Files/"
    if not os.path.exists(working_directory):
        os.makedirs(working_directory)
    partial_extract_pixel_values_only = partial(do_conversion, working_directory=working_directory)
     
    # 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 = 1)
    pool.map(partial_extract_pixel_values_only, file_attribute_collection)

if __name__ == "__main__":
    
    # Do setup
    tool_name = "zeiss_to_tiff"
    utility.do_setup(tool_name)
    logger1 = logging.getLogger('tiff_to_vtk.'+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 converted.')
    parser.add_argument("-m", "--metadatafile", type=str, help='File with image metadata.')
    parser.add_argument("-tx", "--tile_x", type=int, help='X Tile size.')
    parser.add_argument("-ty", "--tile_y", type=int, help='Y Tile Size.')
    parser.add_argument("-dx", "--downsample_x", type=int, help='X Downsample Factor.', default=1)
    parser.add_argument("-dy", "--downsample_y", type=int, help='Y Downsample Factor.', default=1)
    parser.add_argument("-dz", "--downsample_z", type=int, help='Y Downsample Factor.', default=1)
    args = parser.parse_args()
    
    logger1.info('Start Converting Images to VTI at: ' + args.input_folder)
    convert_to_vti(args.input_folder, args.metadatafile, (args.tile_x, args.tile_y), 
                   (args.downsample_x, args.downsample_y, args.downsample_z))
    logger1.info('Completed Converting')