# Copyright (C) 2015 Atsushi Togo
# All rights reserved.
#
# This file is part of spglib.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
# * Redistributions of source code must retain the above copyright
#   notice, this list of conditions and the following disclaimer.
#
# * 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.
#
# * Neither the name of the spglib project nor the names of its
#   contributors may be used to endorse or promote products derived
#   from this software without specific prior written permission.
#
# 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 phonopy import _spglib as spg
import numpy as np

class SpglibError(object):
    message = "no error"

spglib_error = SpglibError()

def get_version():
    _set_no_error()
    return tuple(spg.version())

def get_symmetry(cell, symprec=1e-5, angle_tolerance=-1.0):
    """This gives crystal symmetry operations from a crystal structure.

    Args:
        cell: Crystal structrue given either in Atoms object or tuple.
            In the case given by a tuple, it has to follow the form below,
            (Lattice parameters in a 3x3 array (see the detail below),
             Fractional atomic positions in an Nx3 array,
             Integer numbers to distinguish species in a length N array,
             (optional) Collinear magnetic moments in a length N array),
            where N is the number of atoms.
            Lattice parameters are given in the form:
                [[a_x, a_y, a_z],
                 [b_x, b_y, b_z],
                 [c_x, c_y, c_z]]
        symprec:
            float: Symmetry search tolerance in the unit of length.
        angle_tolerance:
            float: Symmetry search tolerance in the unit of angle deg.
                If the value is negative, an internally optimized routine
                is used to judge symmetry.

    Return:
        A dictionary: Rotation parts and translation parts. Dictionary keys:
            'rotations': Gives the numpy 'intc' array of the rotation matrices.
            'translations': Gives the numpy 'double' array of fractional
                translations with respect to a, b, c axes.

    """
    _set_no_error()

    lattice, positions, numbers, magmoms = _expand_cell(cell)
    if lattice is None:
        return None

    multi = 48 * len(positions)
    rotation = np.zeros((multi, 3, 3), dtype='intc')
    translation = np.zeros((multi, 3), dtype='double')

    # Get symmetry operations
    if magmoms is None:
        dataset = get_symmetry_dataset(cell,
                                       symprec=symprec,
                                       angle_tolerance=angle_tolerance)
        if dataset is None:
            return None
        else:
            return {'rotations': dataset['rotations'],
                    'translations': dataset['translations'],
                    'equivalent_atoms': dataset['equivalent_atoms']}
    else:
        equivalent_atoms = np.zeros(len(magmoms), dtype='intc')
        num_sym = spg.symmetry_with_collinear_spin(rotation,
                                                   translation,
                                                   equivalent_atoms,
                                                   lattice,
                                                   positions,
                                                   numbers,
                                                   magmoms,
                                                   symprec,
                                                   angle_tolerance)
        _set_error_message()
        if num_sym == 0:
            return None
        else:
            return {'rotations': np.array(rotation[:num_sym],
                                          dtype='intc', order='C'),
                    'translations': np.array(translation[:num_sym],
                                             dtype='double', order='C'),
                    'equivalent_atoms': equivalent_atoms}

def get_symmetry_dataset(cell,
                         symprec=1e-5,
                         angle_tolerance=-1.0,
                         hall_number=0):
    """Search symmetry dataset from an input cell.

    Args:
        cell, symprec, angle_tolerance:
            See the docstring of get_symmetry.
        hall_number: If a serial number of Hall symbol (>0) is given,
                     the database corresponding to the Hall symbol is made.

    Return:
        A dictionary is returned. Dictionary keys:
            number (int): International space group number
            international (str): International symbol
            hall (str): Hall symbol
            choice (str): Centring, origin, basis vector setting
            transformation_matrix (3x3 float):
                Transformation matrix from input lattice to standardized
                lattice:
                    L^original = L^standardized * Tmat
            origin shift (3 float):
                Origin shift from standardized to input origin
            rotations (3x3 int), translations (float vector):
                Rotation matrices and translation vectors. Space group
                operations are obtained by
                [(r,t) for r, t in zip(rotations, translations)]
            wyckoffs (n char): Wyckoff letters
            equivalent_atoms (n int): Symmetrically equivalent atoms
            mapping_to_primitive (n int):
                Original cell atom index mapping to primivie cell atom index
            Idealized standardized unit cell:
                std_lattice (3x3 float, row vectors),
                std_positions (Nx3 float), std_types (N int)
            std_rotation_matrix:
                Rigid rotation matrix to rotate from standardized basis
                vectors to idealized standardized basis vectors
                    L^idealized = R * L^standardized
            std_mapping_to_primitive (m int):
                std_positions index mapping to those of primivie cell atoms
            pointgroup (str): Pointgroup symbol

        If it fails, None is returned.

    """
    _set_no_error()

    lattice, positions, numbers, _ = _expand_cell(cell)
    if lattice is None:
        return None

    spg_ds = spg.dataset(lattice, positions, numbers, hall_number,
                         symprec, angle_tolerance)
    if spg_ds is None:
        _set_error_message()
        return None

    keys = ('number',
            'hall_number',
            'international',
            'hall',
            'choice',
            'transformation_matrix',
            'origin_shift',
            'rotations',
            'translations',
            'wyckoffs',
            'site_symmetry_symbols',
            'equivalent_atoms',
            'mapping_to_primitive',
            'std_lattice',
            'std_types',
            'std_positions',
            'std_rotation_matrix',
            'std_mapping_to_primitive',
            # 'pointgroup_number',
            'pointgroup')
    dataset = {}
    for key, data in zip(keys, spg_ds):
        dataset[key] = data

    dataset['international'] = dataset['international'].strip()
    dataset['hall'] = dataset['hall'].strip()
    dataset['choice'] = dataset['choice'].strip()
    dataset['transformation_matrix'] = np.array(
        dataset['transformation_matrix'], dtype='double', order='C')
    dataset['origin_shift'] = np.array(dataset['origin_shift'], dtype='double')
    dataset['rotations'] = np.array(dataset['rotations'],
                                    dtype='intc', order='C')
    dataset['translations'] = np.array(dataset['translations'],
                                       dtype='double', order='C')
    letters = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
    dataset['wyckoffs'] = [letters[x] for x in dataset['wyckoffs']]
    dataset['site_symmetry_symbols'] = [
        s.strip() for s in dataset['site_symmetry_symbols']]
    dataset['equivalent_atoms'] = np.array(dataset['equivalent_atoms'],
                                           dtype='intc')
    dataset['mapping_to_primitive'] = np.array(dataset['mapping_to_primitive'],
                                               dtype='intc')
    dataset['std_lattice'] = np.array(np.transpose(dataset['std_lattice']),
                                      dtype='double', order='C')
    dataset['std_types'] = np.array(dataset['std_types'], dtype='intc')
    dataset['std_positions'] = np.array(dataset['std_positions'],
                                        dtype='double', order='C')
    dataset['std_rotation_matrix'] = np.array(dataset['std_rotation_matrix'],
                                              dtype='double', order='C')
    dataset['std_mapping_to_primitive'] = np.array(
        dataset['std_mapping_to_primitive'], dtype='intc')
    dataset['pointgroup'] = dataset['pointgroup'].strip()

    _set_error_message()
    return dataset

def get_spacegroup(cell, symprec=1e-5, angle_tolerance=-1.0, symbol_type=0):
    """Return space group in international table symbol and number as a string.

    If it fails, None is returned.
    """
    _set_no_error()

    dataset = get_symmetry_dataset(cell,
                                   symprec=symprec,
                                   angle_tolerance=angle_tolerance)
    if dataset is None:
        return None

    spg_type = get_spacegroup_type(dataset['hall_number'])
    if symbol_type == 1:
        return "%s (%d)" % (spg_type['schoenflies'], dataset['number'])
    else:
        return "%s (%d)" % (spg_type['international_short'], dataset['number'])

def get_hall_number_from_symmetry(rotations, translations, symprec=1e-5):
    """Hall number is obtained from a set of symmetry operations.

    If it fails, None is returned.
    """

    r = np.array(rotations, dtype='intc', order='C')
    t = np.array(translations, dtype='double', order='C')
    hall_number = spg.hall_number_from_symmetry(r, t, symprec)
    return hall_number

def get_spacegroup_type(hall_number):
    """Translate Hall number to space group type information.

    If it fails, None is returned.
    """
    _set_no_error()

    keys = ('number',
            'international_short',
            'international_full',
            'international',
            'schoenflies',
            'hall_symbol',
            'choice',
            'pointgroup_schoenflies',
            'pointgroup_international',
            'arithmetic_crystal_class_number',
            'arithmetic_crystal_class_symbol')
    spg_type_list = spg.spacegroup_type(hall_number)
    _set_error_message()

    if spg_type_list is not None:
        spg_type = dict(zip(keys, spg_type_list))
        for key in spg_type:
            if key != 'number' and key != 'arithmetic_crystal_class_number':
                spg_type[key] = spg_type[key].strip()
        return spg_type
    else:
        return None

def get_pointgroup(rotations):
    """Return point group in international table symbol and number.

    The symbols are mapped to the numbers as follows:
    1   "1    "
    2   "-1   "
    3   "2    "
    4   "m    "
    5   "2/m  "
    6   "222  "
    7   "mm2  "
    8   "mmm  "
    9   "4    "
    10  "-4   "
    11  "4/m  "
    12  "422  "
    13  "4mm  "
    14  "-42m "
    15  "4/mmm"
    16  "3    "
    17  "-3   "
    18  "32   "
    19  "3m   "
    20  "-3m  "
    21  "6    "
    22  "-6   "
    23  "6/m  "
    24  "622  "
    25  "6mm  "
    26  "-62m "
    27  "6/mmm"
    28  "23   "
    29  "m-3  "
    30  "432  "
    31  "-43m "
    32  "m-3m "
    """
    _set_no_error()

    # (symbol, pointgroup_number, transformation_matrix)
    pointgroup = spg.pointgroup(np.array(rotations, dtype='intc', order='C'))
    _set_error_message()
    return pointgroup

def standardize_cell(cell,
                     to_primitive=False,
                     no_idealize=False,
                     symprec=1e-5,
                     angle_tolerance=-1.0):
    """Return standardized cell.

    Args:
        cell, symprec, angle_tolerance:
            See the docstring of get_symmetry.
        to_primitive:
            bool: If True, the standardized primitive cell is created.
        no_idealize:
            bool: If True,  it is disabled to idealize lengths and angles of
                  basis vectors and positions of atoms according to crystal
                  symmetry.
    Return:
        The standardized unit cell or primitive cell is returned by a tuple of
        (lattice, positions, numbers).
        If it fails, None is returned.
    """
    _set_no_error()

    lattice, _positions, _numbers, _ = _expand_cell(cell)
    if lattice is None:
        return None

    # Atomic positions have to be specified by scaled positions for spglib.
    num_atom = len(_positions)
    positions = np.zeros((num_atom * 4, 3), dtype='double', order='C')
    positions[:num_atom] = _positions
    numbers = np.zeros(num_atom * 4, dtype='intc')
    numbers[:num_atom] = _numbers
    num_atom_std = spg.standardize_cell(lattice,
                                        positions,
                                        numbers,
                                        num_atom,
                                        to_primitive * 1,
                                        no_idealize * 1,
                                        symprec,
                                        angle_tolerance)
    _set_error_message()

    if num_atom_std > 0:
        return (np.array(lattice.T, dtype='double', order='C'),
                np.array(positions[:num_atom_std], dtype='double', order='C'),
                np.array(numbers[:num_atom_std], dtype='intc'))
    else:
        return None

def refine_cell(cell, symprec=1e-5, angle_tolerance=-1.0):
    """Return refined cell.

    The standardized unit cell is returned by a tuple of
    (lattice, positions, numbers).
    If it fails, None is returned.
    """
    _set_no_error()

    lattice, _positions, _numbers, _ = _expand_cell(cell)
    if lattice is None:
        return None

    # Atomic positions have to be specified by scaled positions for spglib.
    num_atom = len(_positions)
    positions = np.zeros((num_atom * 4, 3), dtype='double', order='C')
    positions[:num_atom] = _positions
    numbers = np.zeros(num_atom * 4, dtype='intc')
    numbers[:num_atom] = _numbers
    num_atom_std = spg.refine_cell(lattice,
                                   positions,
                                   numbers,
                                   num_atom,
                                   symprec,
                                   angle_tolerance)
    _set_error_message()

    if num_atom_std > 0:
        return (np.array(lattice.T, dtype='double', order='C'),
                np.array(positions[:num_atom_std], dtype='double', order='C'),
                np.array(numbers[:num_atom_std], dtype='intc'))
    else:
        return None

def find_primitive(cell, symprec=1e-5, angle_tolerance=-1.0):
    """Primitive cell is searched in the input cell.

    The primitive cell is returned by a tuple of (lattice, positions, numbers).
    If it fails, None is returned.
    """
    _set_no_error()

    lattice, positions, numbers, _ = _expand_cell(cell)
    if lattice is None:
        return None

    num_atom_prim = spg.primitive(lattice,
                                  positions,
                                  numbers,
                                  symprec,
                                  angle_tolerance)
    _set_error_message()

    if num_atom_prim > 0:
        return (np.array(lattice.T, dtype='double', order='C'),
                np.array(positions[:num_atom_prim], dtype='double', order='C'),
                np.array(numbers[:num_atom_prim], dtype='intc'))
    else:
        return None

def get_symmetry_from_database(hall_number):
    """Return symmetry operations corresponding to a Hall symbol.

    The Hall symbol is given by the serial number in between 1 and 530.
    The symmetry operations are given by a dictionary whose keys are
    'rotations' and 'translations'.
    If it fails, None is returned.
    """
    _set_no_error()

    rotations = np.zeros((192, 3, 3), dtype='intc')
    translations = np.zeros((192, 3), dtype='double')
    num_sym = spg.symmetry_from_database(rotations, translations, hall_number)
    _set_error_message()

    if num_sym is None:
        return None
    else:
        return {'rotations':
                np.array(rotations[:num_sym], dtype='intc', order='C'),
                'translations':
                np.array(translations[:num_sym], dtype='double', order='C')}

############
# k-points #
############
def get_grid_point_from_address(grid_address, mesh):
    """Return grid point index by tranlating grid address"""
    _set_no_error()

    return spg.grid_point_from_address(np.array(grid_address, dtype='intc'),
                                       np.array(mesh, dtype='intc'))

def get_ir_reciprocal_mesh(mesh,
                           cell,
                           is_shift=None,
                           is_time_reversal=True,
                           symprec=1e-5,
                           is_dense=False):
    """Return k-points mesh and k-point map to the irreducible k-points.

    The symmetry is serched from the input cell.

    Parameters
    ----------
    mesh : array_like
        Uniform sampling mesh numbers.
        dtype='intc', shape=(3,)
    cell : spglib cell tuple
        Crystal structure.
    is_shift : array_like, optional
        [0, 0, 0] gives Gamma center mesh and value 1 gives half mesh shift.
        Default is None which equals to [0, 0, 0].
        dtype='intc', shape=(3,)
    is_time_reversal : bool, optional
        Whether time reversal symmetry is included or not. Default is True.
    symprec : float, optional
        Symmetry tolerance in distance. Default is 1e-5.
    is_dense : bool, optional
        grid_mapping_table is returned with dtype='uintp' if True. Otherwise
        its dtype='intc'. Default is False.

    Returns
    -------
    grid_mapping_table : ndarray
        Grid point mapping table to ir-gird-points.
        dtype='intc' or 'uintp', shape=(prod(mesh),)
    grid_address : ndarray
        Address of all grid points.
        dtype='intc', shspe=(prod(mesh), 3)

    """
    _set_no_error()

    lattice, positions, numbers, _ = _expand_cell(cell)
    if lattice is None:
        return None

    if is_dense:
        dtype = 'uintp'
    else:
        dtype = 'intc'
    grid_mapping_table = np.zeros(np.prod(mesh), dtype=dtype)
    grid_address = np.zeros((np.prod(mesh), 3), dtype='intc')
    if is_shift is None:
        is_shift = [0, 0, 0]
    if spg.ir_reciprocal_mesh(
            grid_address,
            grid_mapping_table,
            np.array(mesh, dtype='intc'),
            np.array(is_shift, dtype='intc'),
            is_time_reversal * 1,
            lattice,
            positions,
            numbers,
            symprec) > 0:
        return grid_mapping_table, grid_address
    else:
        return None

def get_stabilized_reciprocal_mesh(mesh,
                                   rotations,
                                   is_shift=None,
                                   is_time_reversal=True,
                                   qpoints=None,
                                   is_dense=False):
    """Return k-point map to the irreducible k-points and k-point grid points.

    The symmetry is searched from the input rotation matrices in real space.

    Parameters
    ----------
    mesh : array_like
        Uniform sampling mesh numbers.
        dtype='intc', shape=(3,)
    rotations : array_like
        Rotation matrices with respect to real space basis vectors.
        dtype='intc', shape=(rotations, 3)
    is_shift : array_like
        [0, 0, 0] gives Gamma center mesh and value 1 gives  half mesh shift.
        dtype='intc', shape=(3,)
    is_time_reversal : bool
        Time reversal symmetry is included or not.
    qpoints : array_like
        q-points used as stabilizer(s) given in reciprocal space with respect
        to reciprocal basis vectors.
        dtype='double', shape=(qpoints ,3) or (3,)
    is_dense : bool, optional
        grid_mapping_table is returned with dtype='uintp' if True. Otherwise
        its dtype='intc'. Default is False.

    Returns
    -------
    grid_mapping_table : ndarray
        Grid point mapping table to ir-gird-points.
        dtype='intc', shape=(prod(mesh),)
    grid_address : ndarray
        Address of all grid points. Each address is given by three unsigned
        integers.
        dtype='intc', shape=(prod(mesh), 3)

    """
    _set_no_error()

    if is_dense:
        dtype = 'uintp'
    else:
        dtype = 'intc'
    mapping_table = np.zeros(np.prod(mesh), dtype=dtype)
    grid_address = np.zeros((np.prod(mesh), 3), dtype='intc')
    if is_shift is None:
        is_shift = [0, 0, 0]
    if qpoints is None:
        qpoints = np.array([[0, 0, 0]], dtype='double', order='C')
    else:
        qpoints = np.array(qpoints, dtype='double', order='C')
        if qpoints.shape == (3,):
            qpoints = np.array([qpoints], dtype='double', order='C')

    if spg.stabilized_reciprocal_mesh(
            grid_address,
            mapping_table,
            np.array(mesh, dtype='intc'),
            np.array(is_shift, dtype='intc'),
            is_time_reversal * 1,
            np.array(rotations, dtype='intc', order='C'),
            qpoints) > 0:
        return mapping_table, grid_address
    else:
        return None

def get_grid_points_by_rotations(address_orig,
                                 reciprocal_rotations,
                                 mesh,
                                 is_shift=None,
                                 is_dense=False):
    """Returns grid points obtained after rotating input grid address

    Parameters
    ----------
    address_orig : array_like
        Grid point address to be rotated.
        dtype='intc', shape=(3,)
    reciprocal_rotations : array_like
        Rotation matrices {R} with respect to reciprocal basis vectors.
        Defined by q'=Rq.
        dtype='intc', shape=(rotations, 3, 3)
    mesh : array_like
        dtype='intc', shape=(3,)
    is_shift : array_like, optional
        With (1) or without (0) half grid shifts with respect to grid intervals
        sampled along reciprocal basis vectors. Default is None, which
        gives [0, 0, 0].
    is_dense : bool, optional
        rot_grid_points is returned with dtype='uintp' if True. Otherwise
        its dtype='intc'. Default is False.

    Returns
    -------
    rot_grid_points : ndarray
        Grid points obtained after rotating input grid address
        dtype='intc' or 'uintp', shape=(rotations,)

    """

    _set_no_error()

    if is_shift is None:
        _is_shift = np.zeros(3, dtype='intc')
    else:
        _is_shift = np.array(is_shift, dtype='intc')

    rot_grid_points = np.zeros(len(reciprocal_rotations), dtype='uintp')
    spg.grid_points_by_rotations(
        rot_grid_points,
        np.array(address_orig, dtype='intc'),
        np.array(reciprocal_rotations, dtype='intc', order='C'),
        np.array(mesh, dtype='intc'),
        _is_shift)

    if is_dense:
        return rot_grid_points
    else:
        return np.array(rot_grid_points, dtype='intc')

def get_BZ_grid_points_by_rotations(address_orig,
                                    reciprocal_rotations,
                                    mesh,
                                    bz_map,
                                    is_shift=None,
                                    is_dense=False):
    """Returns grid points obtained after rotating input grid address

    Parameters
    ----------
    address_orig : array_like
        Grid point address to be rotated.
        dtype='intc', shape=(3,)
    reciprocal_rotations : array_like
        Rotation matrices {R} with respect to reciprocal basis vectors.
        Defined by q'=Rq.
        dtype='intc', shape=(rotations, 3, 3)
    mesh : array_like
        dtype='intc', shape=(3,)
    is_shift : array_like, optional
        With (1) or without (0) half grid shifts with respect to grid intervals
        sampled along reciprocal basis vectors. Default is None, which
        gives [0, 0, 0].
    is_dense : bool, optional
        rot_grid_points is returned with dtype='uintp' if True. Otherwise
        its dtype='intc'. Default is False.

    Returns
    -------
    rot_grid_points : ndarray
        Grid points obtained after rotating input grid address
        dtype='intc' or 'uintp', shape=(rotations,)

    """

    _set_no_error()

    if is_shift is None:
        _is_shift = np.zeros(3, dtype='intc')
    else:
        _is_shift = np.array(is_shift, dtype='intc')

    if bz_map.dtype == 'uintp' and bz_map.flags.c_contiguous:
        _bz_map = bz_map
    else:
        _bz_map = np.array(bz_map, dtype='uintp')

    rot_grid_points = np.zeros(len(reciprocal_rotations), dtype='uintp')
    spg.BZ_grid_points_by_rotations(
        rot_grid_points,
        np.array(address_orig, dtype='intc'),
        np.array(reciprocal_rotations, dtype='intc', order='C'),
        np.array(mesh, dtype='intc'),
        _is_shift,
        _bz_map)

    if is_dense:
        return rot_grid_points
    else:
        return np.array(rot_grid_points, dtype='intc')

def relocate_BZ_grid_address(grid_address,
                             mesh,
                             reciprocal_lattice,  # column vectors
                             is_shift=None,
                             is_dense=False):
    """Grid addresses are relocated to be inside first Brillouin zone.

    Number of ir-grid-points inside Brillouin zone is returned.
    It is assumed that the following arrays have the shapes of
        bz_grid_address : (num_grid_points_in_FBZ, 3)
        bz_map (prod(mesh * 2), )

    Note that the shape of grid_address is (prod(mesh), 3) and the
    addresses in grid_address are arranged to be in parallelepiped
    made of reciprocal basis vectors. The addresses in bz_grid_address
    are inside the first Brillouin zone or on its surface. Each
    address in grid_address is mapped to one of those in
    bz_grid_address by a reciprocal lattice vector (including zero
    vector) with keeping element order. For those inside first
    Brillouin zone, the mapping is one-to-one. For those on the first
    Brillouin zone surface, more than one addresses in bz_grid_address
    that are equivalent by the reciprocal lattice translations are
    mapped to one address in grid_address. In this case, those grid
    points except for one of them are appended to the tail of this array,
    for which bz_grid_address has the following data storing:

    |------------------array size of bz_grid_address-------------------------|
    |--those equivalent to grid_address--|--those on surface except for one--|
    |-----array size of grid_address-----|

    Number of grid points stored in bz_grid_address is returned.
    bz_map is used to recover grid point index expanded to include BZ
    surface from grid address. The grid point indices are mapped to
    (mesh[0] * 2) x (mesh[1] * 2) x (mesh[2] * 2) space (bz_map).

    """
    _set_no_error()

    if is_shift is None:
        _is_shift = np.zeros(3, dtype='intc')
    else:
        _is_shift = np.array(is_shift, dtype='intc')
    bz_grid_address = np.zeros((np.prod(np.add(mesh, 1)), 3), dtype='intc')
    bz_map = np.zeros(np.prod(np.multiply(mesh, 2)), dtype='uintp')
    num_bz_ir = spg.BZ_grid_address(
        bz_grid_address,
        bz_map,
        grid_address,
        np.array(mesh, dtype='intc'),
        np.array(reciprocal_lattice, dtype='double', order='C'),
        _is_shift)

    if is_dense:
        return bz_grid_address[:num_bz_ir], bz_map
    else:
        return bz_grid_address[:num_bz_ir], np.array(bz_map, dtype='intc')

def delaunay_reduce(lattice, eps=1e-5):
    """Run Delaunay reduction

    Args:
        lattice: Lattice parameters in the form of
            [[a_x, a_y, a_z],
             [b_x, b_y, b_z],
             [c_x, c_y, c_z]]
        symprec:
            float: Tolerance to check if volume is close to zero or not and
                   if two basis vectors are orthogonal by the value of dot
                   product being close to zero or not.

    Returns:
        if the Delaunay reduction succeeded:
            Reduced lattice parameters are given as a numpy 'double' array:
            [[a_x, a_y, a_z],
             [b_x, b_y, b_z],
             [c_x, c_y, c_z]]
        otherwise None is returned.
    """
    _set_no_error()

    delaunay_lattice = np.array(np.transpose(lattice),
                                dtype='double', order='C')
    result = spg.delaunay_reduce(delaunay_lattice, float(eps))
    _set_error_message()

    if result == 0:
        return None
    else:
        return np.array(np.transpose(delaunay_lattice),
                        dtype='double', order='C')

def niggli_reduce(lattice, eps=1e-5):
    """Run Niggli reduction

    Args:
        lattice: Lattice parameters in the form of
            [[a_x, a_y, a_z],
             [b_x, b_y, b_z],
             [c_x, c_y, c_z]]
        eps:
            float: Tolerance to check if difference of norms of two basis
                   vectors is close to zero or not and if two basis vectors are
                   orthogonal by the value of dot product being close to zero or
                   not. The detail is shown at
                   https://atztogo.github.io/niggli/.

    Returns:
        if the Niggli reduction succeeded:
            Reduced lattice parameters are given as a numpy 'double' array:
            [[a_x, a_y, a_z],
             [b_x, b_y, b_z],
             [c_x, c_y, c_z]]
        otherwise None is returned.
    """
    _set_no_error()

    niggli_lattice = np.array(np.transpose(lattice), dtype='double', order='C')
    result = spg.niggli_reduce(niggli_lattice, float(eps))
    _set_error_message()

    if result == 0:
        return None
    else:
        return np.array(np.transpose(niggli_lattice),
                        dtype='double', order='C')

def get_error_message():
    return spglib_error.message

def _expand_cell(cell):
    if isinstance(cell, tuple):
        lattice = np.array(np.transpose(cell[0]), dtype='double', order='C')
        positions = np.array(cell[1], dtype='double', order='C')
        numbers = np.array(cell[2], dtype='intc')
        if len(cell) > 3:
            magmoms = np.array(cell[3], dtype='double')
        else:
            magmoms = None
    else:
        import warnings
        warnings.warn("ASE Atoms-like input is deprecated.",
                      DeprecationWarning)
        lattice = np.array(cell.get_cell().T, dtype='double', order='C')
        positions = np.array(cell.get_scaled_positions(),
                             dtype='double', order='C')
        numbers = np.array(cell.get_atomic_numbers(), dtype='intc')
        magmoms = None

    if _check(lattice, positions, numbers, magmoms):
        return (lattice, positions, numbers, magmoms)
    else:
        return (None, None, None, None)

def _check(lattice, positions, numbers, magmoms):
    if lattice.shape != (3, 3):
        return False
    if positions.ndim != 2:
        return False
    if positions.shape[1] != 3:
        return False
    if numbers.ndim != 1:
        return False
    if len(numbers) != positions.shape[0]:
        return False
    if magmoms is not None:
        if magmoms.ndim != 1:
            return False
        if len(magmoms) != len(numbers):
            return False
    return True

def _set_error_message():
    spglib_error.message = spg.error_message()

def _set_no_error():
    spglib_error.message = "no error"


from numbers import Integral, Real
from itertools import chain
import string

from six import string_types
import matplotlib.pyplot as plt
import numpy as np

import openmc.checkvalue as cv
import openmc.data

# Supported keywords for continuous-energy cross section plotting
PLOT_TYPES = ['total', 'scatter', 'elastic', 'inelastic', 'fission',
              'absorption', 'capture', 'nu-fission', 'nu-scatter', 'unity',
              'slowing-down power', 'damage']

# Supported keywoards for multi-group cross section plotting
PLOT_TYPES_MGXS = ['total', 'absorption', 'scatter', 'fission',
                   'kappa-fission', 'nu-fission', 'prompt-nu-fission',
                   'deleyed-nu-fission', 'chi', 'chi-prompt', 'chi-delayed',
                   'inverse-velocity', 'beta', 'decay rate', 'unity']
# Create a dictionary which can be used to convert PLOT_TYPES_MGXS to the
# openmc.XSdata attribute name needed to access the data
_PLOT_MGXS_ATTR = {line: line.replace(' ', '_').replace('-', '_')
                   for line in PLOT_TYPES_MGXS}
_PLOT_MGXS_ATTR['scatter'] = 'scatter_matrix'

# Special MT values
UNITY_MT = -1
XI_MT = -2

# MTs to combine to generate associated plot_types
_INELASTIC = [mt for mt in openmc.data.SUM_RULES[3] if mt != 27]
PLOT_TYPES_MT = {'total': openmc.data.SUM_RULES[1],
                 'scatter': [2] + _INELASTIC,
                 'elastic': [2],
                 'inelastic': _INELASTIC,
                 'fission': [18],
                 'absorption': [27], 'capture': [101],
                 'nu-fission': [18],
                 'nu-scatter': [2] + _INELASTIC,
                 'unity': [UNITY_MT],
                 'slowing-down power': [2] + _INELASTIC + [XI_MT],
                 'damage': [444]}
# Operations to use when combining MTs the first np.add is used in reference
# to zero
PLOT_TYPES_OP = {'total': (np.add,),
                 'scatter': (np.add,) * (len(PLOT_TYPES_MT['scatter']) - 1),
                 'elastic': (),
                 'inelastic': (np.add,) * (len(PLOT_TYPES_MT['inelastic']) - 1),
                 'fission': (), 'absorption': (),
                 'capture': (), 'nu-fission': (),
                 'nu-scatter': (np.add,) * (len(PLOT_TYPES_MT['nu-scatter']) - 1),
                 'unity': (),
                 'slowing-down power':
                    (np.add,) * (len(PLOT_TYPES_MT['slowing-down power']) - 2) + (np.multiply,),
                 'damage': ()}

# Types of plots to plot linearly in y
PLOT_TYPES_LINEAR = {'nu-fission / fission', 'nu-scatter / scatter',
                     'nu-fission / absorption', 'fission / absorption'}

# Minimum and maximum energies for plotting (units of eV)
_MIN_E = 1.e-5
_MAX_E = 20.e6

def plot_xs(this, types, divisor_types=None, temperature=294., data_type=None,
            axis=None, sab_name=None, ce_cross_sections=None,
            mg_cross_sections=None, enrichment=None, plot_CE=True, orders=None,
            divisor_orders=None, **kwargs):
    """Creates a figure of continuous-energy cross sections for this item.

    Parameters
    ----------
    this : str or openmc.Material
        Object to source data from
    types : Iterable of values of PLOT_TYPES
        The type of cross sections to include in the plot.
    divisor_types : Iterable of values of PLOT_TYPES, optional
        Cross section types which will divide those produced by types
        before plotting. A type of 'unity' can be used to effectively not
        divide some types.
    temperature : float, optional
        Temperature in Kelvin to plot. If not specified, a default
        temperature of 294K will be plotted. Note that the nearest
        temperature in the library for each nuclide will be used as opposed
        to using any interpolation.
    data_type : {'nuclide', 'element', 'material', 'macroscopic'}, optional
        Type of object to plot. If not specified, a guess is made based on the
        `this` argument.
    axis : matplotlib.axes, optional
        A previously generated axis to use for plotting. If not specified,
        a new axis and figure will be generated.
    sab_name : str, optional
        Name of S(a,b) library to apply to MT=2 data when applicable; only used
        for items which are instances of openmc.Element or openmc.Nuclide
    ce_cross_sections : str, optional
        Location of cross_sections.xml file. Default is None.
    mg_cross_sections : str, optional
        Location of MGXS HDF5 Library file. Default is None.
    enrichment : float, optional
        Enrichment for U235 in weight percent. For example, input 4.95 for
        4.95 weight percent enriched U. Default is None. This is only used for
        items which are instances of openmc.Element
    plot_CE : bool, optional
        Denotes whether or not continuous-energy will be plotted. Defaults to
        plotting the continuous-energy data.
    orders : Iterable of Integral, optional
        The scattering order or delayed group index to use for the
        corresponding entry in types. Defaults to the 0th order for scattering
        and the total delayed neutron data. This only applies to plots of
        multi-group data.
    divisor_orders : Iterable of Integral, optional
        Same as orders, but for divisor_types
    **kwargs
        All keyword arguments are passed to
        :func:`matplotlib.pyplot.figure`.

    Returns
    -------
    fig : matplotlib.figure.Figure
        If axis is None, then a Matplotlib Figure of the generated
        cross section will be returned. Otherwise, a value of
        None will be returned as the figure and axes have already been
        generated.

    """
    cv.check_type("plot_CE", plot_CE, bool)

    if data_type is None:
        if isinstance(this, openmc.Nuclide):
            data_type = 'nuclide'
        elif isinstance(this, openmc.Element):
            data_type = 'element'
        elif isinstance(this, openmc.Material):
            data_type = 'material'
        elif isinstance(this, openmc.Macroscopic):
            data_type = 'macroscopic'
        elif isinstance(this, string_types):
            if this[-1] in string.digits:
                data_type = 'nuclide'
            else:
                data_type = 'element'
        else:
            raise TypeError("Invalid type for plotting")

    if plot_CE:
        # Calculate for the CE cross sections
        E, data = calculate_cexs(this, data_type, types, temperature, sab_name,
                                 ce_cross_sections, enrichment)
        if divisor_types:
            cv.check_length('divisor types', divisor_types, len(types))
            Ediv, data_div = calculate_cexs(this, divisor_types, temperature,
                                            sab_name, ce_cross_sections,
                                            enrichment)

            # Create a new union grid, interpolate data and data_div on to that
            # grid, and then do the actual division
            Enum = E[:]
            E = np.union1d(Enum, Ediv)
            data_new = np.zeros((len(types), len(E)))

            for line in range(len(types)):
                data_new[line, :] = \
                    np.divide(np.interp(E, Enum, data[line, :]),
                              np.interp(E, Ediv, data_div[line, :]))
                if divisor_types[line] != 'unity':
                    types[line] = types[line] + ' / ' + divisor_types[line]
            data = data_new
    else:
        # Calculate for MG cross sections
        E, data = calculate_mgxs(this, types, orders, temperature,
                                 mg_cross_sections, ce_cross_sections,
                                 enrichment)

        if divisor_types:
            cv.check_length('divisor types', divisor_types, len(types))
            Ediv, data_div = calculate_mgxs(this, divisor_types,
                                            divisor_orders, temperature,
                                            mg_cross_sections,
                                            ce_cross_sections, enrichment)

            # Perform the division
            for line in range(len(types)):
                data[line, :] /= data_div[line, :]
                if divisor_types[line] != 'unity':
                    types[line] += ' / ' + divisor_types[line]

    # Generate the plot
    if axis is None:
        fig = plt.figure(**kwargs)
        ax = fig.add_subplot(111)
    else:
        fig = None
        ax = axis
    # Set to loglog or semilogx depending on if we are plotting a data
    # type which we expect to vary linearly
    if set(types).issubset(PLOT_TYPES_LINEAR):
        plot_func = ax.semilogx
    else:
        plot_func = ax.loglog

    # Plot the data
    for i in range(len(data)):
        data[i, :] = np.nan_to_num(data[i, :])
        if np.sum(data[i, :]) > 0.:
            plot_func(E, data[i, :], label=types[i])

    ax.set_xlabel('Energy [eV]')
    if plot_CE:
        ax.set_xlim(_MIN_E, _MAX_E)
    else:
        ax.set_xlim(E[-1], E[0])
    if divisor_types:
        if data_type == 'nuclide':
            ylabel = 'Nuclidic Microscopic Data'
        elif data_type == 'element':
            ylabel = 'Elemental Microscopic Data'
        elif data_type == 'material' or data_type == 'macroscopic':
            ylabel = 'Macroscopic Data'
    else:
        if data_type == 'nuclide':
            ylabel = 'Microscopic Cross Section [b]'
        elif data_type == 'element':
            ylabel = 'Elemental Cross Section [b]'
        elif data_type == 'material' or data_type == 'macroscopic':
            ylabel = 'Macroscopic Cross Section [1/cm]'
    ax.set_ylabel(ylabel)
    ax.legend(loc='best')
    name = this.name if data_type == 'material' else this
    if len(types) > 1:
        ax.set_title('Cross Sections for ' + name)
    else:
        ax.set_title('Cross Section for ' + name)

    return fig

def calculate_cexs(this, data_type, types, temperature=294., sab_name=None,
                   cross_sections=None, enrichment=None):
    """Calculates continuous-energy cross sections of a requested type.

    Parameters
    ----------
    this : str or openmc.Material
        Object to source data from
    data_type : {'nuclide', 'element', material'}
        Type of object to plot
    types : Iterable of values of PLOT_TYPES
        The type of cross sections to calculate
    temperature : float, optional
        Temperature in Kelvin to plot. If not specified, a default
        temperature of 294K will be plotted. Note that the nearest
        temperature in the library for each nuclide will be used as opposed
        to using any interpolation.
    sab_name : str, optional
        Name of S(a,b) library to apply to MT=2 data when applicable.
    cross_sections : str, optional
        Location of cross_sections.xml file. Default is None.
    enrichment : float, optional
        Enrichment for U235 in weight percent. For example, input 4.95 for
        4.95 weight percent enriched U. Default is None
        (natural composition).

    Returns
    -------
    energy_grid : numpy.ndarray
        Energies at which cross sections are calculated, in units of eV
    data : numpy.ndarray
        Cross sections calculated at the energy grid described by energy_grid

    """

    # Check types
    cv.check_type('temperature', temperature, Real)
    if sab_name:
        cv.check_type('sab_name', sab_name, string_types)
    if enrichment:
        cv.check_type('enrichment', enrichment, Real)

    if data_type == 'nuclide':
        energy_grid, xs = _calculate_cexs_nuclide(this, types, temperature,
                                                  sab_name, cross_sections)
        # Convert xs (Iterable of Callable) to a grid of cross section values
        # calculated on @ the points in energy_grid for consistency with the
        # element and material functions.
        data = np.zeros((len(types), len(energy_grid)))
        for line in range(len(types)):
            data[line, :] = xs[line](energy_grid)
    elif data_type == 'element':
        energy_grid, data = _calculate_cexs_elem_mat(this, types, temperature,
                                                     cross_sections, sab_name,
                                                     enrichment)
    elif data_type == 'material':
        energy_grid, data = _calculate_cexs_elem_mat(this, types, temperature,
                                                     cross_sections)
    else:
        raise TypeError("Invalid type")

    return energy_grid, data

def _calculate_cexs_nuclide(this, types, temperature=294., sab_name=None,
                            cross_sections=None):
    """Calculates continuous-energy cross sections of a requested type.

    Parameters
    ----------
    this : openmc.Nuclide
        Nuclide object to source data from
    types : Iterable of str or Integral
        The type of cross sections to calculate; values can either be those
        in openmc.PLOT_TYPES or integers which correspond to reaction
        channel (MT) numbers.
    temperature : float, optional
        Temperature in Kelvin to plot. If not specified, a default
        temperature of 294K will be plotted. Note that the nearest
        temperature in the library for each nuclide will be used as opposed
        to using any interpolation.
    sab_name : str, optional
        Name of S(a,b) library to apply to MT=2 data when applicable.
    cross_sections : str, optional
        Location of cross_sections.xml file. Default is None.

    Returns
    -------
    energy_grid : numpy.ndarray
        Energies at which cross sections are calculated, in units of eV
    data : Iterable of Callable
        Requested cross section functions

    """

    # Parse the types
    mts = []
    ops = []
    yields = []
    for line in types:
        if line in PLOT_TYPES:
            mts.append(PLOT_TYPES_MT[line])
            if line.startswith('nu'):
                yields.append(True)
            else:
                yields.append(False)
            ops.append(PLOT_TYPES_OP[line])
        else:
            # Not a built-in type, we have to parse it ourselves
            cv.check_type('MT in types', line, Integral)
            cv.check_greater_than('MT in types', line, 0)
            mts.append((line,))
            ops.append(())
            yields.append(False)

    # Load the library
    library = openmc.data.DataLibrary.from_xml(cross_sections)

    # Convert temperature to format needed for access in the library
    strT = "{}K".format(int(round(temperature)))
    T = temperature

    # Now we can create the data sets to be plotted
    energy_grid = []
    xs = []
    lib = library.get_by_material(this)
    if lib is not None:
        nuc = openmc.data.IncidentNeutron.from_hdf5(lib['path'])
        # Obtain the nearest temperature
        if strT in nuc.temperatures:
            nucT = strT
        else:
            delta_T = np.array(nuc.kTs) - T * openmc.data.K_BOLTZMANN
            closest_index = np.argmin(np.abs(delta_T))
            nucT = nuc.temperatures[closest_index]

        # Prep S(a,b) data if needed
        if sab_name:
            sab = openmc.data.ThermalScattering.from_hdf5(sab_name)
            # Obtain the nearest temperature
            if strT in sab.temperatures:
                sabT = strT
            else:
                delta_T = np.array(sab.kTs) - T * openmc.data.K_BOLTZMANN
                closest_index = np.argmin(np.abs(delta_T))
                sabT = sab.temperatures[closest_index]

            # Create an energy grid composed the S(a,b) and the nuclide's grid
            grid = nuc.energy[nucT]
            sab_Emax = 0.
            sab_funcs = []
            if sab.elastic_xs:
                elastic = sab.elastic_xs[sabT]
                if isinstance(elastic, openmc.data.CoherentElastic):
                    grid = np.union1d(grid, elastic.bragg_edges)
                    if elastic.bragg_edges[-1] > sab_Emax:
                        sab_Emax = elastic.bragg_edges[-1]
                elif isinstance(elastic, openmc.data.Tabulated1D):
                    grid = np.union1d(grid, elastic.x)
                    if elastic.x[-1] > sab_Emax:
                        sab_Emax = elastic.x[-1]
                sab_funcs.append(elastic)
            if sab.inelastic_xs:
                inelastic = sab.inelastic_xs[sabT]
                grid = np.union1d(grid, inelastic.x)
                if inelastic.x[-1] > sab_Emax:
                        sab_Emax = inelastic.x[-1]
                sab_funcs.append(inelastic)
            energy_grid = grid
        else:
            energy_grid = nuc.energy[nucT]

        for i, mt_set in enumerate(mts):
            # Get the reaction xs data from the nuclide
            funcs = []
            op = ops[i]
            for mt in mt_set:
                if mt == 2:
                    if sab_name:
                        # Then we need to do a piece-wise function of
                        # The S(a,b) and non-thermal data
                        sab_sum = openmc.data.Sum(sab_funcs)
                        pw_funcs = openmc.data.Regions1D(
                            [sab_sum, nuc[mt].xs[nucT]],
                            [sab_Emax])
                        funcs.append(pw_funcs)
                    else:
                        funcs.append(nuc[mt].xs[nucT])
                elif mt in nuc:
                    if yields[i]:
                        # Get the total yield first if available. This will be
                        # used primarily for fission.
                        for prod in chain(nuc[mt].products,
                                          nuc[mt].derived_products):
                            if prod.particle == 'neutron' and \
                                prod.emission_mode == 'total':
                                func = openmc.data.Combination(
                                    [nuc[mt].xs[nucT], prod.yield_],
                                    [np.multiply])
                                funcs.append(func)
                                break
                        else:
                            # Total doesn't exist so we have to create from
                            # prompt and delayed. This is used for scatter
                            # multiplication.
                            func = None
                            for prod in chain(nuc[mt].products,
                                              nuc[mt].derived_products):
                                if prod.particle == 'neutron' and \
                                    prod.emission_mode != 'total':
                                    if func:
                                        func = openmc.data.Combination(
                                            [prod.yield_, func], [np.add])
                                    else:
                                        func = prod.yield_
                            if func:
                                funcs.append(openmc.data.Combination(
                                    [func, nuc[mt].xs[nucT]], [np.multiply]))
                            else:
                                # If func is still None, then there were no
                                # products. In that case, assume the yield is
                                # one as its not provided for some summed
                                # reactions like MT=4
                                funcs.append(nuc[mt].xs[nucT])
                    else:
                        funcs.append(nuc[mt].xs[nucT])
                elif mt == UNITY_MT:
                    funcs.append(lambda x: 1.)
                elif mt == XI_MT:
                    awr = nuc.atomic_weight_ratio
                    alpha = ((awr - 1.) / (awr + 1.))**2
                    xi = 1. + alpha * np.log(alpha) / (1. - alpha)
                    funcs.append(lambda x: xi)
                else:
                    funcs.append(lambda x: 0.)
            xs.append(openmc.data.Combination(funcs, op))
    else:
        raise ValueError(this + " not in library")

    return energy_grid, xs

def _calculate_cexs_elem_mat(this, types, temperature=294.,
                             cross_sections=None, sab_name=None,
                             enrichment=None):
    """Calculates continuous-energy cross sections of a requested type.

    Parameters
    ----------
    this : openmc.Material or openmc.Element
        Object to source data from
    types : Iterable of values of PLOT_TYPES
        The type of cross sections to calculate
    temperature : float, optional
        Temperature in Kelvin to plot. If not specified, a default
        temperature of 294K will be plotted. Note that the nearest
        temperature in the library for each nuclide will be used as opposed
        to using any interpolation.
    cross_sections : str, optional
        Location of cross_sections.xml file. Default is None.
    sab_name : str, optional
        Name of S(a,b) library to apply to MT=2 data when applicable.
    enrichment : float, optional
        Enrichment for U235 in weight percent. For example, input 4.95 for
        4.95 weight percent enriched U. Default is None
        (natural composition).

    Returns
    -------
    energy_grid : numpy.ndarray
        Energies at which cross sections are calculated, in units of eV
    data : numpy.ndarray
        Cross sections calculated at the energy grid described by energy_grid

    """

    if isinstance(this, openmc.Material):
        if this.temperature is not None:
            T = this.temperature
        else:
            T = temperature
        data_type = 'material'
    else:
        T = temperature
        data_type = 'element'

    # Load the library
    library = openmc.data.DataLibrary.from_xml(cross_sections)

    if isinstance(this, openmc.Material):
        # Expand elements in to nuclides with atomic densities
        nuclides = this.get_nuclide_atom_densities()
        # For ease of processing split out the nuclide and its fraction
        nuc_fractions = {nuclide[1][0]: nuclide[1][1]
                         for nuclide in nuclides.items()}
        # Create a dict of [nuclide name] = nuclide object to carry forward
        # with a common nuclides format between openmc.Material and
        # openmc.Element objects
        nuclides = {nuclide[1][0]: nuclide[1][0]
                    for nuclide in nuclides.items()}
    else:
        # Expand elements in to nuclides with atomic densities
        nuclides = this.expand(1., 'ao', enrichment=enrichment,
                               cross_sections=cross_sections)
        # For ease of processing split out the nuclide and its fraction
        nuc_fractions = {nuclide[0]: nuclide[1] for nuclide in nuclides}
        # Create a dict of [nuclide name] = nuclide object to carry forward
        # with a common nuclides format between openmc.Material and
        # openmc.Element objects
        nuclides = {nuclide[0]: nuclide[0] for nuclide in nuclides}

    # Identify the nuclides which have S(a,b) data
    sabs = {}
    for nuclide in nuclides.items():
        sabs[nuclide[0]] = None
    if isinstance(this, openmc.Material):
        for sab_name in this._sab:
            sab = openmc.data.ThermalScattering.from_hdf5(
                library.get_by_material(sab_name)['path'])
            for nuc in sab.nuclides:
                sabs[nuc] = library.get_by_material(sab_name)['path']
    else:
        if sab_name:
            sab = openmc.data.ThermalScattering.from_hdf5(sab_name)
            for nuc in sab.nuclides:
                sabs[nuc] = library.get_by_material(sab_name)['path']

    # Now we can create the data sets to be plotted
    xs = {}
    E = []
    for nuclide in nuclides.items():
        name = nuclide[0]
        nuc = nuclide[1]
        sab_tab = sabs[name]
        temp_E, temp_xs = calculate_cexs(nuc, data_type, types, T, sab_tab,
                                         cross_sections)
        E.append(temp_E)
        # Since the energy grids are different, store the cross sections as
        # a tabulated function so they can be calculated on any grid needed.
        xs[name] = [openmc.data.Tabulated1D(temp_E, temp_xs[line])
                    for line in range(len(types))]

    # Condense the data for every nuclide
    # First create a union energy grid
    energy_grid = E[0]
    for grid in E[1:]:
        energy_grid = np.union1d(energy_grid, grid)

    # Now we can combine all the nuclidic data
    data = np.zeros((len(types), len(energy_grid)))
    for line in range(len(types)):
        if types[line] == 'unity':
            data[line, :] = 1.
        else:
            for nuclide in nuclides.items():
                name = nuclide[0]
                data[line, :] += (nuc_fractions[name] *
                                  xs[name][line](energy_grid))

    return energy_grid, data

def calculate_mgxs(this, data_type, types, orders=None, temperature=294.,
                   cross_sections=None, ce_cross_sections=None,
                   enrichment=None):
    """Calculates multi-group cross sections of a requested type.

    If the data for the nuclide or macroscopic object in the library is
    represented as angle-dependent data then this method will return the
    geometric average cross section over all angles.

    Parameters
    ----------
    this : str or openmc.Material
        Object to source data from
    data_type : {'nuclide', 'element', material', 'macroscopic'}
        Type of object to plot
    types : Iterable of values of PLOT_TYPES_MGXS
        The type of cross sections to calculate
    orders : Iterable of Integral, optional
        The scattering order or delayed group index to use for the
        corresponding entry in types. Defaults to the 0th order for scattering
        and the total delayed neutron data.
    temperature : float, optional
        Temperature in Kelvin to plot. If not specified, a default
        temperature of 294K will be plotted. Note that the nearest
        temperature in the library for each nuclide will be used as opposed
        to using any interpolation.
    cross_sections : str, optional
        Location of MGXS HDF5 Library file. Default is None.
    ce_cross_sections : str, optional
        Location of continuous-energy cross_sections.xml file. Default is None.
        This is used only for expanding an openmc.Element object passed as this
    enrichment : float, optional
        Enrichment for U235 in weight percent. For example, input 4.95 for
        4.95 weight percent enriched U. Default is None
        (natural composition).

    Returns
    -------
    energy_grid : numpy.ndarray
        Energies at which cross sections are calculated, in units of eV
    data : numpy.ndarray
        Cross sections calculated at the energy grid described by energy_grid

    """

    # Check types
    cv.check_type('temperature', temperature, Real)
    if enrichment:
        cv.check_type('enrichment', enrichment, Real)
    cv.check_iterable_type('types', types, string_types)

    cv.check_type("cross_sections", cross_sections, str)
    library = openmc.MGXSLibrary.from_hdf5(cross_sections)

    if data_type in ('nuclide', 'macroscopic'):
        mgxs = _calculate_mgxs_nuc_macro(this, types, library, orders,
                                         temperature)
    elif data_type in ('element', 'material'):
        mgxs = _calculate_mgxs_elem_mat(this, types, library, orders,
                                        temperature, ce_cross_sections,
                                        enrichment)
    else:
        raise TypeError("Invalid type")

    # Convert the data to the format needed
    data = np.zeros((len(types), 2 * library.energy_groups.num_groups))
    energy_grid = np.zeros(2 * library.energy_groups.num_groups)
    for g in range(library.energy_groups.num_groups):
        energy_grid[g * 2: g * 2 + 2] = \
            library.energy_groups.group_edges[g: g + 2]
    # Ensure the energy will show on a log-axis by replacing 0s with a
    # sufficiently small number
    energy_grid[0] = max(energy_grid[0], _MIN_E)

    for line in range(len(types)):
        for g in range(library.energy_groups.num_groups):
            data[line, g * 2: g * 2 + 2] = mgxs[line, g]

    return energy_grid[::-1], data

def _calculate_mgxs_nuc_macro(this, types, library, orders=None,
                              temperature=294.):
    """Determines the multi-group cross sections of a nuclide or macroscopic
    object.

    If the data for the nuclide or macroscopic object in the library is
    represented as angle-dependent data then this method will return the
    geometric average cross section over all angles.

    Parameters
    ----------
    this : openmc.Nuclide or openmc.Macroscopic
        Object to source data from
    types : Iterable of str
        The type of cross sections to calculate; values can either be those
        in openmc.PLOT_TYPES_MGXS
    library : openmc.MGXSLibrary
        MGXS Library containing the data of interest
    orders : Iterable of Integral, optional
        The scattering order or delayed group index to use for the
        corresponding entry in types. Defaults to the 0th order for scattering
        and the total delayed neutron data.
    temperature : float, optional
        Temperature in Kelvin to plot. If not specified, a default
        temperature of 294K will be plotted. Note that the nearest
        temperature in the library for each nuclide will be used as opposed
        to using any interpolation.

    Returns
    -------
    data : numpy.ndarray
        Cross sections calculated at the energy grid described by energy_grid

    """

    # Check the parameters and grab order/delayed groups
    if orders:
        cv.check_iterable_type('orders', orders, Integral,
                               min_depth=len(types), max_depth=len(types))
    else:
        orders = [None] * len(types)
    for i, line in enumerate(types):
        cv.check_type("line", line, str)
        cv.check_value("line", line, PLOT_TYPES_MGXS)
        if orders[i]:
            cv.check_greater_than("order value", orders[i], 0, equality=True)

    xsdata = library.get_by_name(this)

    if xsdata is not None:
        # Obtain the nearest temperature
        t = np.abs(xsdata.temperatures - temperature).argmin()

        # Get the data
        data = np.zeros((len(types), library.energy_groups.num_groups))
        for i, line in enumerate(types):
            if 'fission' in line and not xsdata.fissionable:
                continue
            elif line == 'unity':
                data[i, :] = 1.
            else:
                # Now we have to get the cross section data and properly
                # treat it depending on the requested type.
                # First get the data in a generic fashion
                temp_data = getattr(xsdata, _PLOT_MGXS_ATTR[line])[t]
                shape = temp_data.shape[:]
                # If we have angular data, then want the geometric
                # average over all provided angles.  Since the angles are
                # equi-distant, un-weighted averaging will suffice
                if xsdata.representation == 'angle':
                    temp_data = np.mean(temp_data, axis=(0, 1))

                # Now we can look at the shape of the data to identify how
                # it should be modified to produce an array of values
                # with groups.
                if shape in (xsdata.xs_shapes["[G']"],
                             xsdata.xs_shapes["[G]"]):
                    # Then the data is already an array vs groups so copy
                    # and move along
                    data[i, :] = temp_data
                elif shape == xsdata.xs_shapes["[G][G']"]:
                    # Sum the data over outgoing groups to create our array vs
                    # groups
                    data[i, :] = np.sum(temp_data, axis=1)
                elif shape == xsdata.xs_shapes["[DG]"]:
                    # Then we have a constant vs groups with a value for each
                    # delayed group. The user-provided value of orders tells us
                    # which delayed group we want. If none are provided, then
                    # we sum all the delayed groups together.
                    if orders[i]:
                        if orders[i] < len(shape[0]):
                            data[i, :] = temp_data[orders[i]]
                    else:
                        data[i, :] = np.sum(temp_data[:])
                elif shape in (xsdata.xs_shapes["[DG][G']"],
                               xsdata.xs_shapes["[DG][G]"]):
                    # Then we have an array vs groups with values for each
                    # delayed group. The user-provided value of orders tells us
                    # which delayed group we want. If none are provided, then
                    # we sum all the delayed groups together.
                    if orders[i]:
                        if orders[i] < len(shape[0]):
                            data[i, :] = temp_data[orders[i], :]
                    else:
                        data[i, :] = np.sum(temp_data[:, :], axis=0)
                elif shape == xsdata.xs_shapes["[DG][G][G']"]:
                    # Then we have a delayed group matrix. We will first
                    # remove the outgoing group dependency
                    temp_data = np.sum(temp_data, axis=-1)
                    # And then proceed in exactly the same manner as the
                    # "[DG][G']" or "[DG][G]" shapes in the previous block.
                    if orders[i]:
                        if orders[i] < len(shape[0]):
                            data[i, :] = temp_data[orders[i], :]
                    else:
                        data[i, :] = np.sum(temp_data[:, :], axis=0)
                elif shape == xsdata.xs_shapes["[G][G'][Order]"]:
                    # This is a scattering matrix with angular data
                    # First remove the outgoing group dependence
                    temp_data = np.sum(temp_data, axis=1)
                    # The user either provided a specific order or we resort
                    # to the default 0th order
                    if orders[i]:
                        order = orders[i]
                    else:
                        order = 0
                    # If the order is available, store the data for that order
                    # if it is not available, then the expansion coefficient
                    # is zero and thus we already have the correct value.
                    if order < shape[1]:
                        data[i, :] = temp_data[:, order]
    else:
        raise ValueError("{} not present in provided MGXS "
                         "library".format(this))

    return data

def _calculate_mgxs_elem_mat(this, types, library, orders=None,
                             temperature=294., ce_cross_sections=None,
                             enrichment=None):
    """Determines the multi-group cross sections of an element or material
    object.

    If the data for the nuclide or macroscopic object in the library is
    represented as angle-dependent data then this method will return the
    geometric average cross section over all angles.

    Parameters
    ----------
    this : openmc.Element or openmc.Material
        Object to source data from
    types : Iterable of str
        The type of cross sections to calculate; values can either be those
        in openmc.PLOT_TYPES_MGXS
    library : openmc.MGXSLibrary
        MGXS Library containing the data of interest
    orders : Iterable of Integral, optional
        The scattering order or delayed group index to use for the
        corresponding entry in types. Defaults to the 0th order for scattering
        and the total delayed neutron data.
    temperature : float, optional
        Temperature in Kelvin to plot. If not specified, a default
        temperature of 294K will be plotted. Note that the nearest
        temperature in the library for each nuclide will be used as opposed
        to using any interpolation.
    ce_cross_sections : str, optional
        Location of continuous-energy cross_sections.xml file. Default is None.
        This is used only for expanding the elements
    enrichment : float, optional
        Enrichment for U235 in weight percent. For example, input 4.95 for
        4.95 weight percent enriched U. Default is None
        (natural composition).

    Returns
    -------
    data : numpy.ndarray
        Cross sections calculated at the energy grid described by energy_grid

    """

    if isinstance(this, openmc.Material):
        if this.temperature is not None:
            T = this.temperature
        else:
            T = temperature

        # Check to see if we have nuclides/elements or a macrocopic object
        if this._macroscopic is not None:
            # We have macroscopics
            nuclides = {this._macroscopic: (this._macroscopic, this.density)}
        else:
            # Expand elements in to nuclides with atomic densities
            nuclides = this.get_nuclide_atom_densities()

        # For ease of processing split out nuc and nuc_density
        nuc_fraction = [nuclide[1][1] for nuclide in nuclides.items()]
    else:
        T = temperature
        # Expand elements in to nuclides with atomic densities
        nuclides = this.expand(100., 'ao', enrichment=enrichment,
                               cross_sections=ce_cross_sections)

        # For ease of processing split out nuc and nuc_fractions
        nuc_fraction = [nuclide[1] for nuclide in nuclides]

    nuc_data = []
    for nuclide in nuclides.items():
        nuc_data.append(_calculate_mgxs_nuc_macro(nuclide[0], types, library,
                                                  orders, T))

    # Combine across the nuclides
    data = np.zeros((len(types), library.energy_groups.num_groups))
    for line in range(len(types)):
        if types[line] == 'unity':
            data[line, :] = 1.
        else:
            for n in range(len(nuclides)):
                data[line, :] += nuc_fraction[n] * nuc_data[n][line, :]

    return data


"""
This is a Python implementation of the Denoising Autoencoder approach that we proposed for
the first Multi-target speaker detection and identification Challenge Evaluation (MCE 2018, http://www.mce2018.org)

The basic idea is to train a Denoising Autoencoder to map each individual input ivector
to the mean of all ivectors from that speaker.
The aim of this DAE is to compensate for inter-session variability and increase the discriminative power of the ivectors.

You can find our system description for the MCE 2018 challenge here: http://mce.csail.mit.edu/pdfs/BiometricVox_description.pdf

Part of the code has been adapted from the baseline system at: https://github.com/swshon/multi-speakerID.

 Copyright 2018 Roberto Font
				Biometric Vox S.L.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

from __future__ import print_function
import numpy as np

from mce_utils import load_ivector, length_norm, make_spkvec, calculate_EER, get_trials_label_with_confusion, calculate_EER_with_confusion

import pandas as pd
from keras.models import Model
from keras.layers import Input, Dense, Activation
from keras import metrics
from keras import optimizers

# Neural network definition: a single hidden layer with 'tanh' activation and a linear output layer
def get_DAE(nu=2000):
  iv_dim = 600
  inputs = Input(shape=(iv_dim,))
  x = Dense(nu)(inputs)
  x = Activation('tanh')(x)
  x = Dense(iv_dim)(x)
  out = Activation('linear')(x)
  model = Model(inputs=inputs, outputs=out)

  return model

# Making dictionary to find blacklist pair between train and test dataset
bl_match = np.loadtxt('data/bl_matching.csv',dtype='str')
dev2train={}
dev2id={}
train2dev={}
train2id={}
test2train={}
train2test={}
for iter, line in enumerate(bl_match):
    line_s = line.split(',')
    dev2train[line_s[1].split('_')[-1]]= line_s[3].split('_')[-1]
    dev2id[line_s[1].split('_')[-1]]= line_s[0].split('_')[-1]
    train2dev[line_s[3].split('_')[-1]]= line_s[1].split('_')[-1]
    train2id[line_s[3].split('_')[-1]]= line_s[0].split('_')[-1]
    test2train[line_s[2].split('_')[-1]]= line_s[3].split('_')[-1]
    train2test[line_s[3].split('_')[-1]]= line_s[2].split('_')[-1]

# load test set information
filename = 'data/tst_evaluation_keys.csv'
tst_info = np.loadtxt(filename,dtype='str',delimiter=',',skiprows=1,usecols=range(0,3))
tst_trials = []
tst_trials_label = []
tst_ground_truth =[]
for iter in range(len(tst_info)):
    tst_trials_label.extend([tst_info[iter,0]])
    if tst_info[iter,1]=='background':
        tst_trials = np.append(tst_trials,0)

    else:
        tst_trials = np.append(tst_trials,1)

# Set random seed to make results reproducible
seed = 134
np.random.seed(seed)

# Loading i-vector
trn_bl_id, trn_bl_utt, trn_bl_ivector = load_ivector('data/trn_blacklist.csv')
trn_bg_id, trn_bg_utt, trn_bg_ivector = load_ivector('data/trn_background.csv')
dev_bl_id, dev_bl_utt, dev_bl_ivector = load_ivector('data/dev_blacklist.csv')
dev_bg_id, dev_bg_utt, dev_bg_ivector = load_ivector('data/dev_background.csv')
tst_id, test_utt, tst_ivector = load_ivector('data/tst_evaluation.csv')

# length normalization
trn_bl_ivector = length_norm(trn_bl_ivector)
trn_bg_ivector = length_norm(trn_bg_ivector)
dev_bl_ivector = length_norm(dev_bl_ivector)
dev_bg_ivector = length_norm(dev_bg_ivector)
tst_ivector = length_norm(tst_ivector)

# Inputs to DAE are ivectors and targets are the speaker-level mean ivectors
train_spk_ids = pd.DataFrame({'spk_ids': trn_bg_id})
train_ivs = pd.DataFrame(trn_bg_ivector)

X_train = train_ivs.values
Y_train = (train_ivs.groupby(train_spk_ids['spk_ids']).transform('mean')).values

# DAE training
model = get_DAE()

model.compile(loss='cosine_proximity',
              optimizer = optimizers.Adam(lr=0.001, beta_1=0.9, beta_2=0.999, epsilon=1e-08, decay=1e-06),
              metrics=[metrics.mean_squared_error])

num_examples = X_train.shape[0]
num_epochs = 5
batch_size = 512
num_batch_per_epoch = num_examples / batch_size

model.fit(x=X_train,y=Y_train,batch_size=batch_size,epochs=num_epochs)

# Compute DAE-transformed embeddings from ivectors
trn_bl_embeddings = model.predict(trn_bl_ivector,batch_size=batch_size)
trn_bg_embeddings = model.predict(trn_bg_ivector,batch_size=batch_size)
dev_bl_embeddings = model.predict(dev_bl_ivector,batch_size=batch_size)
dev_bg_embeddings = model.predict(dev_bg_ivector,batch_size=batch_size)
tst_embeddings = model.predict(tst_ivector,batch_size=batch_size)

# Calculating speaker mean vector
spk_mean, spk_mean_label = make_spkvec(trn_bl_embeddings,trn_bl_id)

# length normalization
trn_bl_embeddings = length_norm(trn_bl_embeddings)
trn_bg_embeddings = length_norm(trn_bg_embeddings)
dev_bl_embeddings = length_norm(dev_bl_embeddings)
dev_bg_embeddings = length_norm(dev_bg_embeddings)
tst_embeddings = length_norm(tst_embeddings)

print('Dev set score using train set :')

# making trials of Dev set
dev_embeddings = np.append(dev_bl_embeddings, dev_bg_embeddings,axis=0)
dev_trials = np.append( np.ones([len(dev_bl_id), 1]), np.zeros([len(dev_bg_id), 1]))

# Cosine distance scoring
scores = spk_mean.dot(dev_embeddings.transpose())
dev_scores = np.max(scores,axis=0)

# Top-S detector EER
dev_EER = calculate_EER(dev_trials, dev_scores)

#divide trial label into target and non-target, plus confusion error(blacklist, fail at blacklist detector)
dev_identified_label = spk_mean_label[np.argmax(scores,axis=0)]
dev_trials_label = np.append( dev_bl_id,dev_bg_id)

# Top-1 detector EER
dev_trials_confusion = get_trials_label_with_confusion(dev_identified_label, dev_trials_label, dev2train, dev_trials )
dev_EER_confusion = calculate_EER_with_confusion(dev_scores,dev_trials_confusion)

print('Test set score using train set:')

#Cosine distance scoring on Test set
scores = spk_mean.dot(tst_embeddings.transpose())
tst_scores = np.max(scores,axis=0)

# top-S detector EER
tst_EER = calculate_EER(tst_trials, tst_scores)

#divide trial label into target and non-target, plus confusion error(blacklist, fail at blacklist detector)
tst_identified_label = spk_mean_label[np.argmax(scores,axis=0)]

# Top-1 detector EER
tst_trials_confusion = get_trials_label_with_confusion(tst_identified_label, tst_trials_label, test2train, tst_trials )
tst_EER_confusion = calculate_EER_with_confusion(tst_scores,tst_trials_confusion)

print('Test set score using train + dev set:')

# get dev set id consistent with Train set
dev_bl_id_along_trnset = []
for iter in range(len(dev_bl_id)):
    dev_bl_id_along_trnset.extend([dev2train[dev_bl_id[iter]]])

# Calculating speaker mean vector
spk_mean, spk_mean_label = make_spkvec(np.append(trn_bl_embeddings,dev_bl_embeddings,0),np.append(trn_bl_id,dev_bl_id_along_trnset))

#Cosine distance scoring on Test set
scores = spk_mean.dot(tst_embeddings.transpose())
tst_scores = np.max(scores,axis=0)

# top-S detector EER
tst_EER = calculate_EER(tst_trials, tst_scores)

#divide trial label into target and non-target, plus confusion error(blacklist, fail at blacklist detector)
tst_identified_label = spk_mean_label[np.argmax(scores,axis=0)]

# Top-1 detector EER
tst_trials_confusion = get_trials_label_with_confusion(tst_identified_label, tst_trials_label, test2train,tst_trials )
tst_EER_confusion = calculate_EER_with_confusion(tst_scores,tst_trials_confusion)


"Scripts for generating, solving and sweeping programs"
from __future__ import unicode_literals, print_function
from time import time
import numpy as np
from ..nomials import parse_subs
from ..solution_array import SolutionArray
from ..keydict import KeyDict
from ..small_scripts import maybe_flatten

try:
    from ad import adnumber
except ImportError:
    adnumber = None
    print("Couldn't import ad; automatic differentiation of linked variables"
          " is disabled.")

def evaluate_linked(constants, linked):
    "Evaluates the values and gradients of linked variables."
    if adnumber:
        kdc = KeyDict({k: adnumber(maybe_flatten(v))
                       for k, v in constants.items()})
        kdc.log_gets = True
    kdc_plain = None
    array_calulated, logged_array_gets = {}, {}
    for v, f in linked.items():
        try:
            assert adnumber  # trigger exit if ad not found
            if v.veckey and v.veckey.original_fn:
                if v.veckey not in array_calulated:
                    ofn = v.veckey.original_fn
                    array_calulated[v.veckey] = np.array(ofn(kdc))
                    logged_array_gets[v.veckey] = kdc.logged_gets
                logged_gets = logged_array_gets[v.veckey]
                out = array_calulated[v.veckey][v.idx]
            else:
                logged_gets = kdc.logged_gets
                out = f(kdc)
            constants[v] = out.x
            v.descr["gradients"] = {}
            for key in logged_gets:
                if key.shape:
                    grad = out.gradient(kdc[key])
                    v.gradients[key] = np.array(grad)
                else:
                    v.gradients[key] = out.d(kdc[key])
        except Exception:  # pylint: disable=broad-except
            from .. import settings
            if settings.get("ad_errors_raise", None):
                raise
            if adnumber:
                print("Couldn't auto-differentiate linked variable %s\n  "
                      "(to raise the error directly for debugging purposes,"
                      " set gpkit.settings[\"ad_errors_raise\"] to True)" % v)
            if kdc_plain is None:
                kdc_plain = KeyDict(constants)
            constants[v] = f(kdc_plain)
            v.descr.pop("gradients", None)
        finally:
            if adnumber:
                kdc.logged_gets = set()

def _progify_fctry(program, return_attr=None):
    """Generates function that returns a program() and optionally an attribute.

    Arguments
    ---------
    program: NomialData
        Class to return, e.g. GeometricProgram or SequentialGeometricProgram
    return_attr: string
        attribute to return in addition to the program
    """
    def programify(self, constants=None, **kwargs):
        "Return program version of self"
        if not constants:
            constants, _, linked = parse_subs(self.varkeys, self.substitutions)
            if linked:
                evaluate_linked(constants, linked)
        prog = program(self.cost, self, constants, **kwargs)
        if return_attr:
            return prog, getattr(prog, return_attr)
        return prog
    return programify

def _solve_fctry(genfunction):
    "Returns function for making/solving/sweeping a program."
    def solvefn(self, solver=None, verbosity=1, skipsweepfailures=False,
                **kwargs):
        """Forms a mathematical program and attempts to solve it.

         Arguments
         ---------
         solver : string or function (default None)
             If None, uses the default solver found in installation.
         verbosity : int (default 1)
             If greater than 0 prints runtime messages.
             Is decremented by one and then passed to programs.
         skipsweepfailures : bool (default False)
             If True, when a solve errors during a sweep, skip it.
         **kwargs : Passed to solver

         Returns
         -------
         sol : SolutionArray
             See the SolutionArray documentation for details.

         Raises
         ------
         ValueError if the program is invalid.
         RuntimeWarning if an error occurs in solving or parsing the solution.
         """
        constants, sweep, linked = parse_subs(self.varkeys, self.substitutions)
        solution = SolutionArray()

        # NOTE SIDE EFFECTS: self.program is set below
        if sweep:
            run_sweep(genfunction, self, solution, skipsweepfailures,
                      constants, sweep, linked,
                      solver, verbosity, **kwargs)
        else:
            self.program, progsolve = genfunction(self)
            result = progsolve(solver, verbosity, **kwargs)
            solution.append(result)
        solution.program = self.program
        solution.model = self
        solution.to_arrays()
        if self.cost.units:
            solution["cost"] = solution["cost"] * self.cost.units
        self.solution = solution  # NOTE: SIDE EFFECTS
        return solution
    return solvefn

# pylint: disable=too-many-locals,too-many-arguments
def run_sweep(genfunction, self, solution, skipsweepfailures,
              constants, sweep, linked,
              solver, verbosity, **kwargs):
    "Runs through a sweep."
    # sort sweeps by the eqstr of their varkey
    sweepvars, sweepvals = zip(*sorted(list(sweep.items()),
                                       key=lambda vkval: vkval[0].eqstr))
    if len(sweep) == 1:
        sweep_grids = np.array(list(sweepvals))
    else:
        sweep_grids = np.meshgrid(*list(sweepvals))

    N_passes = sweep_grids[0].size
    sweep_vects = {var: grid.reshape(N_passes)
                   for (var, grid) in zip(sweepvars, sweep_grids)}

    if verbosity > 0:
        print("Solving over %i passes." % N_passes)
        tic = time()

    self.program = []
    for i in range(N_passes):
        constants.update({var: sweep_vect[i]
                          for (var, sweep_vect) in sweep_vects.items()})
        if linked:
            kdc = KeyDict(constants)
            constants.update({v: f(kdc) for v, f in linked.items()})
        program, solvefn = genfunction(self, constants)
        self.program.append(program)  # NOTE: SIDE EFFECTS
        try:
            solution.append(solvefn(solver, verbosity-1, **kwargs))
        except (RuntimeWarning, ValueError):
            if not skipsweepfailures:
                raise RuntimeWarning("solve failed during sweep; program"
                                     " has been saved to m.program[-1]."
                                     " To ignore such failures, solve with"
                                     " skipsweepfailures=True.")
    if not solution:
        raise RuntimeWarning("no sweeps solved successfully.")

    solution["sweepvariables"] = KeyDict()
    ksweep = KeyDict(sweep)
    for var, val in list(solution["constants"].items()):
        if var in ksweep:
            solution["sweepvariables"][var] = val
            del solution["constants"][var]
        elif var not in linked:
            solution["constants"][var] = [val[0]]

    if verbosity > 0:
        soltime = time() - tic
        print("Sweeping took %.3g seconds." % (soltime,))


"""
Modified from PointConv: https://github.com/DylanWusee/pointconv
Author: Jiachen Xu and Jingyu Gong
Date: June 2020
"""

import os
import sys
BASE_DIR = os.path.dirname(__file__)
sys.path.append(BASE_DIR)
sys.path.append(os.path.join(BASE_DIR, '../'))
sys.path.append(os.path.join(BASE_DIR, '../utils'))
import tensorflow as tf
import numpy as np
import tf_util
from PointConv import feature_encoding_layer, feature_decoding_layer, feature_encoding_layer_extra

def placeholder_scene_inputs(batch_size, num_point, num_classes):
    pointclouds_pl = tf.placeholder(tf.float32, shape=(batch_size, num_point, 6))
    labels_pl = tf.placeholder(tf.int32, shape=(batch_size, num_point))
    labels_onehot_pl = tf.placeholder(tf.float32, shape=(batch_size, num_point,num_classes))
    smpws_pl = tf.placeholder(tf.float32, shape=(batch_size, num_point))
    external_scene_encode_pl = tf.placeholder(tf.int32,shape=(batch_size,num_classes))

    cos_loss_weight = tf.placeholder(tf.float32, shape=None)
    return pointclouds_pl, labels_pl, labels_onehot_pl, smpws_pl, external_scene_encode_pl, cos_loss_weight

def get_scene_model(point_cloud, is_training, num_class, sigma, bn_decay=None, weight_decay = None):
    """ Semantic segmentation PointNet, input is BxNx3, output Bxnum_class """

    batch_size = point_cloud.get_shape()[0].value
    num_point = point_cloud.get_shape()[1].value
    end_points = {}
    l0_xyz = point_cloud[:,:,:3]
    l0_points = point_cloud
    l0_points_xyz = point_cloud[:,:,:3]
    l0_points_rgb = point_cloud[:,:,3:]

    # Feature encoding layers
    l1_xyz, l1_points_xyz = feature_encoding_layer(l0_xyz, l0_points_xyz, npoint=2048, radius = 0.1, sigma = sigma, K=8, mlp=[32,32,32], akc_channel = 3, is_training=is_training, bn_decay=bn_decay, weight_decay = weight_decay, scope='layer1_xyz')
    _, l1_points_rgb = feature_encoding_layer_extra(l0_xyz, l1_xyz, l0_points_rgb, npoint=2048, radius = 0.1, sigma = sigma, K=8, mlp=[32,32,32], akc_channel = 3, is_training=is_training, bn_decay=bn_decay, weight_decay = weight_decay, scope='layer1_rgb')
    l1_points = tf.concat([l1_points_xyz, l1_points_rgb], -1)
    l1_xyz_1024, l1_points_1024 = feature_encoding_layer(l1_xyz, l1_points, npoint=1024, radius = 0.1, sigma = sigma, K=8, mlp=[32,32,64], akc_channel=16, is_training=is_training, bn_decay=bn_decay, weight_decay = weight_decay, scope='layer1_1024')
    l1_xyz_512, l1_points_512 = feature_encoding_layer(l1_xyz_1024, l1_points_1024, npoint=512, radius = 0.1, sigma = sigma, K=8, mlp=[64,64,64], akc_channel=32, is_training=is_training, bn_decay=bn_decay, weight_decay = weight_decay, scope='layer1_512')
    l2_xyz, l2_points = feature_encoding_layer(l1_xyz_512, l1_points_512, npoint=256, radius = 0.2, sigma = 2 * sigma, K=8, mlp=[64,64,128], akc_channel=32, is_training=is_training, bn_decay=bn_decay, weight_decay = weight_decay, scope='layer2')
    l2_xyz_128, l2_points_128 = feature_encoding_layer(l2_xyz, l2_points, npoint=128, radius = 0.2, sigma = 2 * sigma, K=8, mlp=[128,128,128], akc_channel=None, is_training=is_training, bn_decay=bn_decay, weight_decay = weight_decay, scope='layer2_128')
    l3_xyz, l3_points = feature_encoding_layer(l2_xyz_128, l2_points_128, npoint=64, radius = 0.4, sigma = 4 * sigma, K=8, mlp=[128,128,256], akc_channel=None, is_training=is_training, bn_decay=bn_decay, weight_decay = weight_decay, scope='layer3')
    l4_xyz, l4_points = feature_encoding_layer(l3_xyz, l3_points, npoint=36, radius = 0.8, sigma = 8 * sigma, K=8, mlp=[256,256,512], akc_channel=None, is_training=is_training, bn_decay=bn_decay, weight_decay = weight_decay, scope='layer4')

    external_l5_xyz,external_l5_points = feature_encoding_layer(l4_xyz, l4_points, npoint=8, radius = 1.6, sigma = 8 * sigma, K=8, mlp=[512,512,512], akc_channel=None, is_training=is_training, bn_decay=bn_decay, weight_decay = weight_decay, scope='external_layer5')
    external_l6_scene_feature = tf.reduce_mean(external_l5_points,axis=1,keepdims=True)
    external_scene_feature = tf_util.dropout(external_l6_scene_feature, keep_prob=0.5, is_training=is_training, scope='external_dp')
    external_scene_feature = tf_util.conv1d(external_scene_feature, num_class, 1, padding='VALID', activation_fn=None, weight_decay=weight_decay, scope='external_fc')

    # Feature decoding layers
    l3_points = feature_decoding_layer(l3_xyz, l4_xyz, l3_points, l4_points, 0.8, 8 * sigma, 8, [512,512], is_training, bn_decay, weight_decay, scope='fa_layer1')
    l2_points_128 = feature_decoding_layer(l2_xyz_128, l3_xyz, l2_points_128, l3_points, 0.4, 4 * sigma, 8, [256,256], is_training, bn_decay, weight_decay, scope='fa_layer2_128')
    l2_points = feature_decoding_layer(l2_xyz, l2_xyz_128, l2_points, l2_points_128, 0.4, 4 * sigma, 8, [256,256], is_training, bn_decay, weight_decay, scope='fa_layer2')
    l1_points_512 = feature_decoding_layer(l1_xyz_512, l2_xyz, l1_points_512, l2_points, 0.2, 2 * sigma, 8, [256,256], is_training, bn_decay, weight_decay, scope='fa_layer3_512')
    l1_points_1024 = feature_decoding_layer(l1_xyz_1024, l1_xyz_512, l1_points_1024, l1_points_512, 0.2, 2 * sigma, 8, [256,256], is_training, bn_decay, weight_decay, scope='fa_layer3_1024')
    l1_points = feature_decoding_layer(l1_xyz, l1_xyz_1024, l1_points, l1_points_1024, 0.2, 2 * sigma, 8, [256,128], is_training, bn_decay, weight_decay, scope='fa_layer3')
    l0_points = feature_decoding_layer(l0_xyz, l1_xyz, l0_points, l1_points, 0.1, sigma, 8, [128,128,128], is_training, bn_decay, weight_decay, scope='fa_layer4')

    # FC layers
    net = tf_util.conv1d(l0_points, 128, 1, padding='VALID', bn=True, is_training=is_training, scope='fc1', bn_decay=bn_decay, weight_decay=weight_decay)
    end_points['feats'] = net
    net = tf_util.dropout(net, keep_prob=0.5, is_training=is_training, scope='dp1')
    net = tf_util.conv1d(net, num_class, 1, padding='VALID', activation_fn=None, weight_decay=weight_decay, scope='fc2')

    return net, end_points, external_scene_feature

def get_mask(label1, label2):
    num_points1 = label1.get_shape().as_list()[1]
    num_points2 = label2.get_shape().as_list()[1]

    label1 = tf.expand_dims(label1, axis=-1)
    label2 = tf.expand_dims(label2, axis=-1)
    label2_transpose = tf.transpose(label2, perm=[0, 2, 1])

    label1 = tf.tile(label1, [1, 1, num_points2])
    label2_transpose = tf.tile(label2_transpose, [1, num_points1, 1])

    # if they have same label, the make value is -1
    mask = (tf.cast(tf.equal(label1, label2_transpose), tf.float32) - 1) * 255

    return mask

def knn1(xyz1, xyz2, k=4, label1=None, label2=None):
    xyz1 = tf.squeeze(xyz1)
    xyz2 = tf.squeeze(xyz2)

    xyz2_transpose = tf.transpose(xyz2, perm=[0, 2, 1])
    point_cloud_inner = tf.matmul(xyz1, xyz2_transpose)

    point_cloud_inner = -2 * point_cloud_inner
    xyz1_square = tf.reduce_sum(tf.square(xyz1), axis=-1, keep_dims=True)

    xyz2_square = tf.reduce_sum(tf.square(xyz2), axis=-1, keep_dims=True)

    xyz2_square_trans = tf.transpose(xyz2_square, perm=[0, 2, 1])

    adj_matrix = xyz1_square + point_cloud_inner + xyz2_square_trans

    neg_adj = -adj_matrix

    if label1 is not None:
        mask = get_mask(label1, label2)
        neg_adj += mask

    _, nn_idx = tf.nn.top_k(neg_adj, k=k)
    return nn_idx

def get_scene_loss(cos_loss_weight, pred, label, label_onehot, smpw, external_scene_feature,external_scene_encode, point_features=None, xyz=None):
    """ pred: BxNxC,
        label: BxN,
	smpw: BxN """
    if point_features is not None:
        pred_possibility = tf.reduce_max(tf.nn.softmax(pred, axis=-1), axis=-1)
        pred_class = tf.cast(tf.argmax(pred, axis=-1), tf.int32)

        batch_size, num_point, num_dims = point_features.get_shape()[:3]
        xyz_lst, feature_lst, label_lst = [], [], []

        def f1(correct_index1, i, possibility, correct_num):
            choice = tf.cond(tf.less(correct_num, 2048), lambda: f4(possibility, correct_num), lambda: f3(possibility))

            choice.set_shape([2048])

            correct_index = tf.squeeze(tf.gather(correct_index1, choice))

            xyz_lst.append(tf.expand_dims(tf.gather(xyz[i], correct_index), axis=0))
            feature_lst.append(tf.expand_dims(tf.gather(point_features[i], correct_index), axis=0))
            label_lst.append(tf.expand_dims(tf.gather(label[i], correct_index), axis=0))

            return tf.constant(False)

        def f2():
            loss_xyz = tf.concat(xyz_lst, axis=0)
            loss_feature = tf.concat(feature_lst, axis=0)
            loss_label = tf.concat(label_lst, axis=0)

            idx = knn1(loss_xyz, xyz, 8, loss_label, label)
            idx_ = tf.range(batch_size) * num_point
            idx_ = tf.reshape(idx_, [batch_size, 1, 1])

            point_features_flat = tf.reshape(point_features, [-1, num_dims])
            point_features_neighbors = tf.gather(point_features_flat, idx+idx_)
            point_features_central = tf.tile(tf.expand_dims(loss_feature, 2), [1, 1, 8, 1])

            point_features_central = tf.stop_gradient(point_features_central)

            norm1 = tf.sqrt(tf.reduce_sum(tf.square(point_features_neighbors), axis=-1))
            norm2 = tf.sqrt(tf.reduce_sum(tf.square(point_features_central), axis=-1))
            product = tf.reduce_sum(point_features_neighbors * point_features_central, axis=-1)
            cos_loss = tf.reduce_mean(product / (norm1 * norm2 + 1e-5))

            return 1 - cos_loss

        def f3(possibility):
            _, choice = tf.nn.top_k(possibility, k=2048)

            return choice

        def f4(possibility, correct_num):
            possibility_max = tf.argmax(possibility, output_type=tf.int32)

            _, choice = tf.nn.top_k(possibility, k=correct_num * 2 // 3)

            choice = tf.pad(choice, [[0, 2048 - correct_num * 2 // 3]], mode="CONSTANT", constant_values=possibility_max)

            return choice

        judge_lst = []
        for i in range(batch_size):
            index = tf.equal(pred_class[i], label[i])
            correct_num = tf.reduce_sum(tf.cast(index, tf.int32))
            correct_index1 = tf.reshape(tf.where(index), [-1])
            possibility = tf.gather(pred_possibility[i], correct_index1)

            judge_lst.append(tf.cond(tf.equal(correct_num, tf.constant(0)),
                                    lambda: tf.constant(True),
                                    lambda: f1(correct_index1, i, possibility, correct_num)))

        cos_loss = tf.cond(tf.equal(tf.reduce_sum(tf.cast(judge_lst, tf.int32)), tf.constant(0)), lambda: f2(), lambda: tf.constant(1.0))

    else:
        cos_loss = tf.constant(0.0)

    external_scene_feature_2d = tf.squeeze(external_scene_feature,[1])
    external_scene_loss = tf.losses.sigmoid_cross_entropy(external_scene_encode,external_scene_feature_2d)
    external_scene_probability = tf.sigmoid(external_scene_feature)
    external_scene_probability = tf.tile(external_scene_probability,[1,pred.get_shape()[1],1])
    external_scene_probability = tf.stop_gradient(external_scene_probability)
    pred = tf.nn.softmax(pred)
    pred = tf.multiply(pred,external_scene_probability)
    softmax_probability_sum = tf.reduce_sum(pred,axis=2,keepdims=True)
    softmax_probability_sum = tf.tile(softmax_probability_sum,[1,1,pred.get_shape()[2]])
    pred = tf.div(pred,softmax_probability_sum+1e-5)
    classify_loss = tf.keras.backend.categorical_crossentropy(label_onehot,pred)
    classify_loss = classify_loss * smpw
    print('add smpw')

    weight_reg = tf.add_n(tf.get_collection('losses'))
    classify_loss_mean = tf.reduce_mean(classify_loss, name='classify_loss_mean')
    total_loss = classify_loss_mean + weight_reg + cos_loss * cos_loss_weight + external_scene_loss
    tf.summary.scalar('classify loss', classify_loss_mean)
    tf.summary.scalar('total loss', total_loss)

    return total_loss, pred

if __name__=='__main__':
    import pdb
    pdb.set_trace()

    with tf.Graph().as_default():
        inputs = tf.zeros((32,2048,3))
        net, _ = get_model(inputs, tf.constant(True), 10, 1.0)
        print(net)