import sys
import argparse
import logging
import os

from xml.etree.ElementTree import Element, SubElement, Comment, tostring
from xml.dom import minidom

import SimpleITK as sitk
import numpy as np
from six import iteritems

from allensdk.internal.core.lims_pipeline_module import PipelineModule, run_module
from allensdk.internal.mouse_connectivity.tissuecyte_stitching.stitcher import Stitcher
from allensdk.internal.mouse_connectivity.tissuecyte_stitching.tile import Tile
import allensdk.core.json_utilities as ju

# TODO this ought to be installed with the actual python build? 
# need to consult with sysadmins/refactor jp2 project build
sys.path.append('/shared/bioapps/itk/itk_shared/jp2/build')
import jpeg_twok


logging.getLogger().setLevel(logging.INFO)
logging.captureWarnings(True)


def get_missing_tile_paths(missing_tiles):

    paths = []

    for index, path in iteritems(missing_tiles):
        spath = ','.join(map(str, path))
        logging.info('writing missing tile path for tile {0} as {1}'.format(index, spath))
        paths.append(spath)

    return paths


def read_image(file_name):
    logging.info('reading image from {0}'.format(file_name))
    image = sitk.ReadImage(str(file_name))
    return np.flipud(sitk.GetArrayFromImage(image)).T


def normalize_image_by_median(image):

    median = np.median(image)

    if median != 0:
        image = np.divide(median, image)
        image[np.isnan(image)] = 0
        image[np.isinf(image)] = 0

    return image


def load_average_tile(path):
    tile = read_image(path)
    return normalize_image_by_median(tile)


def get_average_tiles(average_tile_paths):

    average_tiles = {}    
    for key, path in iteritems(average_tile_paths):
        key = int(key) - 1

        try:
            average_tiles[key] = load_average_tile(path)
            logging.info('found average tile for channel {0} (zero-indexed)'.format(key))        
        except(IOError, OSError, RuntimeError) as err:
            average_tiles[key] = None
            logging.info('did not find average tile for channel {0} (zero-indexed)'.format(key))
        
    return average_tiles


def generate_tiles(tiles):

    for tile_params in tiles:
        tile = tile_params.copy()

        try:
            tile['image'] = read_image(tile['path'])
            tile['is_missing'] = False
        except (IOError, OSError, RuntimeError) as err:
            tile['image'] = None
            tile['is_missing'] = True
        
        tile['channel'] = tile['channel'] - 1

        tile_obj = Tile(**tile)
        del tile
        yield tile_obj


def write_output(arr, spacing, path):
    jpeg_twok.write(arr, path)


def main():

    output_json = args.output_json
    output_directory = os.path.dirname(output_json)

    slice_path = os.path.join(output_directory, data['slice_fname'])
    
    tiles = generate_tiles(data['tiles'])
    average_tiles = get_average_tiles(data['average_tile_paths'])

    stitcher = Stitcher(data['image_dimensions'], tiles, average_tiles, data['channels'])
    image, missing = stitcher.run()
    del tiles
    missing_tile_paths = get_missing_tile_paths(missing)

    write_output(np.ascontiguousarray(image), data['spacing'], slice_path)

    module_outputs = {'slice_fname': slice_path, 
                      'missing_tile_paths': missing_tile_paths}
    ju.write(output_json, module_outputs)
        

if __name__ == '__main__':

    logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')

    parser = argparse.ArgumentParser()
    parser.add_argument('input_json', type=str)
    parser.add_argument('output_json', type=str)
    args = parser.parse_args()

    data = ju.read(args.input_json)

    main()
