# ===========================================================================
# This module is adpated from: https://github.com/fchollet/keras
# Revision: @bec2701
# Original work Copyright (c) 2014-2015 keras contributors
# Some idea are also borrowed from Lasagne library
# Original work Copyright (c) 2014-2015 Lasagne contributors
# Some work are adapted from tensorfuse library
# Original work Copyright (c) [dementrock](https://github.com/dementrock)
# Modified work Copyright 2016-2017 TrungNT
# ===========================================================================

import tensorflow as tf
import numpy as np
import os
from collections import OrderedDict

from .. import config
from .numpy_backend import get_random_magic_seed

from six.moves import range, zip

_FLOATX = config.floatX()
_EPSILON = config.epsilon()

# ===========================================================================
# INTERNAL UTILS
# ===========================================================================
_SESSION = None

def get_session():
    global _SESSION
    if _SESSION is None:
        if not os.environ.get('OMP_NUM_THREADS'):
            _SESSION = tf.Session('')
        else:
            nb_thread = int(os.environ.get('OMP_NUM_THREADS'))
            _SESSION = tf.Session(config=tf.ConfigProto(intra_op_parallelism_threads=nb_thread))
    return _SESSION

def _set_session(session):
    global _SESSION
    _SESSION = session

# From Theano

def _format_as(use_list, use_tuple, outputs):
    """
    Formats the outputs according to the flags `use_list` and `use_tuple`.
    If `use_list` is True, `outputs` is returned as a list (if `outputs`
    is not a list or a tuple then it is converted in a one element list).
    If `use_tuple` is True, `outputs` is returned as a tuple (if `outputs`
    is not a list or a tuple then it is converted into a one element tuple).
    Otherwise (if both flags are false), `outputs` is returned.
    """
    assert not (use_list and use_tuple), \
        "Both flags cannot be simultaneously True"
    if (use_list or use_tuple) and not isinstance(outputs, (list, tuple)):
        if use_list:
            return [outputs]
        else:
            return (outputs,)
    elif not (use_list or use_tuple) and isinstance(outputs, (list, tuple)):
        assert len(outputs) == 1, \
            "Wrong arguments. Expected a one element list"
        return outputs[0]
    elif use_list or use_tuple:
        if use_list:
            return list(outputs)
        else:
            return tuple(outputs)
    else:
        return outputs

def _wrap_into_list(x):
    """
    Wrap the input into a list if it is not already a list.
    """
    if x is None:
        return []
    elif not isinstance(x, (list, tuple)):
        return [x]
    else:
        return list(x)
# ===========================================================================
# VARIABLE MANIPULATION
# ===========================================================================

def variable(value, dtype=_FLOATX, name=None, broadcastable=None):
    v = tf.Variable(np.asarray(value, dtype=dtype), name=name)
    get_session().run(v.initializer)
    return v

def zeros_var(shape, dtype=_FLOATX, name=None):
    return variable(np.zeros(shape), dtype, name)

def ones_var(shape, dtype=_FLOATX, name=None):
    return variable(np.ones(shape), dtype, name)

def is_variable(v):
    return isinstance(v, tf.python.Variable)

_PLACEHOLDER_ID = 0
_PLACEHOLDER_SHAPE = {}

def placeholder(shape=None, ndim=None, dtype=_FLOATX, name=None):
    # name must match: [A-Za-z0-9.][A-Za-z0-9_.\-/]*
    if not shape:
        if ndim:
            shape = [None for _ in range(ndim)]

    # ====== Modify add name prefix ====== #
    global _PLACEHOLDER_ID
    name_prefix = 'ID.%02d.' % _PLACEHOLDER_ID
    _PLACEHOLDER_ID += 1
    if name is None:
        name = ''
    name = name_prefix + name
    placeholder = tf.placeholder(dtype, shape=shape, name=name)
    _PLACEHOLDER_SHAPE[placeholder.name] = shape
    return placeholder

def is_expression(v):
    return isinstance(v, tf.python.Tensor)

def is_placeholder(v):
    if is_expression(v) and v.name in _PLACEHOLDER_SHAPE:
        return True
    return False

def eval(x):
    '''Run a graph.
    '''
    if isinstance(x, tf.TensorShape):
        return x.as_list()
    return x.eval(session=get_session())

# ===========================================================================
# Shape operators
# ===========================================================================

def shape(x):
    return x.get_shape()

def int_shape(x):
    shape = x.get_shape()
    return tuple([i.__int__() for i in shape])

def ndim(x):
    return len(x.get_shape())

def broadcastable(x):
    return None

def addbroadcast(x, *axes):
    return x

# ===========================================================================
# Predefined data
# ===========================================================================

def zeros(shape, dtype=_FLOATX, name=None):
    return tf.zeros(shape, dtype=dtype, name=name)

def ones(shape, dtype=_FLOATX, name=None):
    return tf.ones(shape, dtype=dtype, name=name)

def ones_like(x, name=None):
    return tf.ones_like(x, name=name)

def zeros_like(x, name=None):
    return tf.zeros_like(x, name=name)

def count_params(x):
    '''Return number of scalars in a tensor.
    '''
    shape = x.get_shape()
    return np.prod([shape[i]._value for i in range(len(shape))])

def cast(x, dtype):
    if 'tensorflow.' in str(x.__class__):
        return tf.cast(x, dtype)
    return np.cast[dtype](x)

def castX(x):
    return cast(x, _FLOATX)

# ===========================================================================
# LINEAR ALGEBRA
# ===========================================================================

def dot(x, y):
    return tf.matmul(x, y)

def transpose(x):
    return tf.transpose(x)

def gather(reference, indices):
    '''
    # Arguments
        reference: a tensor.
        indices: an int tensor of indices.

    # Returns
        a tensor of same type as `reference`.
    '''
    return tf.gather(reference, indices)

# ===========================================================================
# ELEMENT-WISE OPERATIONS
# ===========================================================================
def normalize_axis(axis, ndim):
    if type(axis) is tuple:
        axis = list(axis)
    if type(axis) is list:
        for i, a in enumerate(axis):
            if a is not None and a < 0:
                axis[i] = a % ndim
    else:
        if axis is not None and axis < 0:
            axis = axis % ndim
    return axis

def max(x, axis=None, keepdims=False):
    axis = normalize_axis(axis, ndim(x))
    return tf.reduce_max(x, reduction_indices=axis, keep_dims=keepdims)

def min(x, axis=None, keepdims=False):
    axis = normalize_axis(axis, ndim(x))
    return tf.reduce_min(x, reduction_indices=axis, keep_dims=keepdims)

def sum(x, axis=None, keepdims=False):
    '''Sum of the values in a tensor, alongside the specified axis.
    '''
    axis = normalize_axis(axis, ndim(x))
    return tf.reduce_sum(x, reduction_indices=axis, keep_dims=keepdims)

def prod(x, axis=None, keepdims=False):
    '''Multiply the values in a tensor, alongside the specified axis.
    '''
    axis = normalize_axis(axis, ndim(x))
    return tf.reduce_prod(x, reduction_indices=axis, keep_dims=keepdims)

def var(x, axis=None, keepdims=False):
    axis = normalize_axis(axis, ndim(x))
    if x.dtype.base_dtype == tf.bool:
        x = tf.cast(x, _FLOATX)
    m = tf.reduce_mean(x, reduction_indices=axis, keep_dims=True)
    devs_squared = tf.square(x - m)
    return tf.reduce_mean(devs_squared,
                          reduction_indices=axis,
                          keep_dims=keepdims)

def std(x, axis=None, keepdims=False):
    axis = normalize_axis(axis, ndim(x))
    if x.dtype.base_dtype == tf.bool:
        x = tf.cast(x, _FLOATX)
    m = tf.reduce_mean(x, reduction_indices=axis, keep_dims=True)
    devs_squared = tf.square(x - m)
    return tf.sqrt(tf.reduce_mean(devs_squared,
                                  reduction_indices=axis,
                                  keep_dims=keepdims))

def mean(x, axis=None, keepdims=False):
    axis = normalize_axis(axis, ndim(x))
    if x.dtype.base_dtype == tf.bool:
        x = tf.cast(x, _FLOATX)
    return tf.reduce_mean(x, reduction_indices=axis, keep_dims=keepdims)

def any(x, axis=None, keepdims=False):
    '''Bitwise reduction (logical OR).

    Return array of uint8 (0s and 1s).
    '''
    axis = normalize_axis(axis, ndim(x))
    x = tf.cast(x, tf.bool)
    x = tf.reduce_any(x, reduction_indices=axis, keep_dims=keepdims)
    return tf.cast(x, tf.uint8)

def argmax(x, axis=-1):
    if axis < 0:
        axis = axis % len(x.get_shape())
    return tf.argmax(x, axis)

def argsort(x, axis=-1):
    raise NotImplementedError

def argtop_k(x, k=1):
    ''' See also: tf.nn.in_top_k '''
    return tf.nn.top_k(x, k)[1]

def argmin(x, axis=-1):
    if axis < 0:
        axis = axis % len(x.get_shape())
    return tf.argmin(x, axis)

def square(x):
    return tf.square(x)

def abs(x):
    return tf.abs(x)

def sqrt(x):
    x = tf.clip_by_value(x, tf.cast(0., dtype=_FLOATX),
                         tf.cast(np.inf, dtype=_FLOATX))
    return tf.sqrt(x)

def exp(x):
    return tf.exp(x)

def log(x):
    return tf.log(x)

def round(x):
    return tf.round(x)

def pow(x, a):
    return tf.pow(x, a)

def clip(x, min_value, max_value):
    if max_value < min_value:
        max_value = min_value
    return tf.clip_by_value(x, tf.cast(min_value, dtype=_FLOATX),
                            tf.cast(max_value, dtype=_FLOATX))

def maximum(x, y):
    return tf.maximum(x, y)

def minimum(x, y):
    return tf.minimum(x, y)

# ===========================================================================
# SHAPE OPERATIONS
# ===========================================================================
def reverse(x, axis=-1):
    '''Apply [::-1] to appropriate axis'''
    ndim = len(x.get_shape())
    dims = [False] * ndim
    if axis < 0:
        axis = axis % ndim
    dims[axis] = True
    return tf.reverse(x, dims)

def concatenate(tensors, axis=-1):
    if axis < 0:
        axis = axis % len(tensors[0].get_shape())
    return tf.concat(axis, tensors)

def reshape(x, shape):
    return tf.reshape(x, shape)

def dimshuffle(x, pattern):
    '''
    # Arguments
        pattern: should be a tuple or list of
            dimension indices, e.g. [0, 2, 1].
    '''
    if 'x' in pattern:
        x = tf.transpose(x, perm=[i for i in pattern if i != 'x'])
    for i, p in enumerate(pattern):
        if p == 'x':
            x = tf.expand_dims(x, i)
    return x

def resize_images(X, height_factor, width_factor, dim_ordering):
    '''Resize the images contained in a 4D tensor of shape
    - [batch, channels, height, width] (for 'th' dim_ordering)
    - [batch, height, width, channels] (for 'tf' dim_ordering)
    by a factor of (height_factor, width_factor). Both factors should be
    positive integers.
    '''
    if dim_ordering == 'th':
        new_height = shape(X)[2].value * height_factor
        new_width = shape(X)[3].value * width_factor
        X = dimshuffle(X, [0, 2, 3, 1])
        X = tf.image.resize_nearest_neighbor(X, (new_height, new_width))
        return dimshuffle(X, [0, 3, 1, 2])
    elif dim_ordering == 'tf':
        new_height = shape(X)[1].value * height_factor
        new_width = shape(X)[2].value * width_factor
        return tf.image.resize_nearest_neighbor(X, (new_height, new_width))
    else:
        raise Exception('Invalid dim_ordering: ' + dim_ordering)

def repeat_elements(x, rep, axis):
    '''Repeats the elements of a tensor along an axis, like np.repeat

    If x has shape (s1, s2, s3) and axis=1, the output
    will have shape (s1, s2 * rep, s3)
    '''
    x_shape = x.get_shape().as_list()
    # slices along the repeat axis
    splits = tf.split(axis, x_shape[axis], x)
    # repeat each slice the given number of reps
    x_rep = [s for s in splits for i in range(rep)]
    return tf.concat(axis, x_rep)

def repeat(x, n):
    '''Repeat a 2D tensor:

    if x has shape (samples, dim) and n=2,
    the output will have shape (samples, 2, dim)
    '''
    assert ndim(x) == 2
    tensors = [x] * n
    stacked = tf.pack(tensors)
    return tf.transpose(stacked, (1, 0, 2))

def tile(x, n):
    return tf.tile(x, n)

def flatten(x, outdim=2):
    '''Turn a n-D tensor into a m-D tensor (m < n) where
    the first dimension is conserved.
    '''
    if outdim == 1:
        pattern = [-1]
    else:
        pattern = [-1, np.prod(x.get_shape()[(outdim - 1):].as_list())]
    return tf.reshape(x, pattern)

def expand_dims(x, dim=-1):
    '''Add a 1-sized dimension at index "dim".
    '''
    return tf.expand_dims(x, dim)

def squeeze(x, axis):
    '''Remove a 1-dimension from the tensor at index "axis".
    '''
    return tf.squeeze(x, [axis])

def temporal_padding(x, padding=1):
    '''Pad the middle dimension of a 3D tensor
    with "padding" zeros left and right.
    '''
    pattern = [[0, 0], [padding, padding], [0, 0]]
    return tf.pad(x, pattern)

def spatial_2d_padding(x, padding=(1, 1), dim_ordering='th'):
    '''Pad the 2nd and 3rd dimensions of a 4D tensor
    with "padding[0]" and "padding[1]" (resp.) zeros left and right.
    '''
    if dim_ordering == 'th':
        pattern = [[0, 0], [0, 0],
                   [padding[0], padding[0]], [padding[1], padding[1]]]
    else:
        pattern = [[0, 0],
                   [padding[0], padding[0]], [padding[1], padding[1]],
                   [0, 0]]
    return tf.pad(x, pattern)

def stack(*x):
    return tf.pack(x)

# ===========================================================================
# VALUE MANIPULATION
# ===========================================================================

def get_value(x, borrow=False):
    '''Technically the same as eval() for TF.
    '''
    return x.eval(session=get_session())

def set_value(x, value):
    tf.assign(x, np.asarray(value)).op.run(session=get_session())

def set_subtensor(x, y):
    raise NotImplementedError

# ===========================================================================
# GRAPH MANIPULATION
# ===========================================================================
_GLOBALS_UPDATES = OrderedDict()

def add_global_updates(variable, value):
    '''trick to update tensorflow variables anywhere
    This dictionary will be reseted after each time you create a function
    '''
    _GLOBALS_UPDATES[variable] = value

def reset_global_updates():
    global _GLOBALS_UPDATES
    _GLOBALS_UPDATES = OrderedDict()

class Function(object):

    def __init__(self, inputs, outputs, updates=[]):
        assert type(inputs) in {list, tuple}
        if type(outputs) not in {list, tuple}:
            outputs = [outputs]
            self._return_list = False
        else:
            self._return_list = True
        if isinstance(updates, OrderedDict):
            updates = updates.items()
        assert type(updates) in {list, tuple}
        self.inputs = list(inputs)
        self.outputs = list(outputs)
        with tf.control_dependencies(self.outputs):
            self.updates = [tf.assign(p, new_p) for (p, new_p) in updates]
        # ====== add global_update ====== #
        self.global_update = [tf.assign(p, new_p) for (p, new_p) in _GLOBALS_UPDATES.items()]
        reset_global_updates()

    def __call__(self, *inputs):
        assert type(inputs) in {list, tuple}
        names = [v.name for v in self.inputs]
        feed_dict = dict(zip(names, inputs))
        session = get_session()
        # ====== add global updates ====== #
        updated = session.run(self.outputs + self.updates + self.global_update,
            feed_dict=feed_dict)
        if self._return_list:
            return updated[:len(self.outputs)]
        return updated[0]

def function(inputs, outputs, updates=[]):
    return Function(inputs, outputs, updates=updates)

def gradients(loss, variables, consider_constant=None, known_grads=None):
    """
    Return symbolic gradients for one or more variables with respect to some
    cost.

    For more information about how automatic differentiation works in Theano,
    see :mod:`gradient`. For information on how to implement the gradient of
    a certain Op, see :func:`grad`.

    Parameters
    ----------
    cost : scalar (0-dimensional) tensor variable or None
        Value with respect to which we are differentiating.  May be
        `None` if known_grads is provided.
    wrt : variable or list of variables
        term[s] for which we want gradients
    consider_constant : list of expressions(variables)
        expressions not to backpropagate through
    known_grads : dict, optional
        A dictionary mapping variables to their gradients. This is
        useful in the case where you know the gradient on some
        variables but do not know the original cost.
    Returns
    -------
    variable or list/tuple of variables (matches `wrt`)
        symbolic expression of gradient of `cost` with respect to each
        of the `wrt` terms.  If an element of `wrt` is not
        differentiable with respect to the output, then a zero
        variable is returned.

    """
    if consider_constant is not None:
        for i in consider_constant:
            tf.stop_gradient(i)
        raise NotImplementedError
    grad = tf.gradients(loss, variables)
    if known_grads is not None:
        grad = [known_grads[i] if i in known_grads else j
                for i, j in zip(variables, grad)]
    return grad

def grad_clip(x, clip):
    '''
    This clip the gradient of expression, used on forward pass but clip the
    gradient on backward pass

    This is an elemwise operation.

    Parameters
    ----------
    x: expression
        the variable we want its gradient inputs clipped
    lower_bound: float
        The lower bound of the gradient value
    upper_bound: float
        The upper bound of the gradient value.

    Example
    -------
    >>> x = theano.tensor.scalar()
    >>>
    >>> z = theano.tensor.grad(grad_clip(x, -1, 1)**2, x)
    >>> z2 = theano.tensor.grad(x**2, x)
    >>>
    >>> f = theano.function([x], outputs = [z, z2])
    >>>
    >>> print(f(2.0))  # output (1.0, 4.0)

    Note
    ----
    We register an opt in tensor/opt.py that remove the GradClip.
    So it have 0 cost in the forward and only do work in the grad.

    '''
    # TODO: no implementation for grad_clipping on tensorflow on forward pass
    return x

def jacobian(expression, wrt):
    # copying theano's implementation, which is based on scan
    #from theano.tensor import arange
    # Check inputs have the right format
    assert is_variable(expression), \
        "tensor.jacobian expects a Variable as `expression`"
    assert expression.ndim < 2, \
        ("tensor.jacobian expects a 1 dimensional variable as "
         "`expression`. If not use flatten to make it a vector")
    assert not is_variable(expression.shape[0]), \
        "shape of the expression must be known"

    using_list = isinstance(wrt, list)
    using_tuple = isinstance(wrt, tuple)

    if isinstance(wrt, (list, tuple)):
        wrt = list(wrt)
    else:
        wrt = [wrt]

    if expression.ndim == 0:
        # expression is just a scalar, use grad
        return _format_as(using_list, using_tuple, gradients(expression, wrt))

    def inner_function(*args):
        idx = args[0]
        expr = args[1]
        rvals = []
        for inp in args[2:]:
            try:
                rval = gradients(expr[idx], inp)
            except Exception as e:
                import ipdb; ipdb.set_trace()
            if rval is None:
                import ipdb; ipdb.set_trace()
            rvals.append(rval)
        return rvals
    # Computing the gradients does not affect the random seeds on any random
    # generator used n expression (because during computing gradients we are
    # just backtracking over old values. (rp Jan 2012 - if anyone has a
    # counter example please show me)
    jacobs, updates = scan(inner_function,
                           sequences=[range(expression.shape[0])],
                           non_sequences=[expression] + wrt,
                           n_steps=expression.shape[0])
    assert not updates
    return _format_as(using_list, using_tuple, jacobs)

def hessian(expression, wrt):
    raise NotImplementedError

# ===========================================================================
# CONTROL FLOW
# ===========================================================================

def scan(step_fn, sequences=None, outputs_info=None, non_sequences=None,
    n_steps=None, truncate_gradient=-1, go_backwards=False):
    from operator import itemgetter
    # n_steps must be provided under cgt or tensorflow
    if n_steps is None:
        raise ValueError(
            'n_steps must be provided for scan to work under TensorFlow')
    sequences = _wrap_into_list(sequences)
    non_sequences = _wrap_into_list(non_sequences)
    if outputs_info is not None:
        outputs_info = _wrap_into_list(outputs_info)
    if go_backwards and n_steps < 0:
        go_backwards = False
        n_steps = -n_steps
    if go_backwards or n_steps < 0:
        go_backwards = True
        n_steps = abs(n_steps)
    step_outputs = []
    cur_output = outputs_info
    loop_range = range(n_steps - 1, -1, -1) if go_backwards else range(n_steps)
    for i in loop_range:
        # Only pass output if needed
        if outputs_info is not None:
            cur_output = step_fn(*(map(itemgetter(i), sequences) + cur_output + non_sequences))
        else:
            cur_output = step_fn(*(map(itemgetter(i), sequences) + non_sequences))
        step_outputs.append(cur_output)
    outputs = []
    try:
        if len(step_outputs) > 0:
            if outputs_info is None:
                for i in range(len(step_outputs[0])):
                    outputs.append(tf.pack(map(itemgetter(i), step_outputs)))
                #outputs = step_outputs
            else:
                for i in range(len(outputs_info)):
                    outputs.append(tf.pack(map(itemgetter(i), step_outputs)))
        else:
            import ipdb; ipdb.set_trace()
    except Exception as e:
        raise e
    # This is quite ugly, but unfortunately it's what theano does
    if len(outputs) > 1:
        # update is not supported yet
        return outputs, None
    elif len(outputs) == 1:
        return outputs[0], None
    else:
        return None, None

def loop(step_fn, n_steps, sequences=None, outputs_info=None, non_sequences=None,
         go_backwards=False):
    """
    Helper function to unroll for loops. Can be used to unroll theano.scan.
    The parameter names are identical to theano.scan, please refer to here
    for more information.

    Note that this function does not support the truncate_gradient
    setting from theano.scan.

    Parameters
    ----------
    step_fn : function
        Function that defines calculations at each step.

    sequences : TensorVariable or list of TensorVariables
        List of TensorVariable with sequence data. The function iterates
        over the first dimension of each TensorVariable.

    outputs_info : list of TensorVariables
        List of tensors specifying the initial values for each recurrent
        value. Specify output_info to None for non-arguments to
        the step_function

    non_sequences: list of TensorVariables
        List of theano.shared variables that are used in the step function.

    n_steps: int
        Number of steps to unroll.

    go_backwards: bool
        If true the recursion starts at sequences[-1] and iterates
        backwards.

    Returns
    -------
    List of TensorVariables. Each element in the list gives the recurrent
    values at each time step.

    """
    if not isinstance(sequences, (list, tuple)):
        sequences = [] if sequences is None else [sequences]

    # When backwards reverse the recursion direction
    counter = range(n_steps)
    if go_backwards:
        counter = counter[::-1]

    # ====== check if outputs_info is None ====== #
    output = []
    if outputs_info is not None:
        prev_vals = outputs_info
    else:
        prev_vals = []
    output_idx = [i for i in range(len(prev_vals)) if prev_vals[i] is not None]
    # ====== check if non_sequences is None ====== #
    if non_sequences is None:
        non_sequences = []
    # ====== Main loop ====== #
    for i in counter:
        step_input = [s[i] for s in sequences] + \
                     [prev_vals[idx] for idx in output_idx] + \
            non_sequences
        out_ = step_fn(*step_input)
        # The returned values from step can be either a TensorVariable,
        # a list, or a tuple.  Below, we force it to always be a list.
        if isinstance(out_, tf.python.Tensor):
            out_ = [out_]
        if isinstance(out_, tuple):
            out_ = list(out_)
        output.append(out_)
        prev_vals = output[-1]

    # iterate over each scan output and convert it to same format as scan:
    # [[output11, output12,...output1n],
    # [output21, output22,...output2n],...]
    output_scan = []
    for i in range(len(output[0])):
        l = map(lambda x: x[i], output)
        output_scan.append(tf.pack(l))
    return output_scan

def rnn(step_function, inputs, initial_states,
        go_backwards=False, mask=None, constants=None):
    '''Iterates over the time dimension of a tensor.
    # Arguments
        inputs: tensor of temporal data of shape (samples, time, ...)
            (at least 3D).
        step_function:
            Parameters:
                input: tensor with shape (samples, ...) (no time dimension),
                    representing input for the batch of samples at a certain
                    time step.
                states: list of tensors.
            Returns:
                output: tensor with shape (samples, ...) (no time dimension),
                new_states: list of tensors, same length and shapes
                    as 'states'.
        initial_states: tensor with shape (samples, ...) (no time dimension),
            containing the initial values for the states used in
            the step function.
        go_backwards: boolean. If True, do the iteration over
            the time dimension in reverse order.
        mask: binary tensor with shape (samples, time, 1),
            with a zero for every element that is masked.
        constants: a list of constant values passed at each step.
    # Returns
        A tuple (last_output, outputs, new_states).
            last_output: the latest output of the rnn, of shape (samples, ...)
            outputs: tensor with shape (samples, time, ...) where each
                entry outputs[s, t] is the output of the step function
                at time t for sample s.
            new_states: list of tensors, latest states returned by
                the step function, of shape (samples, ...).
    '''
    ndim = len(inputs.get_shape())
    assert ndim >= 3, "Input should be at least 3D."
    axes = [1, 0] + list(range(2, ndim))
    inputs = tf.transpose(inputs, (axes))
    input_list = tf.unpack(inputs)
    if constants is None:
        constants = []

    states = initial_states
    successive_states = []
    successive_outputs = []
    if go_backwards:
        input_list.reverse()

    if mask is not None:
        # Transpose not supported by bool tensor types, hence round-trip to uint8.
        mask = tf.cast(mask, tf.uint8)
        if len(mask.get_shape()) == ndim - 1:
            mask = expand_dims(mask)
        mask = tf.cast(tf.transpose(mask, axes), tf.bool)
        mask_list = tf.unpack(mask)

        for input, mask_t in zip(input_list, mask_list):
            output, new_states = step_function(input, states + constants)

            # tf.select needs its condition tensor to be the same shape as its two
            # result tensors, but in our case the condition (mask) tensor is
            # (nsamples, 1), and A and B are (nsamples, ndimensions). So we need to
            # broadcast the mask to match the shape of A and B. That's what the
            # tile call does, is just repeat the mask along its second dimension
            # ndimensions times.
            tiled_mask_t = tf.tile(mask_t, tf.pack([1, tf.shape(output)[1]]))

            if len(successive_outputs) == 0:
                prev_output = zeros_like(output)
            else:
                prev_output = successive_outputs[-1]

            output = tf.select(tiled_mask_t, output, prev_output)

            return_states = []
            for state, new_state in zip(states, new_states):
                # (see earlier comment for tile explanation)
                tiled_mask_t = tf.tile(mask_t, tf.pack([1, tf.shape(new_state)[1]]))
                return_states.append(tf.select(tiled_mask_t, new_state, state))

            states = return_states
            successive_outputs.append(output)
            successive_states.append(states)
    else:
        for input in input_list:
            output, states = step_function(input, states + constants)
            successive_outputs.append(output)
            successive_states.append(states)

    last_output = successive_outputs[-1]
    outputs = tf.pack(successive_outputs)
    new_states = successive_states[-1]

    axes = [1, 0] + list(range(2, len(outputs.get_shape())))
    outputs = tf.transpose(outputs, axes)
    return last_output, outputs, new_states

def switch(condition, then_expression, else_expression):
    '''Switch between two operations depending on a scalar value.

    # Arguments
        condition: scalar tensor.
        then_expression: TensorFlow operation.
        else_expression: TensorFlow operation.
    '''
    return tf.python.control_flow_ops.cond(condition,
                                           lambda: then_expression,
                                           lambda: else_expression)

# ===========================================================================
# NN OPERATIONS
# ===========================================================================
def relu(x, alpha=0., max_value=None):
    '''ReLU.
    # Arguments
        alpha: slope of negative section.
        max_value: saturation threshold.
    '''
    negative_part = tf.nn.relu(-x)
    x = tf.nn.relu(x)
    if max_value is not None:
        x = tf.clip_by_value(x, tf.cast(0., dtype=_FLOATX),
                             tf.cast(max_value, dtype=_FLOATX))
    if isinstance(alpha, (tuple, list, np.ndarray)) or np.isscalar(alpha):
        alpha = tf.constant(alpha, dtype=_FLOATX)
    x -= alpha * negative_part
    return x

def linear(x):
    return x

def softmax(x):
    return tf.nn.softmax(x)

def softplus(x):
    return tf.nn.softplus(x)

def categorical_crossentropy(output, target, from_logits=False):
    '''Note: tf.nn.softmax_cross_entropy_with_logits
    expects logits, Keras expects probabilities.
    '''
    if not from_logits:
        # scale preds so that the class probas of each sample sum to 1
        output /= tf.reduce_sum(output,
                                reduction_indices=len(output.get_shape()) - 1,
                                keep_dims=True)
        # manual computation of crossentropy
        output = tf.clip_by_value(output, tf.cast(_EPSILON, dtype=_FLOATX),
                                  tf.cast(1. - _EPSILON, dtype=_FLOATX))
        return - tf.reduce_sum(target * tf.log(output),
                               reduction_indices=len(output.get_shape()) - 1)
    else:
        return tf.nn.softmax_cross_entropy_with_logits(output, target)

def binary_crossentropy(output, target, from_logits=False):
    '''Note: tf.nn.sigmoid_cross_entropy_with_logits
    expects logits, Keras expects probabilities.
    '''
    if not from_logits:
        # transform back to logits
        output = tf.clip_by_value(output, tf.cast(_EPSILON, dtype=_FLOATX),
                                  tf.cast(1. - _EPSILON, dtype=_FLOATX))
        output = tf.log(output / (1 - output))
    return tf.nn.sigmoid_cross_entropy_with_logits(output, target)

def sigmoid(x):
    return tf.nn.sigmoid(x)

def hard_sigmoid(x):
    x = (0.2 * x) + 0.5
    x = tf.clip_by_value(x, tf.cast(0., dtype=_FLOATX),
                         tf.cast(1., dtype=_FLOATX))
    return x

def tanh(x):
    return tf.nn.tanh(x)

def dropout(x, level, rescale=True, noise_shape=None,
    seed=None, rng=None):
    """Computes dropout.

    With probability `keep_prob`, outputs the input element scaled up by
    `1 / keep_prob`, otherwise outputs `0`.  The scaling is so that the expected
    sum is unchanged.

    By default, each element is kept or dropped independently.  If `noise_shape`
    is specified, it must be
    [broadcastable](http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html)
    to the shape of `x`, and only dimensions with `noise_shape[i] == shape(x)[i]`
    will make independent decisions.  For example, if `shape(x) = [k, l, m, n]`
    and `noise_shape = [k, 1, 1, n]`, each batch and channel component will be
    kept independently and each row and column will be kept or not kept together.

    Parameters
    ----------
    x: A tensor.
    level: float(0.-1.)
        probability dropout values in given tensor
    rescale: bool
        whether rescale the outputs by dividing the retain probablity
    noise_shape: A 1-D `Tensor` of type `int32`, representing the
      shape for randomly generated keep/drop flags.
    seed: int
        A Python integer. Used to create random seeds. See
    rng: `tensor.rng`
        random generator from tensor class
    """
    retain_prob = 1. - level
    if isinstance(rng, _RandomWrapper):
        seed = rng._rng.randint(10e6)
    elif seed is None:
        seed = get_random_magic_seed()

    if noise_shape is not None:
        # from tensorflow.python.ops import array_ops
        # shape_x = array_ops.shape(x)
        noise_shape = tuple([shape(x)[i].value if j is None or j < 0 else j
                            for i, j in enumerate(noise_shape)])
    # the dummy 1. works around a TF bug
    # (float32_ref vs. float32 incomptability)
    x = tf.nn.dropout(x * 1., retain_prob, noise_shape=noise_shape, seed=seed)
    if not rescale:
        x = x * retain_prob
    return x

# ==================== Regularizations ==================== #

def l2_normalize(x, axis):
    if axis < 0:
        axis = axis % len(x.get_shape())
    return tf.nn.l2_normalize(x, dim=axis)

def l2_regularize(x):
    return sum(tf.square(x))

def l1_regularize(x):
    return sum(tf.abs(x))

def jacobian_regularize(hidden, params):
    ''' Computes the jacobian of the hidden layer with respect to
    the input, reshapes are necessary for broadcasting the
    element-wise product on the right axis
    '''
    hidden = hidden * (1 - hidden)
    L = expand_dims(hidden, 1) * expand_dims(params, 0)
    # Compute the jacobian and average over the number of samples/minibatch
    L = sum(mean(tf.pow(L, 2), axis=0)) # avr over all samples in batch
    return mean(L)

def kl_gaussian(mean_, logsigma,
                prior_mean=0., prior_logsigma=0.,
                regularizer_scale=1.):
    ''' KL-divergence between two gaussians.
    Useful for Variational AutoEncoders. Use this as an activation regularizer
    Parameters:
    -----------
    mean, logsigma: parameters of the input distributions
    prior_mean, prior_logsigma: paramaters of the desired distribution (note the
        log on logsigma)
    regularizer_scale: Rescales the regularization cost. Keep this 1 for most cases.

    Note
    ----
    origin implementation from seya:
    https://github.com/Philip-Bachman/ICML-2015/blob/master/LogPDFs.py
    Copyright (c) Philip Bachman
    '''
    gauss_klds = 0.5 * (prior_logsigma - logsigma +
            ((tf.exp(logsigma) + pow((mean_ - prior_mean), 2.0)) / tf.exp(prior_logsigma)) - 1.0)
    return mean(gauss_klds)

def correntropy_regularize(x, sigma=1.):
    '''
    Note
    ----
    origin implementation from seya:
    https://github.com/EderSantana/seya/blob/master/seya/regularizers.py
    Copyright (c) EderSantana
    '''
    return -sum(mean(tf.exp(tf.pow(x, 2) / sigma), axis=0)) / tf.sqrt(2 * np.pi * sigma)

# ===========================================================================
# CONVOLUTIONS
# ===========================================================================

def conv2d(x, kernel, strides=(1, 1), border_mode='valid', dim_ordering='th',
           image_shape=None, filter_shape=None):
    '''Runs on cuDNN if available.

    # Arguments
        border_mode: string, "same" or "valid".
        dim_ordering: whether to use Theano or TensorFlow dimension ordering
        in inputs/kernels/ouputs.
    '''
    if border_mode == 'same':
        padding = 'SAME'
    elif border_mode == 'valid':
        padding = 'VALID'
    else:
        raise Exception('Invalid border mode: ' + str(border_mode))

    strides = (1,) + strides + (1,)

    if _FLOATX == 'float64':
        # tf conv2d only supports float32
        x = tf.cast(x, 'float32')
        kernel = tf.cast(kernel, 'float32')

    if dim_ordering == 'th':
        # TF uses the last dimension as channel dimension,
        # instead of the 2nd one.
        # TH input shape: (samples, input_depth, rows, cols)
        # TF input shape: (samples, rows, cols, input_depth)
        # TH kernel shape: (depth, input_depth, rows, cols)
        # TF kernel shape: (rows, cols, input_depth, depth)
        x = tf.transpose(x, (0, 2, 3, 1))
        kernel = tf.transpose(kernel, (2, 3, 1, 0))
        x = tf.nn.conv2d(x, kernel, strides, padding=padding)
        x = tf.transpose(x, (0, 3, 1, 2))
    elif dim_ordering == 'tf':
        x = tf.nn.conv2d(x, kernel, strides, padding=padding)
    else:
        raise Exception('Unknown dim_ordering: ' + str(dim_ordering))

    if _FLOATX == 'float64':
        x = tf.cast(x, 'float64')
    return x

def conv3d(x, kernel, strides=(1, 1, 1), border_mode='valid', dim_ordering='th',
           image_shape=None, filter_shape=None):
    raise NotImplementedError

def pool2d(x, pool_size, strides=(1, 1),
           border_mode='valid', dim_ordering='th', pool_mode='max'):
    '''
    # Arguments
        pool_size: tuple of 2 integers.
        strides: tuple of 2 integers.
        border_mode: one of "valid", "same".
        dim_ordering: one of "th", "tf".
    '''
    if border_mode == 'same':
        padding = 'SAME'
    elif border_mode == 'valid':
        padding = 'VALID'
    else:
        raise Exception('Invalid border mode: ' + str(border_mode))

    strides = (1,) + strides + (1,)
    pool_size = (1,) + pool_size + (1,)

    if _FLOATX == 'float64':
        # tf max_pool only supports float32
        x = tf.cast(x, 'float32')

    if dim_ordering in {'tf', 'th'}:
        if dim_ordering == 'th':
            # TF uses the last dimension as channel dimension,
            # instead of the 2nd one.
            # TH input shape: (samples, input_depth, rows, cols)
            # TF input shape: (samples, rows, cols, input_depth)
            # TH kernel shape: (depth, input_depth, rows, cols)
            # TF kernel shape: (rows, cols, input_depth, depth)
            x = tf.transpose(x, (0, 2, 3, 1))
        if pool_mode == 'max':
            x = tf.nn.max_pool(x, pool_size, strides, padding=padding)
        elif pool_mode == 'avg':
            x = tf.nn.avg_pool(x, pool_size, strides, padding=padding)
        else:
            raise Exception('Invalid pooling mode: ' + str(pool_mode))
        if dim_ordering == 'th':
            x = tf.transpose(x, (0, 3, 1, 2))
    else:
        raise Exception('Unknown dim_ordering: ' + str(dim_ordering))

    if _FLOATX == 'float64':
        x = tf.cast(x, 'float64')
    return x

def pool3d(x, pool_size, strides=(1, 1, 1),
           border_mode='valid', dim_ordering='th', pool_mode='max'):
    raise NotImplementedError

# ===========================================================================
# RANDOMNESS
# ===========================================================================

class _RandomWrapper(object):

    def __init__(self, rng):
        super(_RandomWrapper, self).__init__()
        self._rng = np.random.RandomState(rng)
        self._state = np.random.RandomState(rng)

    def randint(self):
        return self._state.randint(10e6)

    def normal(self, shape, mean, std, dtype=_FLOATX):
        return tf.random_normal(shape=shape, mean=mean, stddev=std,
            dtype=dtype, seed=self._rng.randint(10e6))

    def uniform(self, shape, low, high, dtype=_FLOATX):
        return tf.random_uniform(shape=shape, minval=low, maxval=high,
                             dtype=dtype, seed=self._rng.randint(10e6))

    def binomial(self, shape, p, dtype=_FLOATX):
        return tf.cast(
            tf.less(
                tf.random_uniform(shape=shape, minval=0., maxval=1.,
                             dtype=_FLOATX, seed=self._rng.randint(10e6)),
                p),
            dtype)

def rng(seed=None):
    if seed is None:
        seed = get_random_magic_seed()
    return _RandomWrapper(seed)

def random_normal(shape, mean=0.0, std=1.0, dtype=_FLOATX, seed=None):
    if seed is None:
        seed = get_random_magic_seed()
    return tf.random_normal(shape, mean=mean, stddev=std,
                            dtype=dtype, seed=seed)

def random_uniform(shape, low=0.0, high=1.0, dtype=_FLOATX, seed=None):
    if seed is None:
        seed = get_random_magic_seed()
    return tf.random_uniform(shape, minval=low, maxval=high,
                             dtype=dtype, seed=seed)

def random_binomial(shape, p, dtype=_FLOATX, seed=None):
    if seed is None:
        seed = get_random_magic_seed()
    return tf.cast(
        tf.less(
            tf.random_uniform(shape=shape, minval=0., maxval=1.,
                         dtype=dtype, seed=seed),
            p),
        dtype)

# ===========================================================================
# Comparator
# ===========================================================================

def eq(x, y):
    """a == b"""
    return tf.equal(x, y)

def neq(x, y):
    """a != b"""
    return tf.not_equal(x, y)

def gt(a, b):
    """a > b"""
    return tf.greater(a, b)

def ge(a, b):
    """a >= b"""
    return tf.greater_equal(a, b)

def lt(a, b):
    """a < b"""
    return tf.less(a, b)

def le(a, b):
    """a <= b"""
    return tf.less_equal(a, b)

def one_hot(x, nb_class):
    ''' x: 1D-integer vector '''
    shape = x.get_shape()
    ret = tf.zeros((shape[0].value, nb_class), dtype=_FLOATX)
    ret[np.arange(shape[0].value), x] = 1
    return ret

def one_hot_max(x, axis=-1):
    '''
    Example
    -------
    >>> Input: [[0.0, 0.0, 0.5],
    >>>         [0.0, 0.3, 0.1],
    >>>         [0.6, 0.0, 0.2]]
    >>> Output: [[0.0, 0.0, 1.0],
    >>>         [0.0, 1.0, 0.0],
    >>>         [1.0, 0.0, 0.0]]
    '''
    if axis < 0:
        axis = axis % len(x.get_shape())
    shape = x.get_shape()[axis].value
    return tf.cast(
        tf.equal(tf.cast(tf.range(shape), 'int64'),
                expand_dims(tf.argmax(x, axis))),
        _FLOATX
    )

def apply_mask(x, mask):
    '''
    x : 3D tensor
    mask : 2D tensor

    Example
    -------
    >>> Input: [128, 500, 120]
    >>> Mask:  [1, 1, 0]
    >>> Output: [128, 500, 0]
    '''
    return tf.mul(x, tf.expand_dims(mask, -1))


import math
from collections import OrderedDict

import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F

'''
AdaCos and Ad margin loss taken from https://github.com/4uiiurz1/pytorch-adacos
'''

class DropClassBase(nn.Module):

    def __init__(self, num_classes):
        '''
        DropClass class which other classifier heads should inherit from

        This is to package the useful wrapper scripts for which classes to include/ignore

        The class has two main modes, called via .drop() and .nodrop(), which sets which method will be
        called by .forward()

        forward_drop defines the ordinary behaviour
        forward_nodrop defines the behaviour in which only the remaining class columns are used
        '''
        super(DropClassBase, self).__init__()
        self.n_classes = num_classes
        self.dropmode = False # Default is the normal behaviour
        self.set_ignored_classes([])
        self.combined_class_label = None

    def forward(self, input, label=None):
        '''
        input: (batch_size, num_features): FloatTensor
        label (optional): (batch_size): LongTensor
        '''
        if self.dropmode:
            if label is not None:
                assert (torch.max(label) < len(self.rem_classes)), 'Contains label out of range of allowed classes: Have they been converted?'
            return self.forward_drop(input, label=label)
        else:
            return self.forward_nodrop(input, label=label)

    def drop(self):
        self.dropmode = True

    def nodrop(self):
        self.dropmode = False

    def forward_drop(self, input, label=None):
        raise NotImplementedError

    def forward_nodrop(self, input, label=None):
        raise NotImplementedError

    def set_ignored_classes(self, ignored:list):
        if len(ignored) != 0:
            assert min(ignored) >= 0
            assert max(ignored) < self.n_classes
        self.ignored = sorted(list(set(ignored)))
        self.rem_classes = sorted(set(np.arange(self.n_classes)) - set(ignored))
        self.ldict = OrderedDict({k:v for v, k in enumerate(self.rem_classes)}) #mapping of original label to new index
        self.idict = OrderedDict({k:v for k, v in enumerate(self.rem_classes)}) #mapping of remaining indexes to original label

    def set_remaining_classes(self, remaining:list):
        assert min(remaining) >= 0
        assert max(remaining) < self.n_classes
        self.rem_classes = sorted(set(remaining))
        self.ignored = sorted(set(np.arange(self.n_classes)) - set(remaining))
        self.ldict = OrderedDict({k:v for v, k in enumerate(self.rem_classes)}) #mapping of original label to new index
        self.idict = OrderedDict({k:v for k, v in enumerate(self.rem_classes)}) #mapping of remaining indexes to original label

    def get_mini_labels(self, label:list):
        # convert list of labels into new indexes for ignored classes
        mini_labels = torch.LongTensor(list(map(lambda x: self.ldict[x], label)))
        return mini_labels

    def get_orig_labels(self, label:list):
        # convert list of mini_labels into original class labels
        # assert not self.combined_class_label, 'Combined classes means original labels not recoverable'
        orig_labels = list(map(lambda x: self.idict[x], label))
        return orig_labels

    def set_remaining_classes_comb(self, remaining:list):
        # remaining must not include the combined class
        assert self.combined_class_label is not None, 'combined_class_label has not been set'
        assert min(remaining) >= 0
        assert max(remaining) < self.n_classes
        remaining.append(self.combined_class_label)
        self.rem_classes = sorted(set(remaining))
        self.ignored = sorted(set(np.arange(self.n_classes)) - set(remaining)) # not really ignored, just combined
        self.ldict = OrderedDict({k:v for v, k in enumerate(self.rem_classes)})
        for k in self.ignored:
            self.ldict[k] = self.combined_class_label # set all ignored classes to the combined class label
        self.idict = OrderedDict({k:v for k, v in enumerate(self.rem_classes)}) # not the original mapping for comb classes

class DropAffine(DropClassBase):

    def __init__(self, num_features, num_classes):
        super(DropAffine, self).__init__(num_classes)
        self.fc = nn.Linear(num_features, num_classes)
        self.reset_parameters()

    def reset_parameters(self):
        self.fc.reset_parameters()

    def forward_nodrop(self, input, label=None):
        W = self.fc.weight
        b = self.fc.bias
        logits = F.linear(input, W, b)
        return logits

    def forward_drop(self, input, label=None):
        W = self.fc.weight[self.rem_classes]
        b = self.fc.bias[self.rem_classes]
        logits = F.linear(input, W, b)
        return logits

class L2SoftMax(DropClassBase):

    def __init__(self, num_features, num_classes):
        super(L2SoftMax, self).__init__(num_classes)
        self.W = nn.Parameter(torch.FloatTensor(num_classes, num_features))
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.xavier_uniform_(self.W)

    def forward_nodrop(self, input, label=None):
        x = F.normalize(input)
        W = F.normalize(self.W)
        logits = F.linear(x, W)
        return logits

    def forward_drop(self, input, label=None):
        x = F.normalize(input)
        W = F.normalize(self.W[self.rem_classes])
        logits = F.linear(x, W)
        return logits

class SoftMax(DropClassBase):

    def __init__(self, num_features, num_classes):
        super(SoftMax, self).__init__(num_classes)
        self.W = nn.Parameter(torch.FloatTensor(num_classes, num_features))
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.xavier_uniform_(self.W)

    def forward_nodrop(self, input, label=None):
        x = input
        W = self.W
        logits = F.linear(x, W)
        return logits

    def forward_drop(self, input, label=None):
        x = input
        W = self.W[self.rem_classes]
        logits = F.linear(x, W)
        return logits

class XVecHead(DropClassBase):

    def __init__(self, num_features, num_classes, hidden_features=None):
        super(XVecHead, self).__init__(num_classes)
        hidden_features = num_features if not hidden_features else hidden_features
        self.fc_hidden = nn.Linear(num_features, hidden_features)
        self.nl = nn.LeakyReLU()
        self.bn = nn.BatchNorm1d(hidden_features)
        self.fc = nn.Linear(hidden_features, num_classes)
        self.reset_parameters()

    def reset_parameters(self):
        self.fc.reset_parameters()

    def forward_nodrop(self, input, label=None):
        input = self.fc_hidden(input)
        input = self.nl(input)
        input = self.bn(input)
        W = self.fc.weight
        b = self.fc.bias
        logits = F.linear(input, W, b)
        return logits

    def forward_drop(self, input, label=None):
        input = self.fc_hidden(input)
        input = self.nl(input)
        input = self.bn(input)
        W = self.fc.weight[self.rem_classes]
        b = self.fc.bias[self.rem_classes]
        logits = F.linear(input, W, b)
        return logits

class AMSMLoss(DropClassBase):

    def __init__(self, num_features, num_classes, s=30.0, m=0.4):
        super(AMSMLoss, self).__init__(num_classes)
        self.num_features = num_features
        self.n_classes = num_classes
        self.s = s
        self.m = m
        self.W = nn.Parameter(torch.FloatTensor(num_classes, num_features))
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.xavier_uniform_(self.W)

    def forward_nodrop(self, input, label=None):
        # normalize features
        x = F.normalize(input)
        # normalize weights
        W = F.normalize(self.W)
        # dot product
        logits = F.linear(x, W)
        if label is None:
            return logits
        # add margin
        target_logits = logits - self.m
        one_hot = torch.zeros_like(logits)
        one_hot.scatter_(1, label.view(-1, 1).long(), 1)
        output = logits * (1 - one_hot) + target_logits * one_hot
        # feature re-scale
        output *= self.s

        return output

    def forward_drop(self, input, label=None):
        # normalize features
        x = F.normalize(input)
        # normalize weights
        W = F.normalize(self.W[self.rem_classes])
        # dot product
        logits = F.linear(x, W)
        if label is None:
            return logits
        # add margin
        target_logits = logits - self.m
        one_hot = torch.zeros_like(logits)
        one_hot.scatter_(1, label.view(-1, 1).long(), 1)
        output = logits * (1 - one_hot) + target_logits * one_hot
        # feature re-scale
        output *= self.s

        return output

class SphereFace(DropClassBase):

    def __init__(self, num_features, num_classes, s=30.0, m=1.35):
        super(SphereFace, self).__init__(num_classes)
        self.num_features = num_features
        self.n_classes = num_classes
        self.s = s
        self.m = m
        self.W = nn.Parameter(torch.FloatTensor(num_classes, num_features))
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.xavier_uniform_(self.W)

    def forward_nodrop(self, input, label=None):
        # normalize features
        x = F.normalize(input)
        # normalize weights
        W = F.normalize(self.W)
        # dot product
        logits = F.linear(x, W)
        if label is None:
            return logits
        # add margin
        theta = torch.acos(torch.clamp(logits, -1.0 + 1e-7, 1.0 - 1e-7))
        target_logits = torch.cos(self.m * theta)
        one_hot = torch.zeros_like(logits)
        one_hot.scatter_(1, label.view(-1, 1).long(), 1)
        output = logits * (1 - one_hot) + target_logits * one_hot
        # feature re-scale
        output *= self.s

        return output

    def forward_drop(self, input, label=None):
        # normalize features
        x = F.normalize(input)
        # normalize weights
        W = F.normalize(self.W[self.rem_classes])
        # dot product
        logits = F.linear(x, W)
        if label is None:
            return logits
        # add margin
        theta = torch.acos(torch.clamp(logits, -1.0 + 1e-7, 1.0 - 1e-7))
        target_logits = torch.cos(self.m * theta)
        one_hot = torch.zeros_like(logits)
        one_hot.scatter_(1, label.view(-1, 1).long(), 1)
        output = logits * (1 - one_hot) + target_logits * one_hot
        # feature re-scale
        output *= self.s

        return output

class ArcFace(DropClassBase):

    def __init__(self, num_features, num_classes, s=30.0, m=0.50):
        super(ArcFace, self).__init__(num_classes)
        self.num_features = num_features
        self.n_classes = num_classes
        self.s = s
        self.m = m
        self.W = nn.Parameter(torch.FloatTensor(num_classes, num_features))
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.xavier_uniform_(self.W)

    def forward_nodrop(self, input, label=None):
        # normalize features
        x = F.normalize(input)
        # normalize weights
        W = F.normalize(self.W)
        # dot product
        logits = F.linear(x, W)
        if label is None:
            return logits
        # add margin
        theta = torch.acos(torch.clamp(logits, -1.0 + 1e-7, 1.0 - 1e-7))
        target_logits = torch.cos(theta + self.m)
        one_hot = torch.zeros_like(logits)
        one_hot.scatter_(1, label.view(-1, 1).long(), 1)
        output = logits * (1 - one_hot) + target_logits * one_hot
        # feature re-scale
        output *= self.s

        return output

    def forward_drop(self, input, label=None):
        # normalize features
        x = F.normalize(input)
        # normalize weights
        W = F.normalize(self.W[self.rem_classes])
        # dot product
        logits = F.linear(x, W)
        if label is None:
            return logits
        # add margin
        theta = torch.acos(torch.clamp(logits, -1.0 + 1e-7, 1.0 - 1e-7))
        target_logits = torch.cos(theta + self.m)
        one_hot = torch.zeros_like(logits)
        one_hot.scatter_(1, label.view(-1, 1).long(), 1)
        output = logits * (1 - one_hot) + target_logits * one_hot
        # feature re-scale
        output *= self.s

        return output

class AdaCos(DropClassBase):

    def __init__(self, num_features, num_classes, m=0.50):
        super(AdaCos, self).__init__(num_classes)
        self.num_features = num_features
        self.n_classes = num_classes
        self.s = math.sqrt(2) * math.log(num_classes - 1)
        self.m = m
        self.W = nn.Parameter(torch.FloatTensor(num_classes, num_features))
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.xavier_uniform_(self.W)

    def forward_nodrop(self, input, label=None):
        # normalize features
        x = F.normalize(input)
        # normalize weights
        W = F.normalize(self.W)
        # dot product
        logits = F.linear(x, W)
        if label is None:
            return logits
        # feature re-scale
        theta = torch.acos(torch.clamp(logits, -1.0 + 1e-7, 1.0 - 1e-7))
        one_hot = torch.zeros_like(logits)
        one_hot.scatter_(1, label.view(-1, 1).long(), 1)
        with torch.no_grad():
            B_avg = torch.where(one_hot < 1, torch.exp(self.s * logits), torch.zeros_like(logits))
            B_avg = torch.sum(B_avg) / input.size(0)
            theta_med = torch.median(theta[one_hot == 1])
            self.s = torch.log(B_avg) / torch.cos(torch.min(math.pi/4 * torch.ones_like(theta_med), theta_med))
        output = self.s * logits

        return output

    def forward_drop(self, input, label=None):
        # normalize features
        x = F.normalize(input)
        # normalize weights
        W = F.normalize(self.W[self.rem_classes])
        # dot product
        logits = F.linear(x, W)
        if label is None:
            return logits
        # feature re-scale
        theta = torch.acos(torch.clamp(logits, -1.0 + 1e-7, 1.0 - 1e-7))
        one_hot = torch.zeros_like(logits)
        one_hot.scatter_(1, label.view(-1, 1).long(), 1)
        with torch.no_grad():
            B_avg = torch.where(one_hot < 1, torch.exp(self.s * logits), torch.zeros_like(logits))
            B_avg = torch.sum(B_avg) / input.size(0)
            theta_med = torch.median(theta[one_hot == 1])
            self.s = torch.log(B_avg) / torch.cos(torch.min(math.pi/4 * torch.ones_like(theta_med), theta_med))
        output = self.s * logits

        return output

class DisturbLabelLoss(nn.Module):

    def __init__(self, device, disturb_prob=0.1):
        super(DisturbLabelLoss, self).__init__()
        self.disturb_prob = disturb_prob
        self.ce = nn.CrossEntropyLoss()
        self.device = device

    def forward(self, pred, target):
        with torch.no_grad():
            disturb_indexes = torch.rand(len(pred)) < self.disturb_prob
            target[disturb_indexes] = torch.randint(pred.shape[-1], (int(disturb_indexes.sum()),)).to(self.device)
        return self.ce(pred, target)

class LabelSmoothingLoss(nn.Module):

    def __init__(self, smoothing=0.1, dim=-1):
        super(LabelSmoothingLoss, self).__init__()
        self.confidence = 1.0 - smoothing
        self.smoothing = smoothing
        self.dim = dim

    def forward(self, pred, target):
        pred = pred.log_softmax(dim=self.dim)
        with torch.no_grad():
            true_dist = torch.zeros_like(pred)
            true_dist.fill_(self.smoothing / (pred.shape[-1] - 1))
            true_dist.scatter_(1, target.data.unsqueeze(1), self.confidence)
        return torch.mean(torch.sum(-true_dist * pred, dim=self.dim))


#!/usr/bin/python
# -*- coding: UTF-8 -*-

import matplotlib
matplotlib.use('Agg')

# Pytorch requirements
import unicodedata
import string
import re
import random

import torch
import torch.nn as nn
from torch.nn import init
from torch.autograd import Variable
from torch import optim
import torch.nn.functional as F
from MatchingLayer import MatchingLayer
import numpy as np

if torch.cuda.is_available():
    dtype = torch.cuda.FloatTensor
    dtype_l = torch.cuda.LongTensor
else:
    dtype = torch.FloatTensor
    dtype_l = torch.cuda.LongTensor

def sinkhorn_knopp(A, iterations=1):
    A_size = A.size()
    for it in range(iterations):
        A = A.view(A_size[0]*A_size[1], A_size[2])
        A = F.softmax(A)
        A = A.view(*A_size).permute(0, 2, 1)
        A = A.view(A_size[0]*A_size[1], A_size[2])
        A = F.softmax(A)
        A = A.view(*A_size).permute(0, 2, 1)
    return A

def gmul(input):
    W, x = input
    # x is a tensor of size (bs, N, num_features)
    # W is a tensor of size (bs, N, N, J)
    x_size = x.size()
    W_size = W.size()
    N = W_size[-2]
    J = W_size[-1]
    W = W.split(1, 3)
    W = torch.cat(W, 1).squeeze(3) # W is now a tensor of size (bs, J*N, N)
    output = torch.bmm(W, x) # output has size (bs, J*N, num_features)
    output = output.split(N, 1)
    output = torch.cat(output, 2) # output has size (bs, N, J*num_features)
    return output

class Gconv_last(nn.Module):
    def __init__(self, feature_maps, J):
        super(Gconv_last, self).__init__()
        self.num_inputs = J*feature_maps[0]
        self.num_outputs = feature_maps[2]
        self.fc = nn.Linear(self.num_inputs, self.num_outputs)

    def forward(self, input):
        W = input[0]
        x = gmul(input) # out has size (bs, N, num_inputs)
        x_size = x.size()
        x = x.contiguous()
        x = x.view(x_size[0]*x_size[1], -1)
        x = self.fc(x) # has size (bs*N, num_outputs)
        x = x.view(*x_size[:-1], self.num_outputs)
        return W, x

class Gconv(nn.Module):
    def __init__(self, feature_maps, J):
        super(Gconv, self).__init__()
        self.num_inputs = J*feature_maps[0]
        self.num_outputs = feature_maps[2]
        self.fc1 = nn.Linear(self.num_inputs, self.num_outputs // 2)
        self.fc2 = nn.Linear(self.num_inputs, self.num_outputs // 2)
        self.bn = nn.BatchNorm1d(self.num_outputs)

    def forward(self, input):
        W = input[0]
        x = gmul(input) # out has size (bs, N, num_inputs)
        x_size = x.size()
        x = x.contiguous()
        x = x.view(-1, self.num_inputs)
        x1 = F.relu(self.fc1(x)) # has size (bs*N, num_outputs)
        x2 = self.fc2(x)
        x = torch.cat((x1, x2), 1)
        x = self.bn(x)
        x = x.view(*x_size[:-1], self.num_outputs)
        return W, x

class GNN(nn.Module):
    def __init__(self, num_features, num_layers, J):
        super(GNN, self).__init__()
        self.num_features = num_features
        self.num_layers = num_layers
        self.featuremap_in = [1, 1, num_features]
        self.featuremap_mi = [num_features, num_features, num_features]
        self.featuremap_end = [num_features, num_features, num_features]
        self.layer0 = Gconv(self.featuremap_in, J)
        for i in range(num_layers):
            module = Gconv(self.featuremap_mi, J)
            self.add_module('layer{}'.format(i + 1), module)
        self.layerlast = Gconv_last(self.featuremap_end, J)

    def forward(self, input):
        cur = self.layer0(input)
        for i in range(self.num_layers):
            cur = self._modules['layer{}'.format(i+1)](cur)
        out = self.layerlast(cur)
        return out[1]

class GNN_Matcher(nn.Module):
    def __init__(self, num_features, num_layers, J, matching_layer):
        super(GNN_Matcher, self).__init__()
        self.num_features = num_features
        self.num_layers = num_layers
        self.featuremap_in = [1, 1, num_features]
        self.featuremap_mi = [num_features, num_features, num_features]
        self.featuremap_end = [num_features, num_features, num_features]
        self.layer0 = Gconv(self.featuremap_in, J)
        self.matching_layer = matching_layer
        for i in range(num_layers):
            module = Gconv(self.featuremap_mi, J)
            self.add_module('layer{}'.format(i + 1), module)
        self.layerlast = Gconv_last(self.featuremap_end, J)

    def forward(self, input):
        cur = self.layer0(input)
        for i in range(self.num_layers):
            cur = self._modules['layer{}'.format(i+1)](cur)

        out = self.layerlast(cur)

        x = out[1]
        outsize = x.size()

        x = x.sum(2)   # sum along the axis of 'features' from the GNN
        outsize = x.size()

        """
        print("after sum")
        print("outsize")
        print( outsize )
        """

        ns = np.ceil(np.sqrt(x.size()[1])).astype(int)

        x = self.matching_layer(x.view(1,-1)).view(outsize)
        #x = self.matching_layer(x.view(1,ns,ns))#.view(outsize)

        return x.view(1,ns,ns)
        #return out[1]

class Siamese_GNN(nn.Module):
    def __init__(self, num_features, num_layers, J):
        super(Siamese_GNN, self).__init__()
        self.gnn = GNN(num_features, num_layers, J)

    def forward(self, g1, g2):
        """
        print("Inside Siamese GNN forward")
        print("g1 = ")
        print( g1 )
        print("g2 = ")
        print( g2 )
        """
        emb1 = self.gnn(g1)
        emb2 = self.gnn(g2)
        """
        print("emb1 = ")
        print( emb1 )
        print("emb2 = ")
        print( emb2 )
        print("emb2.size() = ")
        print( emb2.size() )
        """

        # embx are tensors of size (bs, N, num_features)
        out = torch.bmm(emb1, emb2.permute(0, 2, 1))
        return out # out has size (bs, N, N)

class Siamese_Matcher(nn.Module):
    def __init__(self, num_features, num_layers, J, matching_layer):
        super(Siamese_Matcher, self).__init__()
        self.gnn = GNN(num_features, num_layers, J)
        self.matching_layer = matching_layer  #MatchingLayer(nNodes=n, eps=1e-4)

    def forward(self, g1, g2):

        emb1 = self.gnn(g1)
        emb2 = self.gnn(g2)

        # embx are tensors of size (bs, N, num_features)
        out = torch.bmm(emb1, emb2.permute(0, 2, 1))
        outsize = out.size()
        out = self.matching_layer(out.view(1,-1)).view(outsize)
        # adjust training loop to accomodate this output
        return out # out has size (bs, N, N)

if __name__ == '__main__':
    # test modules
    bs =  4
    num_features = 10
    num_layers = 5
    N = 8
    x = torch.ones((bs, N, num_features))
    W1 = torch.eye(N).unsqueeze(0).unsqueeze(-1).expand(bs, N, N, 1)
    W2 = torch.ones(N).unsqueeze(0).unsqueeze(-1).expand(bs, N, N, 1)
    J = 2
    W = torch.cat((W1, W2), 3)
    input = [Variable(W), Variable(x)]
    ######################### test gmul ##############################
    # feature_maps = [num_features, num_features, num_features]
    # out = gmul(input)
    # print(out[0, :, num_features:])
    ######################### test gconv ##############################
    # feature_maps = [num_features, num_features, num_features]
    # gconv = Gconv(feature_maps, J)
    # _, out = gconv(input)
    # print(out.size())
    ######################### test gnn ##############################
    # x = torch.ones((bs, N, 1))
    # input = [Variable(W), Variable(x)]
    # gnn = GNN(num_features, num_layers, J)
    # out = gnn(input)
    # print(out.size())
    ######################### test siamese gnn ##############################
    x = torch.ones((bs, N, 1))
    input1 = [Variable(W), Variable(x)]
    input2 = [Variable(W.clone()), Variable(x.clone())]
    siamese_gnn = Siamese_GNN(num_features, num_layers, J)
    out = siamese_gnn(input1, input2)
    print(out.size())


"""
.. class:: SpectraStackingEBOSS

.. moduleauthor:: Johan Comparat <johan.comparat__at__gmail.com>

The class SpectraStacking is dedicated to stacking spectra from SDSS-IV eBOSS

Current version

"""
import os
import astropy.io.fits as fits
import numpy as n
from scipy.interpolate import interp1d
import spectres as sp

maskLambda = n.loadtxt(os.path.join(os.environ['GIT_ARCHETYPES'],'data',"dr12-sky-mask.txt"), unpack=True)

get_path_to_spectrum_v5_13_0 = lambda plate, mjd, fiberid : os.path.join(
	os.environ['HOME'],
	'SDSS',
	'v5_13_0',
	'spectra',
	str(int(plate)).zfill(4),
	"spec-"+str(int(plate)).zfill(4)+"-"+str(int(mjd)).zfill(5)+"-"+str(int(fiberid)).zfill(4)+".fits"
	)

get_path_to_spectrum_26 = lambda plate, mjd, fiberid : os.path.join(
	os.environ['HOME'],
	'SDSS',
	'26',
	'spectra',
	str(int(plate)).zfill(4),
	"spec-"+str(int(plate)).zfill(4)+"-"+str(int(mjd)).zfill(5)+"-"+str(int(fiberid)).zfill(4)+".fits" )

line_list_abs = n.array([2249.88, 2260.78, 2344.21, 2374.46, 2382.76, 2576.88, 2586.65, 2594.50, 2600.17, 2606.46, 2796.35, 2803.53, 2852.96])
line_list_abs_names = n.array(['FeII', 'FeII', 'FeII', 'FeII', 'FeII', 'MnII', 'FeII', 'MnII','FeII', 'MnII', 'MgII','MgII','MgI'])
line_list_em = n.array([2327, 2365.55, 2396.36, 2612.65,2626.45])
line_list_em_names = n.array(['CII]', 'FeII*', 'FeII*', 'FeII*', 'FeII*'])

class SpectraStackingEBOSS:
	"""
	The model luminosity function class
	:param in_file: file containing spectra ids to be stacked
	:param Resolution: Resolution
	:param out_file: where to output stacks
	"""
	def __init__(self, in_file, out_file, dLambda = 0.0001, dV=-9999, l_start=2.9, l_end=4.04, KZ_input=False, PBKT_input=False, csv_input=False):
		print( "input list:", in_file )
		self.in_file = in_file
		if KZ_input :
			print('KZ input')
			self.mjds, self.plates, self.fiberids, self.redshifts = n.loadtxt(self.in_file, unpack=True)
		elif PBKT_input :
			print('PBKT input')
			self.plates, self.mjds, self.fiberids, self.redshifts, self.weights = n.loadtxt(self.in_file, unpack=True)
		elif csv_input :
			print('csv input list')
			self.plates, self.mjds, self.fiberids, self.redshifts = n.loadtxt(self.in_file, unpack=True, delimiter=',', skiprows=1)
		else:
			print('regular input list')
			self.plates, self.mjds, self.fiberids, self.redshifts = n.loadtxt(self.in_file, unpack=True)

		print('N spectra = ', len(self.plates))
		self.out_file = out_file
		self.dLambda = dLambda
		#self.wave= 10**n.arange(2.6, 4.0211892990699383, dLambda) # 1500,10500
		self.wave= 10**n.arange(l_start, l_end, dLambda) # 1500,10500
		print('wavelength array', self.wave)
		self.R = int(1/n.mean((self.wave[1:] -self.wave[:-1])/ self.wave[1:]))
		print( "R=", n.median(self.R) )
		self.dV = dV
		self.survey="eBOSS"
		self.N_angstrom_masked = 20.
		#self.run2d = run2d
		#self.run1d = self.run2d
		#self.topdirBOSS = os.path.join(os.environ['BOSS_SPECTRO_REDUX'], run2d)

	def stack_function(self,specMatrix,specMatrixWeight):
		"""Creates the stack.
		:param specMatrix: matrix of observed spectra
		:param specMatrixWeight: matrix of the statistical weights used in the LF.
		"""
		stackMed = n.ones_like(n.empty(len(self.wave)))*self.dV
		stackMean = n.ones_like(n.empty(len(self.wave)))*self.dV
		#stackMeanWeighted = n.ones_like(n.empty(len(self.wave)))*self.dV
		stackVar = n.ones_like(n.empty(len(self.wave)))*self.dV
		stackN = n.ones_like(n.empty(len(self.wave)))*self.dV
		jackknifes = n.ones_like(n.empty((len(self.wave),10)))*self.dV
		for i in range(len(specMatrix.T)):
				pt=specMatrix.T[i]
				wt=specMatrixWeight.T[i]
				sel=(pt!=self.dV)
				# jackknife sub-sampling
				rd=n.random.random(len(pt))
				aim=n.arange(0,1.01,0.1)
				jks=n.array([ (rd>aim[jj])&(rd<aim[jj+1]) for jj in range(len(aim)-1) ])
				if len(pt[sel])>1:
						stackMed[i] = n.median(pt[sel])
						stackMean[i] = n.mean(pt[sel])
						#stackMeanWeighted[i] = n.average(pt[sel],weights=wt[sel])
						stackN[i] = len(pt[sel])
						inter = n.array([ n.median( pt[sel & (seK==False)] ) for seK in jks ])
						jackknifes[i] = inter
						stackVar[i] = n.std(inter)

		wavelength = fits.Column(name="wavelength",format="D", unit="Angstrom", array= self.wave)
		medianStack=fits.Column(name="medianStack",format="D", unit="erg/s/cm2/Angstrom", array= n.array(stackMed))
		meanStack=fits.Column(name="meanStack",format="D", unit="erg/s/cm2/Angstrom", array= n.array(stackMean))
		#meanWeightedStack=fits.Column(name="meanWeightedStack",format="D", unit= "erg/s/cm2/Angstrom", array= n.array(stackMeanWeighted))
		jackknifeSpectra=fits.Column(name="jackknifeSpectra",format="10D", unit="erg/s/cm2/Angstrom", array= n.array(jackknifes))
		jackknifStackErrors=fits.Column(name="jackknifStackErrors",format="D", unit="erg/s/cm2/Angstrom", array= n.array(stackVar))
		NspectraPerPixel=fits.Column(name="NspectraPerPixel",format="D", unit="", array= n.array(stackN))
		return  wavelength, medianStack, meanStack, jackknifStackErrors, jackknifeSpectra, NspectraPerPixel
		#return  wavelength, medianStack, meanStack, meanWeightedStack, jackknifStackErrors, jackknifeSpectra, NspectraPerPixel

	def convertSpectrum(self,redshift):
		"""
		Shifts the spectrum in the rest-frame and creates a spectrum with the sampling desired.
		Uses the spectres package from A.C. Carnall
		:param redshift: redshift of the spectrum
		return the new flux and erro flux array
		"""
		nwave=self.wavelength/(1+redshift)

		#inL=(self.wave>nwave.min())&(self.wave<nwave.max())
		#outL=(inL==False)

		#points=interp1d(nwave,nwave * self.fluxl)
		#pts=points(self.wave[inL]) / self.wave[inL]
		#res=n.ones_like(self.wave)*self.dV
		#res[inL]=pts

		#pointsErr=interp1d(nwave,nwave * self.fluxlErr)
		#ptsErr=pointsErr(self.wave[inL]) / self.wave[inL]
		#resErr=n.ones_like(self.wave)*self.dV
		#resErr[inL]=ptsErr

		#return res, resErr
		wavelength_spectrum = n.hstack((
			self.wave[0]-10,
			self.wave[0]-5,
			n.min(nwave)-10,
			n.min(nwave)-5,
			nwave,
			n.max(nwave)+5,
			n.max(nwave)+10,
			self.wave[-1]+5,
			self.wave[-1]+10
			))
		#
		flux_spectrum = n.hstack((
			self.dV,self.dV,self.dV,self.fluxl[0],
			self.fluxl,
			self.fluxl[-1],self.dV,self.dV,self.dV
			))
		#
		flux_error_spectrum = n.hstack((
			self.dV,self.dV,self.dV,self.dV,
			self.fluxlErr,
			self.dV,self.dV,self.dV,self.dV
			))
		#
		final_spectrum, final_spectrum_err = sp.spectres(
			self.wave,
			wavelength_spectrum,
			flux_spectrum,
			flux_error_spectrum )
		return final_spectrum, final_spectrum_err

	def getSpectra(self, path_to_spectrum):
		hdulist = fits.open(path_to_spectrum)
		wave = 10**hdulist[1].data['loglam']
		flux = hdulist[1].data['flux']
		ivar = hdulist[1].data['ivar']
		ratio = n.min(abs(10000.*n.log10(n.outer(wave, 1./maskLambda))), axis=1)
		margin = 1.5
		veto_sky = ratio <= margin
		selection = (veto_sky) & (ivar<=0) & (flux<0.)& (n.isinf(ivar)) & (n.isinf(flux))
		flux[selection] = n.zeros_like(ivar[selection])
		ivar[selection] = n.zeros_like(ivar[selection])
		out_sel = (flux>0)&(ivar>0)
		self.fluxl =flux[out_sel]
		self.fluxlErr=ivar[out_sel]**(-0.5)
		self.wavelength = wave[out_sel] #/(1+z)

	def fit_UV_continuum(self,x,y,yerr,degree=3):
		"""
		We then mask out
		absorption and emission features and fit a cubic polyno-
		mial function through the rest of the spectrum.
		Using
		the best-fit polynomial function as an estimate of the un-
		derlying continuum, F lambda
		we normalize the observed spectrum to obtain the continuum-normalized spectrum
		"""
		self.bad_flags = n.ones(len(x))

		# masking sky contaminated pixels
		maskLambda = n.loadtxt(os.path.join(os.environ['GIT_SPM'],'data',"dr12-sky-mask.txt"), unpack=True)
		ratio = n.min(abs(10000.*n.log10(n.outer(x, 1./maskLambda))), axis=1)
		margin = 1.5
		veto_sky = ( ratio <= margin )

		# UV mask
		UV_mask = (x>2000)&(x<3600)

		# UV line mask
		ratio = n.min(abs(10000.*n.log10(n.outer(x, 1./line_list_abs))), axis=1)
		margin = 8
		veto_line_abs = ( ratio <= margin )

		ratio = n.min(abs(10000.*n.log10(n.outer(x, 1./line_list_em))), axis=1)
		margin = 8
		veto_line_em = ( ratio <= margin )

		# MASKING BAD DATA
		bad_data = n.isnan(y) | n.isinf(y) | (y <= 0.0) | n.isnan(yerr) | n.isinf(yerr)
		# creating new arrays
		x = x[(UV_mask)&(veto_sky==False)&(bad_data==False)&(veto_line_abs==False)&(veto_line_em==False)]
		y = y[(UV_mask)&(veto_sky==False)&(bad_data==False)&(veto_line_abs==False)&(veto_line_em==False)]
		yerr = yerr[(UV_mask)&(veto_sky==False)&(bad_data==False)&(veto_line_abs==False)&(veto_line_em==False)]

		out=n.polyfit(x, y, degree, w=1/yerr)
		return out

	def createStackMatrix_Weighted(self):
		"""
		Function that constructs the stack matrix UV normed
		"""
		# loop over the file with N sorted with luminosity
		specMatrix, specMatrixErr, specMatrixWeight=[],[],[]
		print('plate, mjd, fiber, z, weights',self.plates[:10], self.mjds[:10], self.fiberids[:10], self.redshifts[:10], self.weights[:10])
		for plate, mjd, fiber, redshift, weight in zip(self.plates, self.mjds, self.fiberids, self.redshifts, self.weights):
			try:
				#print(plate, mjd, fiber, redshift)
				if plate > 3006 :
					path_to_spectrum = get_path_to_spectrum_v5_13_0(plate, mjd, fiber)
				else:
					path_to_spectrum = get_path_to_spectrum_26(plate, mjd, fiber)

				if os.path.isfile(path_to_spectrum):
					self.getSpectra(path_to_spectrum)
					pts,ptsErr = self.convertSpectrum(redshift)
					specMatrix.append(pts)
					specMatrixErr.append(ptsErr)
					specMatrixWeight.append(n.ones_like(pts)*weight)
				else: # for ELG spectra in v5_10_7
					path_to_spectrum = get_path_to_spectrum_v5_13_0(plate, mjd, fiber)
					if os.path.isfile(path_to_spectrum):
						self.getSpectra(path_to_spectrum)
						pts,ptsErr = self.convertSpectrum(redshift)
						specMatrix.append(pts)
						specMatrixErr.append(ptsErr)
						specMatrixWeight.append(n.ones_like(pts)*weight)
			except(ValueError,FileNotFoundError):
				print('value / file not found error !',plate, mjd, fiber)

		specMatrixWeight=n.array(specMatrixWeight)
		specMatrix=n.array(specMatrix)
		specMatrixErr=n.array(specMatrixErr)
		n.savetxt(self.out_file+'.specMatrix.dat', specMatrix)
		n.savetxt(self.out_file+'.specMatrixErr.dat', specMatrixErr)
		n.savetxt(self.out_file+'.specMatrixWeight.dat', specMatrixWeight)

	def createStackMatrix(self):
		"""
		Function that constructs the stack matrix UV normed
		"""
		# loop over the file with N sorted with luminosity
		specMatrix, specMatrixErr, specMatrixWeight=[],[],[]
		print(self.plates, self.mjds, self.fiberids, self.redshifts)
		for plate, mjd, fiber, redshift in zip(self.plates, self.mjds, self.fiberids, self.redshifts):
			try:
				print(plate, mjd, fiber, redshift)
				if plate > 3006 :
					path_to_spectrum = get_path_to_spectrum_v5_13_0(plate, mjd, fiber)
				else:
					path_to_spectrum = get_path_to_spectrum_26(plate, mjd, fiber)

				if os.path.isfile(path_to_spectrum):
					self.getSpectra(path_to_spectrum)
					pts,ptsErr = self.convertSpectrum(redshift)
					specMatrix.append(pts)
					specMatrixErr.append(ptsErr)
					weight=1.
					specMatrixWeight.append(n.ones_like(pts)*weight)
				else: # for ELG spectra in v5_10_7
					path_to_spectrum = get_path_to_spectrum_v5_13_0(plate, mjd, fiber)
					if os.path.isfile(path_to_spectrum):
						self.getSpectra(path_to_spectrum)
						pts,ptsErr = self.convertSpectrum(redshift)
						specMatrix.append(pts)
						specMatrixErr.append(ptsErr)
						weight=1.
						specMatrixWeight.append(n.ones_like(pts)*weight)
			except(ValueError,FileNotFoundError):
				print('value / file not found error !',plate, mjd, fiber)

		specMatrixWeight=n.array(specMatrixWeight)
		specMatrix=n.array(specMatrix)
		specMatrixErr=n.array(specMatrixErr)
		n.savetxt(self.out_file+'.specMatrix.dat', specMatrix)
		n.savetxt(self.out_file+'.specMatrixErr.dat', specMatrixErr)
		n.savetxt(self.out_file+'.specMatrixWeight.dat', specMatrixWeight)

	def createStackMatrix_UVnormed(self):
		"""
		Function that constructs the stack matrix UV normed
		"""
		# loop over the file with N sorted with luminosity
		specMatrix, specMatrixErr, specMatrixWeight=[],[],[]

		for plate, mjd, fiber, redshift in zip(self.plates, self.mjds, self.fiberids, self.redshifts):
			try:
				#print(plate, mjd, fiber, redshift)
				if plate > 3006 :
					path_to_spectrum = get_path_to_spectrum_v5_13_0(plate, mjd, fiber)
				else:
					path_to_spectrum = get_path_to_spectrum_26(plate, mjd, fiber)
				if os.path.isfile(path_to_spectrum):
					self.getSpectra(path_to_spectrum)
					pts,ptsErr = self.convertSpectrum(redshift)
					pfit = self.fit_UV_continuum(self.wave, pts ,ptsErr)
					Fcont = n.polyval(pfit, self.wave)
					specMatrix.append(pts/Fcont)
					specMatrixErr.append(ptsErr/Fcont)
					specMatrix.append(pts)
					specMatrixErr.append(ptsErr)
					weight=1.
					specMatrixWeight.append(n.ones_like(pts)*weight)
				else: # get ELG spectra in v5_10_7
					path_to_spectrum = get_path_to_spectrum_v5_13_0(plate, mjd, fiber)
					if os.path.isfile(path_to_spectrum):
						self.getSpectra(path_to_spectrum)
						pts,ptsErr = self.convertSpectrum(redshift)
						pfit = self.fit_UV_continuum(self.wave, pts ,ptsErr)
						Fcont = n.polyval(pfit, self.wave)
						specMatrix.append(pts/Fcont)
						specMatrixErr.append(ptsErr/Fcont)
						specMatrix.append(pts)
						specMatrixErr.append(ptsErr)
						weight=1.
						specMatrixWeight.append(n.ones_like(pts)*weight)

			except(ValueError,TypeError,FileNotFoundError):
				print('value or type error !',plate, mjd, fiber)

		specMatrixWeight=n.array(specMatrixWeight)
		specMatrix=n.array(specMatrix)
		specMatrixErr=n.array(specMatrixErr)
		n.savetxt(self.out_file+'.specMatrix.dat', specMatrix)
		n.savetxt(self.out_file+'.specMatrixErr.dat', specMatrixErr)
		n.savetxt(self.out_file+'.specMatrixWeight.dat', specMatrixWeight)

	def stackSpectra(self):
		"""
		Stacks
		"""
		# loop over the file with N sorted with luminosity
		self.specMatrix = n.loadtxt(self.out_file+'.specMatrix.dat')
		#specMatrixErr = n.loadtxt(self.out_file+'.specMatrixErr.dat')
		self.specMatrixWeight = n.loadtxt(self.out_file+'.specMatrixWeight.dat')
		print( "now stacks" )
		#wavelength, medianStack, meanStack, meanWeightedStack, jackknifStackErrors, jackknifeSpectra, NspectraPerPixel = self.stack_function( specMatrix ,specMatrixWeight)
		#cols = fits.ColDefs([wavelength, medianStack, meanStack, meanWeightedStack, jackknifStackErrors, jackknifeSpectra, NspectraPerPixel])
		wavelength, medianStack, meanStack, jackknifStackErrors, jackknifeSpectra, NspectraPerPixel = self.stack_function( self.specMatrix ,self.specMatrixWeight)
		cols = fits.ColDefs([wavelength, medianStack, meanStack, jackknifStackErrors, jackknifeSpectra, NspectraPerPixel])
		tbhdu = fits.BinTableHDU.from_columns(cols)
		prihdr = fits.Header()
		prihdr['author'] = "JC"
		prihdr['survey'] = self.survey
		prihdr['in_file'] = os.path.basename(self.in_file)[:-4]
		prihdr['Nspec'] = len(self.plates)
		prihdu = fits.PrimaryHDU(header=prihdr)
		thdulist = fits.HDUList([prihdu, tbhdu])
		if os.path.isfile(self.out_file):
			os.remove(self.out_file)
		print( "stack written to", self.out_file )
		thdulist.writeto(self.out_file)


'''
Andrew Till
Summer 2014

Utility functions for materials.

References:
mcnp5: MCNP 5 Manual, Volume I, LA-UR-03-1987, Appendix G (MCNP Data Libraries, including S(alpha,beta))
mcnpdata: Listing of Available ACE Data Tables [for MCNP6], LA-UR-13-21822 (An updated version of the above)
njoy2012: The NJOY Nuclear Data Processing System, Version 2012, LA-UR-12-27079 (NJOY 2012 manual)
endf6: ENDF-6 Formats Manual, BNL-90365-2009 Rev.2 (Description of the ENDF-6 format)

Naming conventions:
'Thermal name' is specified in materials_materials.py and applies to a material (e.g., ZrH is 'zrh').
'Element thermal name' includes the (bound) thermal treatment of the material and the applicable
    element (e.g., H in Zr is 'hzrh'). Uses Hollerith strings specified in Table 21 of MATXSR chapter
    in NJOY manual (ref: njoy2012).
    Elsewhere, this is called Sab.
'Thermal XS name' includes the (bound) thermal treatment of the material, the applicable element,
    and the applicable type of thermal XS (e.g., the inelastic bound thermal XS for H in ZrH is 'hzrhinel')
'MCNP thermal name' is the string used by MCNP to refer to a bound thermal treatment of a material along
    with a consituent element (e.g., the bound thermal XS for H in ZrH is 'h/zr')
NB: 'free' is used for thermal name, element thermal name, and thermal xs name for free-gas thermal
    treatment
NB: ''/None/'none' is used for thermal name; 'none' is used for element thermal name; [] is used
    for the thermal xs list (for no thermal treatment)
'''

#TPL
import numpy as np

###############################################################################
def get_nearest_points(value, sortedArray):
    '''If value in sortedArray, return [value], else return two nearest points in sortedArray'''
    nearestIndex = np.argmin(np.abs(sortedArray - value))
    nearestPoints = [sortedArray[nearestIndex]]
    lastIndex = len(sortedArray) - 1
    if value < sortedArray[nearestIndex] and nearestIndex != 0:
        nearestPoints.append(sortedArray[nearestIndex-1])
    elif value > sortedArray[nearestIndex] and nearestIndex != lastIndex:
        nearestPoints.append(sortedArray[nearestIndex+1])
    return nearestPoints

def thin_list(inList, useLinSpacing=True):
    '''Thin a list by returning the value of the point which is nearest to its two neighbors.'''
    sortedList = np.array(sorted(inList))
    # Use [2:] and [:-2] to get sum of distances to neighbors. Add 1 for the indexing to be correct.
    if useLinSpacing:
        distance = sortedList[2:] - sortedList[:-2]
    else:
        distance = sortedList[2:] / sortedList[:-2]
    index = np.argmin(distance) + 1
    return sortedList[index]

###############################################################################
def calc_chord_length(fuelRadius):
    if fuelRadius == 'unshielded':
        return 1.e10
    elif fuelRadius:
        surfaceArea = 2 * np.pi * fuelRadius
        volume = np.pi * fuelRadius * fuelRadius
        chordLength = surfaceArea / (4 * volume)
        return chordLength
    else:
        return 0.0

def has_bondarenko_iteration(Z):
    if Z <= 5:
        return False
    else:
        return True

def is_fissionable((Z,A)):
    if Z >= 89:
        return True
    elif (Z,A) in [(88, 223), (88, 226)]:
        return True
    else:
        return False

def avogadros_number():
    '''In units of atoms / mole, but multiplied by 1E-24'''
    return 0.60221413

def get_inelastic_thermal_mt_list():
    '''Ref: Culled from Table 4 in the NJOY manual (ref: njoy2012)'''
    return set([221, 222, 223, 225, 227, 228, 229, 231, 233, 235, 237, 239, 241, 243, 245])

###############################################################################
def format_zaid():
    '''Use format_ZAID()(**dict) to apply to a dictionary'''
    return '{Z:d}{A:03d}'.format

def format_zaid_leading_zeros():
    return '{Z:02d}{A:03d}'.format

def format_thermal_filename():
    return 'endf_th_{Z:02d}_{thermalName}_vii1'.format

def get_nuclide_dirr(sym, A, elementThermalName, metastableStr=''):
    '''Returns the directory name for a nuclide'''
    nuclideName = '{0}-{1}{2}'.format(sym.lower(), A, metastableStr)
    et2t = get_element_thermal_name_to_thermal_name_dict()
    if elementThermalName in et2t:
        thermalName = et2t[elementThermalName]
        nuclideName = '{0}-{1}'.format(nuclideName, thermalName)
    return nuclideName

def get_ace_extension(elementThermalName):
    '''Returns the ACE extension given an element thermal name. The extension is 9 (as in .90c)
    if the *nuclide* does not have a bound thermal treatment.'''
    if elementThermalName in get_non_bound_names():
        return 9
    else:
        et2t = get_element_thermal_name_to_thermal_name_dict()
        t2ext = get_thermal_name_to_ace_ext_dict()
        return t2ext[et2t[elementThermalName]]

###############################################################################
def get_thermal_name_to_element_thermal_name_dict():
    '''Returns a dict that maps tuple (Z, thermal_name) to element thermal name (a Hollerith string).
    Free thermal treatment is not in dict'''
    return {(1,'poly'): 'hpoly', (1,'h2o'): 'hh2o', (8,'uo2'): 'ouo2', (92,'uo2'): 'uuo2', (1,'zrh'): 'hzrh', (40,'zrh'): 'zrzrh', (6,'graphite'): 'graph', (13,'al'): 'al', (26,'fe'): 'fe'}

def get_element_thermal_name_to_nuclide_list_dict():
    '''Returns a dict that maps element thermal name to the (Z,A)'s of the nuclides to which it applies.
    Free thermal treatment is not in dict. Ref: mcnp5, mcnpdata'''
    return {'hpoly': [(1,1)], 'hh2o': [(1,1)], 'ouo2': [(8,16), (8,17), (8,18)], 'uuo2': [(92,238)], 'hzrh': [(1,1)], 'zrzrh': [(40,0), (40,90), (40,91), (40,92), (40,94), (40,96)], 'graph': [(6,0), (6,12)], 'al': [(13,27)], 'fe': [(26,56)]}

def get_element_thermal_name_to_inelastic_mt_number_dict():
    '''Returns a dict that maps element thermal name to the MT number corresponding to the
    inelastic thermal XS for that element thermal name. See Tables 4 and 25 in NJOY manual (ref: njoy2012).
    Should correspond to 'inel' endf numbers in Readgroupr.py's get_endf_mt_list() function.'''
    return {'free': 221, 'hh2o': 222, 'hpoly': 223, 'ouo2': 239, 'uuo2': 241, 'hzrh': 225, 'zrzrh': 235, 'graph': 229, 'al': 243, 'fe': 245}

def get_thermal_name_to_nuclide_list_dict():
    '''Returns a dict that maps thermal name to the (Z,A)'s of the nuclides to which it applies.
    Free thermal treatment is not in dict'''
    # Derive this dictionary from previous information
    Zthermal2elem = get_thermal_name_to_element_thermal_name_dict()
    elem2ZAs = get_element_thermal_name_to_nuclide_list_dict()
    # First, initialize each thermal name as an empty set of nuclides
    thermal2ZAs = {}
    for (Z,thermalName) in Zthermal2elem:
        thermal2ZAs[thermalName] = set()
    # Then, populate nuclides that correspond to that thermal name
    for (Z,thermalName) in Zthermal2elem:
        elem = Zthermal2elem[(Z,thermalName)]
        ZAList = elem2ZAs[elem]
        thermal2ZAs[thermalName].update(ZAList)
    return thermal2ZAs

def get_element_thermal_name_to_thermal_name_dict():
    '''Returns a dict that maps element thermal name (e.g., hh2o) to thermal name (h2o).
    Free thermal treatment is not in dict.'''
    # Derive this dictionary from previous information
    return {et: t for (Z,t), et in get_thermal_name_to_element_thermal_name_dict().items()}

###############################################################################
def get_non_bound_names():
    '''Returns a set of thermal names that do not use bound cross sections'''
    return set(['free', '', None, 'none'])

def get_element_thermal_name_to_thermal_xs_list_dict():
    '''Returns a dict that maps element thermal name to a list of thermal XS names.
    See Tables 4 and 25 in NJOY manual (ref: njoy2012)'''
    return {'free': ['free'], 'none': [], 'hpoly': ['hpolyinel', 'hpolyelas'], 'hh2o': ['hh2o'], 'ouo2': ['ouo2inel', 'ouo2elas'], 'uuo2': ['uuo2inel', 'uuo2elas'], 'hzrh': ['hzrhinel', 'hzrhelas'], 'zrzrh': ['zrzrhinel', 'zrzrhelas'], 'graph': ['graphinel', 'graphelas'], 'al': ['alinel', 'alelas'], 'fe': ['feinel', 'feelas'],}

def get_element_thermal_name_to_mat_number_dict():
    '''Returns a dict that maps element thermal name to thermal MAT number.
    See Table 4 in THERMR chapter of NJOY manual (ref: njoy2012).
    Change: Manuals say Al is material 45, but the data file has material 53'''
    return {'free': 0, 'none': 0, 'hpoly': 37, 'hh2o': 1, 'ouo2': 75, 'uuo2': 48, 'hzrh': 7, 'zrzrh': 58, 'graph': 31, 'al': 53, 'fe': 56}

def get_element_thermal_name_to_bnl_id_dict():
    '''Returns a dict that maps element thermal name to BNL's ID. Used in URL for automatic downloads.
    Free thermal treatment not in dict.
    Can be check at http://www.nndc.bnl.gov/sigma/tree/index.html. Last verified: 05/31/2016.'''
    return {'hpoly': 15390, 'hh2o': 15391, 'ouo2': 15395, 'uuo2': 15402, 'hzrh': 15392, 'zrzrh': 15403, 'graph': 15389, 'al': 15383, 'fe': 15384}

def get_element_thermal_name_to_endf_filename_dict():
    '''Returns a dict that maps element thermal name to bound thermal ENDF file name'''
    # Derive this dictionary from previous information
    ZName2elem = get_thermal_name_to_element_thermal_name_dict()
    ZName2filename = format_thermal_filename()
    elem2filename = {}
    for (Z,thermalName) in ZName2elem:
        elem = ZName2elem[(Z,thermalName)]
        elem2filename[elem] = ZName2filename(Z=Z, thermalName=thermalName)
    for key in ['free', None]:
        elem2filename[key] = None
    return elem2filename

###############################################################################
def get_thermal_name_to_ace_ext_dict():
    '''Returns a dict that maps thermal name to ACE extension (e.g., the 9 in 1001.92c).
    Reasonable extension options include poly (CH2) being 6, benzene (C6H6) being 7, and d2o (D2O) being 8, though these are currently unsupported'''
    return {'h2o': 0, 'uo2': 1, 'zrh': 2, 'graphite': 3, 'al': 4, 'fe': 5, 'free': 9}

def get_element_thermal_name_to_mcnp_thermal_name_dict():
    '''Returns a dict that maps element thermal name to MCNP thermal material name. Ref: mcnp5, mcnpdata'''
    return {'hpoly': 'poly', 'hh2o': 'lwtr', 'zrzrh': 'zr/h', 'hzrh': 'h/zr', 'uuo2': 'u/o2', 'ouo2': 'o2/u', 'graph': 'grph', 'fe': 'fe56', 'al': 'al27'}

def get_mcnp_thermal_name_to_main_za_dict():
    '''Returns a dict that maps MCNP bound thermal material name to the (Z,A) of the nuclide from which to read PENDF tape (for ACE only). Uses nuclide with highest atom fraction.'''
    return {'poly': (1,1), 'lwtr': (1,1), 'o2/u': (8,16), 'u/o2': (92,238), 'grph': (6,0), 'h/zr': (1,1), 'zr/h': (40,90), 'al27': (13,27), 'fe56': (26,56), }

def get_mcnp_thermal_name_to_zaid_list_dict():
    '''Returns a dict that maps MCNP bound thermal material name to a ZZAAA string of the nuclides that use this thermal option. Ref: mcnp5, mcnpdata'''
    # Derive this dictionary from previous information
    elem2mcnp = get_element_thermal_name_to_mcnp_thermal_name_dict()
    elem2ZAs = get_element_thermal_name_to_nuclide_list_dict()
    ZA2zaid = format_zaid()
    mcnp2zaid = {}
    for elem in elem2mcnp:
        mcnpName = elem2mcnp[elem]
        ZAList = elem2ZAs[elem]
        zaidList = [ZA2zaid(Z=Z, A=A) for (Z,A) in ZAList]
        mcnp2zaid[mcnpName] = zaidList
    return mcnp2zaid

###############################################################################
def print_newline(verbosity=False):
    if verbosity:
        print ''