import logging

from allensdk.core.simple_tree import SimpleTree

from allensdk.internal.mouse_connectivity.interval_unionize.tissuecyte_unionizer import TissuecyteUnionizer
import data_utilities as du


def get_ancestor_id_map(structures):


    tree = SimpleTree( structures, 
                       lambda st: int(st['id']), 
                       lambda st: st['parent_structure_id'])
    ancestor_id_map = tree.value_map( lambda st: st['id'], 
                                      lambda st: tree.ancestor_ids([st['id']])[0] )
    for k in ancestor_id_map.keys():
        ancestor_id_map[-k] = map(lambda x: -x, ancestor_id_map[k])
      
    return ancestor_id_map


def get_volume_scale(image_resolution, voxel_depth):
    return image_resolution ** 2 * 10 ** -9 * voxel_depth


def run(input_data):

    logging.info('making ancestor id map')
    ancestor_id_map = get_ancestor_id_map(input_data['structures'])

    logging.info('computing volume scale factor')
    volume_scale = get_volume_scale(input_data['image_resolution'], input_data['reference_spacing'])  
    logging.info('volume scale factor : {0}'.format(volume_scale))

    logging.info('reference shape : {0}'.format(input_data['reference_shape']))
    logging.info('reference spacing : {0}'.format(input_data['reference_spacing']))
    logging.info('image_series_id : {0}'.format(input_data['image_series_id']))

    annotation = du.load_annotation(input_data['annotation_path'], input_data['grid_paths']['data_mask'])

    unionizer = TissuecyteUnionizer()
    unionizer.setup_interval_map(annotation)
    del annotation

    signal_arrays = du.get_injection_data(input_data['grid_paths']['injection_fraction'],
                                          input_data['grid_paths']['injection_density'], 
                                          input_data['grid_paths']['injection_energy'])
    signal_arrays.update(du.get_projection_data(input_data['grid_paths']['projection_density'],
                                                input_data['grid_paths']['projection_energy'], 
                                                input_data['grid_paths']['aav_exclusion_fraction']))
    signal_arrays.update(du.get_sum_pixels(input_data['grid_paths']['sum_pixels']))
    signal_arrays.update(du.get_sum_pixel_intensities(input_data['grid_paths']['sum_pixel_intensities'], 
                                                      input_data['grid_paths']['injection_sum_pixel_intensities']))

    for k, v in signal_arrays.iteritems():
        logging.info('sorting {0} array'.format(k))
        signal_arrays[k] = v.flat[unionizer.sort]
    
    logging.info('computing unionizes from directly annotated voxels')
    raw_unionizes = unionizer.direct_unionize(signal_arrays, pre_sorted=True)
    
    logging.info('propagating data to ancestor structures')
    raw_unionizes = TissuecyteUnionizer.propagate_unionizes(raw_unionizes, 
                                                            ancestor_id_map)

    logging.info('propagating data to bilateral unionizes')
    bilateral = TissuecyteUnionizer.propagate_to_bilateral(raw_unionizes)

    cooked_unionizes = unionizer.postprocess_unionizes(raw_unionizes, 
                                                       image_series_id=input_data['image_series_id'], 
                                                       output_spacing_iso=input_data['reference_spacing'], 
                                                       volume_scale=volume_scale, 
                                                       target_shape=input_data['reference_shape'],
                                                       sort=unionizer.sort)

    cooked_bilateral = unionizer.postprocess_unionizes(bilateral, 
                                                       image_series_id=input_data['image_series_id'], 
                                                       output_spacing_iso=input_data['reference_spacing'], 
                                                       volume_scale=volume_scale, 
                                                       target_shape=input_data['reference_shape'], 
                                                       sort=unionizer.sort)
    for item in cooked_bilateral:
        item['hemisphere_id'] = 3
        cooked_unionizes.append(item)
  
    logging.info('computed {0} unionize records'.format(len(cooked_unionizes)))
    return cooked_unionizes
