# Allen Institute Software License - This software license is the 2-clause BSD
# license plus a third clause that prohibits redistribution for commercial
# purposes without further permission.
#
# Copyright 2017. Allen Institute. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Redistributions for commercial purposes are not permitted without the
# Allen Institute's written permission.
# For purposes of this license, commercial purposes is the incorporation of the
# Allen Institute's software into anything for which you will charge fees or
# other compensation. Contact terms@alleninstitute.org for commercial licensing
# opportunities.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
#
from __future__ import division, print_function, absolute_import
import re
import operator as op
from six import iteritems, string_types

import numpy as np

from .simple_tree import SimpleTree


class StructureTree( SimpleTree ):
                 
    def __init__(self, nodes):
        '''A tree whose nodes are brain structures and whose edges indicate 
        physical containment.
        
        Parameters
        ----------
        nodes : list of dict
            Each specifies a structure. Fields are:
            
            'acronym' : str
                Abbreviated name for the structure.
            'rgb_triplet' : str
                Canonical RGB uint8 color assigned to this structure
            'graph_id' : int
                Specifies the structure graph containing this structure.
            'graph_order' : int
                Canonical position in the flattened structure graph.
            'id': int
                Unique structure specifier.
            'name' : str
                Full name of structure.
            'structure_id_path' : list of int
                This structure's ancestors (inclusive) from the root of the 
                tree.
            'structure_set_ids' : list of int
                Unique identifiers of structure sets to which this structure 
                belongs.
        
        '''
        
        super(StructureTree, self).__init__(nodes,
                                            lambda s: int(s['id']),
                                            lambda s: s['structure_id_path'][-2] \
                                                if len(s['structure_id_path']) > 1 \
                                                and s['structure_id_path'] is not None \
                                                and np.isfinite(s['structure_id_path'][-2]) \
                                                else None)
                                                
                                             
    def get_structures_by_id(self, structure_ids):
        '''Obtain a list of brain structures from their structure ids
        
        Parameters
        ----------
        structure_ids : list of int
            Get structures corresponding to these ids.
            
        Returns
        -------
        list of dict : 
            Each item describes a structure.
        
        '''
    
        return self.filter_nodes(lambda x: x['id'] in structure_ids)
        
    
    def get_structures_by_name(self, names):
        '''Obtain a list of brain structures from their names,
        
        Parameters
        ----------
        names : list of str
            Get structures corresponding to these names.
            
        Returns
        -------
        list of dict : 
            Each item describes a structure.
            
        '''
        
        return self.filter_nodes(lambda x: x['name'] in names)
        
        
    def get_structures_by_acronym(self, acronyms):
        '''Obtain a list of brain structures from their acronyms
        
        Parameters
        ----------
        names : list of str
            Get structures corresponding to these acronyms.
            
        Returns
        -------
        list of dict : 
            Each item describes a structure.
            
        '''
        
        return self.filter_nodes(lambda x: x['acronym'] in acronyms)
        
        
    def get_structures_by_set_id(self, structure_set_ids):
        '''Obtain a list of brain structures from by the sets that contain 
        them.
        
        Parameters
        ----------
        structure_set_ids : list of int
            Get structures belonging to these structure sets.
            
        Returns
        -------
        list of dict : 
            Each item describes a structure.
            
        '''
        
        overlap = lambda x: (set(structure_set_ids) & set(x['structure_set_ids']))
        return self.filter_nodes(overlap)
        
        
    def get_colormap(self):
        '''Get a dictionary mapping structure ids to colors across all nodes.
        
        Returns
        -------
        dict : 
            Keys are structure ids. Values are RGB lists of integers.
        
        '''
    
        return self.value_map(lambda x: x['id'], 
                              lambda y: y['rgb_triplet'])

                              
                              
    def get_name_map(self):
        '''Get a dictionary mapping structure ids to names across all nodes.
        
        Returns
        -------
        dict : 
            Keys are structure ids. Values are structure name strings.
        
        '''
    
        return self.value_map(lambda x: x['id'], 
                              lambda y: y['name'])
        
        
    def get_id_acronym_map(self):
        '''Get a dictionary mapping structure acronyms to ids across all nodes.
        
        Returns
        -------
        dict : 
            Keys are structure acronyms. Values are structure ids.
        
        '''
        
        return self.value_map(lambda x: x['acronym'], 
                              lambda y: y['id'])
        
        
    def get_ancestor_id_map(self):
        '''Get a dictionary mapping structure ids to ancestor ids across all 
        nodes. 
        
        Returns
        -------
        dict : 
            Keys are structure ids. Values are lists of ancestor ids.
        
        '''

        return self.value_map(lambda x: x['id'], 
                              lambda y: self.ancestor_ids([y['id']])[0])
        
        
    def structure_descends_from(self, child_id, parent_id):
        '''Tests whether one structure descends from another. 
        
        Parameters
        ----------
        child_id : int
            Id of the putative child structure.
        parent_id : int
            Id of the putative parent structure.
            
        Returns
        -------
        bool :
            True if the structure specified by child_id is a descendant of 
            the one specified by parent_id. Otherwise False.
        
        '''
    
        return parent_id in self.ancestor_ids([child_id])[0]
    
    
    def get_structure_sets(self):
        '''Lists all unique structure sets that are assigned to at least one 
        structure in the tree.
        
        Returns
        -------
        list of int : 
            Elements are ids of structure sets.
        
        '''
        
        return set(reduce(op.add, map(lambda x: x['structure_set_ids'], 
                                      self.node())))
        
        
    def has_overlaps(self, structure_ids):
        '''Determine if a list of structures contains structures along with 
        their ancestors
        
        Parameters
        ----------
        structure_ids : list of int
            Check this set of structures for overlaps
            
        Returns
        -------
        set : 
            Ids of structures that are the ancestors of other structures in 
            the supplied set.
        
        '''
    
        ancestor_ids = reduce(op.add, 
                              map(lambda x: x[1:], 
                                  self.ancestor_ids(structure_ids)))
        return (set(ancestor_ids) & set(structure_ids))
        

    @staticmethod
    def clean_structures(structures, whitelist=None, data_transforms=None, renames=None):
        '''Convert structures_with_sets query results into a form that can be 
        used to construct a StructureTree
        
        Parameters
        ----------
        structures : list of dict
            Each element describes a structure. Should have a structure id path 
            field (str values) and a structure_sets field (list of dict).
        whitelist : list of str, optional
            Only these fields will be included in the final structure record. Default is 
            the output of StructureTree.whitelist.
        data_transforms : dict, optional
            Keys are str field names. Values are functions which will be applied to the 
            data associated with those fields. Default is to map colors from hex to rgb and 
            convert the structure id path to a list of int.
        renames : dict, optional
            Controls the field names that appear in the output structure records. Default is 
            to map 'color_hex_triplet' to 'rgb_triplet'.
            
        Returns
        -------
        list of dict : 
            structures, after conversion of structure_id_path and structure_sets 
        
        '''

        if whitelist is None:
            whitelist = StructureTree.whitelist()

        if data_transforms is None:
            data_transforms = StructureTree.data_transforms()

        if renames is None:        
            renames = StructureTree.renames()
            whitelist.extend(renames.values())

        for ii, st in enumerate(structures):

            StructureTree.collect_sets(st)
            record = {}

            for name in whitelist:
                
                if name not in st:
                    continue
                data = st[name]
            
                if name in data_transforms:
                    data = data_transforms[name](data)

                if name in renames:
                    name = renames[name]

                record[name] = data

            structures[ii] = record

        return structures
       
    @staticmethod
    def data_transforms():
        return  {'color_hex_triplet': StructureTree.hex_to_rgb, 
                 'structure_id_path': StructureTree.path_to_list}


    @staticmethod
    def renames():
        return {'color_hex_triplet': 'rgb_triplet'}

    @staticmethod
    def whitelist():
        return ['acronym', 'color_hex_triplet', 'graph_id', 'graph_order', 'id', 
                'name', 'structure_id_path', 'structure_set_ids'] 
    
        
    @staticmethod
    def hex_to_rgb(hex_color):
        '''Convert a hexadecimal color string to a uint8 triplet
        
        Parameters
        ----------
        hex_color : string 
            Must be 6 characters long, unless it is 7 long and the first 
            character is #. If hex_color is a triplet of int, it will be 
            returned unchanged.
            
        Returns
        -------
        list of int : 
            3 characters long - 1 per two characters in the input string.
        
        '''
        
        if not isinstance(hex_color, string_types):
            return list(hex_color)

        if hex_color[0] == '#':
            hex_color = hex_color[1:]
        
        return [int(hex_color[a * 2: a*2 + 2], 16) for a in xrange(3)] 
    

    @staticmethod
    def path_to_list(path):
        '''Structure id paths are sometimes formatted as "/"-seperated strings. 
        This method converts them to a list of integers, if needed.
        '''

        if not isinstance(path, string_types):
            return list(path)

        return [int(stid) for stid in path.split('/') if stid != '']


    @staticmethod
    def collect_sets(structure):
        '''Structure sets may be specified by full records or id. This method 
        collects all of the structure set records/ids in a structure record and 
        replaces them with a single list of id records.
        '''

        if not 'structure_sets' in structure:
            structure['structure_sets'] = []
        if not 'structure_set_ids' in structure:
            structure['structure_set_ids'] = []    
        
        structure['structure_set_ids'].extend([sts['id'] for sts 
                                               in structure['structure_sets']])
        structure['structure_set_ids'] = list(set(structure['structure_set_ids']))
        
        del structure['structure_sets']
    
