import logging
import os
import functools

import SimpleITK as sitk
from scipy.misc import imread, imsave
import numpy as np
import pandas as pd

from allensdk.internal.core.lims_pipeline_module import PipelineModule

from allensdk.internal.mouse_connectivity.projection_thumbnail.volume_utilities import sitk_get_diagonal_length
from allensdk.internal.mouse_connectivity.projection_thumbnail.generate_projection_strip import run, apply_colormap
from allensdk.internal.mouse_connectivity.projection_thumbnail.visualization_utilities import convert_discrete_colormap


PERMUTATION = [2, 1, 0]
FLIP = [True, False, False]


def write_depth_image(image, path):
    image = sitk.GetImageFromArray(image)
    sitk.WriteImage(image, str(path))
    

def load_background_image(path):
    background = sitk.ReadImage(str(path))
    background = sitk.GetArrayFromImage(background)
    bg_split = np.split(background, 3, axis=-1)
    return bg_split[0] / 255.0


def no_pad(volume):
    shape = [int(np.ceil(sitk_get_diagonal_length(volume))), 0, 0]
    shape[1] = volume.GetSize()[1]
    shape[2] = volume.GetSize()[2]
    return shape


def pad(volume):
    shape = [int(np.ceil(sitk_get_diagonal_length(volume))), 0, 0]
    shape[1] = int(np.floor(np.linalg.norm([volume.GetSize()[1], volume.GetSize()[0]])))
    shape[2] = int(np.floor(np.linalg.norm([volume.GetSize()[2], volume.GetSize()[0]])))
    return shape


def main():

    module = PipelineModule()
    input_data = module.input_data()

    output_dir = os.path.dirname(module.args.output_json)

    logging.info('reading data volume from {0}'.format(input_data['volume_path']))
    volume = sitk.ReadImage(str(input_data['volume_path']))
    volume = sitk.PermuteAxes(volume, PERMUTATION)
    volume = sitk.Flip(volume, FLIP)

    logging.info('reading colormap from {0}'.format(input_data['colormap_path']))
    colormap = pd.read_csv(input_data['colormap_path'], header=None, 
                           names=['red', 'green', 'blue'], delim_whitespace=True)
    colormap = convert_discrete_colormap(colormap.values, 'projection')

    output_data = {'output_file_paths': []}
    for rot in input_data['rotations']:
    
        rot['write_depth_sheet'] = functools.partial(write_depth_image, 
                                                     path=str(os.path.join(output_dir, rot['depth_path'])))
        output_data['output_file_paths'].append(os.path.join(output_dir, rot['depth_path']))

        if isinstance(rot['window_size'], (str, unicode)):
            if rot['window_size'] == 'no_pad':
                rot['window_size'] = no_pad(volume)
            elif rot['window_size'] == 'pad':
                rot['window_size'] = pad(volume)
            else:
                raise ValueError('did not understand window size option {0}'.format(rot['window_size']))
        logging.info('window_size: {0}'.format(rot['window_size']))
        
        for out_image in rot['output_images']:
            out_image['write'] = functools.partial(imsave, os.path.join(output_dir, out_image['path']))
            output_data['output_file_paths'].append(os.path.join(output_dir, out_image['path']))            

            if 'background_path' in out_image:
                out_image['background'] = load_background_image(out_image['background_path'])
            else:
                out_image['background'] = None

    run(volume, input_data['min_threshold'], input_data['max_threshold'], 
        input_data['rotations'], colormap)
    module.write_output_data(output_data)


if __name__ == '__main__':
    main()
