# -*- coding: utf-8 -*-
"""
Multi-lib backend for POT

The goal is to write backend-agnostic code. Whether you're using Numpy, PyTorch,
or Jax, POT code should work nonetheless.
To achieve that, POT provides backend classes which implements functions in their respective backend
imitating Numpy API. As a convention, we use nx instead of np to refer to the backend.

Examples
--------

>>> from ot.utils import list_to_array
>>> from ot.backend import get_backend
>>> def f(a, b):  # the function does not know which backend to use
...     a, b = list_to_array(a, b)  # if a list in given, make it an array
...     nx = get_backend(a, b)  # infer the backend from the arguments
...     c = nx.dot(a, b)  # now use the backend to do any calculation
...     return c
"""

# Author: Remi Flamary <remi.flamary@polytechnique.edu>
#         Nicolas Courty <ncourty@irisa.fr>
#
# License: MIT License

import numpy as np
import scipy.special as scipy
from scipy.sparse import issparse, coo_matrix, csr_matrix

try:
    import torch
    torch_type = torch.Tensor
except ImportError:
    torch = False
    torch_type = float

try:
    import jax
    import jax.numpy as jnp
    import jax.scipy.special as jscipy
    jax_type = jax.numpy.ndarray
except ImportError:
    jax = False
    jax_type = float

str_type_error = "All array should be from the same type/backend. Current types are : {}"

def get_backend_list():
    """Returns the list of available backends"""
    lst = [NumpyBackend(), ]

    if torch:
        lst.append(TorchBackend())

    if jax:
        lst.append(JaxBackend())

    return lst

def get_backend(*args):
    """Returns the proper backend for a list of input arrays

        Also raises TypeError if all arrays are not from the same backend
    """
    # check that some arrays given
    if not len(args) > 0:
        raise ValueError(" The function takes at least one parameter")
    # check all same type
    if not len(set(type(a) for a in args)) == 1:
        raise ValueError(str_type_error.format([type(a) for a in args]))

    if isinstance(args[0], np.ndarray):
        return NumpyBackend()
    elif isinstance(args[0], torch_type):
        return TorchBackend()
    elif isinstance(args[0], jax_type):
        return JaxBackend()
    else:
        raise ValueError("Unknown type of non implemented backend.")

def to_numpy(*args):
    """Returns numpy arrays from any compatible backend"""

    if len(args) == 1:
        return get_backend(args[0]).to_numpy(args[0])
    else:
        return [get_backend(a).to_numpy(a) for a in args]

class Backend():
    """
    Backend abstract class.
    Implementations: :py:class:`JaxBackend`, :py:class:`NumpyBackend`, :py:class:`TorchBackend`

    - The `__name__` class attribute refers to the name of the backend.
    - The `__type__` class attribute refers to the data structure used by the backend.
    """

    __name__ = None
    __type__ = None
    __type_list__ = None

    rng_ = None

    def __str__(self):
        return self.__name__

    # convert to numpy
    def to_numpy(self, a):
        """Returns the numpy version of a tensor"""
        raise NotImplementedError()

    # convert from numpy
    def from_numpy(self, a, type_as=None):
        """Creates a tensor cloning a numpy array, with the given precision (defaulting to input's precision) and the given device (in case of GPUs)"""
        raise NotImplementedError()

    def set_gradients(self, val, inputs, grads):
        """Define the gradients for the value val wrt the inputs """
        raise NotImplementedError()

    def zeros(self, shape, type_as=None):
        r"""
        Creates a tensor full of zeros.

        This function follows the api from :any:`numpy.zeros`

        See: https://numpy.org/doc/stable/reference/generated/numpy.zeros.html
        """
        raise NotImplementedError()

    def ones(self, shape, type_as=None):
        r"""
        Creates a tensor full of ones.

        This function follows the api from :any:`numpy.ones`

        See: https://numpy.org/doc/stable/reference/generated/numpy.ones.html
        """
        raise NotImplementedError()

    def arange(self, stop, start=0, step=1, type_as=None):
        r"""
        Returns evenly spaced values within a given interval.

        This function follows the api from :any:`numpy.arange`

        See: https://numpy.org/doc/stable/reference/generated/numpy.arange.html
        """
        raise NotImplementedError()

    def full(self, shape, fill_value, type_as=None):
        r"""
        Creates a tensor with given shape, filled with given value.

        This function follows the api from :any:`numpy.full`

        See: https://numpy.org/doc/stable/reference/generated/numpy.full.html
        """
        raise NotImplementedError()

    def eye(self, N, M=None, type_as=None):
        r"""
        Creates the identity matrix of given size.

        This function follows the api from :any:`numpy.eye`

        See: https://numpy.org/doc/stable/reference/generated/numpy.eye.html
        """
        raise NotImplementedError()

    def sum(self, a, axis=None, keepdims=False):
        r"""
        Sums tensor elements over given dimensions.

        This function follows the api from :any:`numpy.sum`

        See: https://numpy.org/doc/stable/reference/generated/numpy.sum.html
        """
        raise NotImplementedError()

    def cumsum(self, a, axis=None):
        r"""
        Returns the cumulative sum of tensor elements over given dimensions.

        This function follows the api from :any:`numpy.cumsum`

        See: https://numpy.org/doc/stable/reference/generated/numpy.cumsum.html
        """
        raise NotImplementedError()

    def max(self, a, axis=None, keepdims=False):
        r"""
        Returns the maximum of an array or maximum along given dimensions.

        This function follows the api from :any:`numpy.amax`

        See: https://numpy.org/doc/stable/reference/generated/numpy.amax.html
        """
        raise NotImplementedError()

    def min(self, a, axis=None, keepdims=False):
        r"""
        Returns the maximum of an array or maximum along given dimensions.

        This function follows the api from :any:`numpy.amin`

        See: https://numpy.org/doc/stable/reference/generated/numpy.amin.html
        """
        raise NotImplementedError()

    def maximum(self, a, b):
        r"""
        Returns element-wise maximum of array elements.

        This function follows the api from :any:`numpy.maximum`

        See: https://numpy.org/doc/stable/reference/generated/numpy.maximum.html
        """
        raise NotImplementedError()

    def minimum(self, a, b):
        r"""
        Returns element-wise minimum of array elements.

        This function follows the api from :any:`numpy.minimum`

        See: https://numpy.org/doc/stable/reference/generated/numpy.minimum.html
        """
        raise NotImplementedError()

    def dot(self, a, b):
        r"""
        Returns the dot product of two tensors.

        This function follows the api from :any:`numpy.dot`

        See: https://numpy.org/doc/stable/reference/generated/numpy.dot.html
        """
        raise NotImplementedError()

    def abs(self, a):
        r"""
        Computes the absolute value element-wise.

        This function follows the api from :any:`numpy.absolute`

        See: https://numpy.org/doc/stable/reference/generated/numpy.absolute.html
        """
        raise NotImplementedError()

    def exp(self, a):
        r"""
        Computes the exponential value element-wise.

        This function follows the api from :any:`numpy.exp`

        See: https://numpy.org/doc/stable/reference/generated/numpy.exp.html
        """
        raise NotImplementedError()

    def log(self, a):
        r"""
        Computes the natural logarithm, element-wise.

        This function follows the api from :any:`numpy.log`

        See: https://numpy.org/doc/stable/reference/generated/numpy.log.html
        """
        raise NotImplementedError()

    def sqrt(self, a):
        r"""
        Returns the non-ngeative square root of a tensor, element-wise.

        This function follows the api from :any:`numpy.sqrt`

        See: https://numpy.org/doc/stable/reference/generated/numpy.sqrt.html
        """
        raise NotImplementedError()

    def power(self, a, exponents):
        r"""
        First tensor elements raised to powers from second tensor, element-wise.

        This function follows the api from :any:`numpy.power`

        See: https://numpy.org/doc/stable/reference/generated/numpy.power.html
        """
        raise NotImplementedError()

    def norm(self, a):
        r"""
        Computes the matrix frobenius norm.

        This function follows the api from :any:`numpy.linalg.norm`

        See: https://numpy.org/doc/stable/reference/generated/numpy.linalg.norm.html
        """
        raise NotImplementedError()

    def any(self, a):
        r"""
        Tests whether any tensor element along given dimensions evaluates to True.

        This function follows the api from :any:`numpy.any`

        See: https://numpy.org/doc/stable/reference/generated/numpy.any.html
        """
        raise NotImplementedError()

    def isnan(self, a):
        r"""
        Tests element-wise for NaN and returns result as a boolean tensor.

        This function follows the api from :any:`numpy.isnan`

        See: https://numpy.org/doc/stable/reference/generated/numpy.isnan.html
        """
        raise NotImplementedError()

    def isinf(self, a):
        r"""
        Tests element-wise for positive or negative infinity and returns result as a boolean tensor.

        This function follows the api from :any:`numpy.isinf`

        See: https://numpy.org/doc/stable/reference/generated/numpy.isinf.html
        """
        raise NotImplementedError()

    def einsum(self, subscripts, *operands):
        r"""
        Evaluates the Einstein summation convention on the operands.

        This function follows the api from :any:`numpy.einsum`

        See: https://numpy.org/doc/stable/reference/generated/numpy.einsum.html
        """
        raise NotImplementedError()

    def sort(self, a, axis=-1):
        r"""
        Returns a sorted copy of a tensor.

        This function follows the api from :any:`numpy.sort`

        See: https://numpy.org/doc/stable/reference/generated/numpy.sort.html
        """
        raise NotImplementedError()

    def argsort(self, a, axis=None):
        r"""
        Returns the indices that would sort a tensor.

        This function follows the api from :any:`numpy.argsort`

        See: https://numpy.org/doc/stable/reference/generated/numpy.argsort.html
        """
        raise NotImplementedError()

    def searchsorted(self, a, v, side='left'):
        r"""
        Finds indices where elements should be inserted to maintain order in given tensor.

        This function follows the api from :any:`numpy.searchsorted`

        See: https://numpy.org/doc/stable/reference/generated/numpy.searchsorted.html
        """
        raise NotImplementedError()

    def flip(self, a, axis=None):
        r"""
        Reverses the order of elements in a tensor along given dimensions.

        This function follows the api from :any:`numpy.flip`

        See: https://numpy.org/doc/stable/reference/generated/numpy.flip.html
        """
        raise NotImplementedError()

    def clip(self, a, a_min, a_max):
        """
        Limits the values in a tensor.

        This function follows the api from :any:`numpy.clip`

        See: https://numpy.org/doc/stable/reference/generated/numpy.clip.html
        """
        raise NotImplementedError()

    def repeat(self, a, repeats, axis=None):
        r"""
        Repeats elements of a tensor.

        This function follows the api from :any:`numpy.repeat`

        See: https://numpy.org/doc/stable/reference/generated/numpy.repeat.html
        """
        raise NotImplementedError()

    def take_along_axis(self, arr, indices, axis):
        r"""
        Gathers elements of a tensor along given dimensions.

        This function follows the api from :any:`numpy.take_along_axis`

        See: https://numpy.org/doc/stable/reference/generated/numpy.take_along_axis.html
        """
        raise NotImplementedError()

    def concatenate(self, arrays, axis=0):
        r"""
        Joins a sequence of tensors along an existing dimension.

        This function follows the api from :any:`numpy.concatenate`

        See: https://numpy.org/doc/stable/reference/generated/numpy.concatenate.html
        """
        raise NotImplementedError()

    def zero_pad(self, a, pad_width):
        r"""
        Pads a tensor.

        This function follows the api from :any:`numpy.pad`

        See: https://numpy.org/doc/stable/reference/generated/numpy.pad.html
        """
        raise NotImplementedError()

    def argmax(self, a, axis=None):
        r"""
        Returns the indices of the maximum values of a tensor along given dimensions.

        This function follows the api from :any:`numpy.argmax`

        See: https://numpy.org/doc/stable/reference/generated/numpy.argmax.html
        """
        raise NotImplementedError()

    def mean(self, a, axis=None):
        r"""
        Computes the arithmetic mean of a tensor along given dimensions.

        This function follows the api from :any:`numpy.mean`

        See: https://numpy.org/doc/stable/reference/generated/numpy.mean.html
        """
        raise NotImplementedError()

    def std(self, a, axis=None):
        r"""
        Computes the standard deviation of a tensor along given dimensions.

        This function follows the api from :any:`numpy.std`

        See: https://numpy.org/doc/stable/reference/generated/numpy.std.html
        """
        raise NotImplementedError()

    def linspace(self, start, stop, num):
        r"""
        Returns a specified number of evenly spaced values over a given interval.

        This function follows the api from :any:`numpy.linspace`

        See: https://numpy.org/doc/stable/reference/generated/numpy.linspace.html
        """
        raise NotImplementedError()

    def meshgrid(self, a, b):
        r"""
        Returns coordinate matrices from coordinate vectors (Numpy convention).

        This function follows the api from :any:`numpy.meshgrid`

        See: https://numpy.org/doc/stable/reference/generated/numpy.meshgrid.html
        """
        raise NotImplementedError()

    def diag(self, a, k=0):
        r"""
        Extracts or constructs a diagonal tensor.

        This function follows the api from :any:`numpy.diag`

        See: https://numpy.org/doc/stable/reference/generated/numpy.diag.html
        """
        raise NotImplementedError()

    def unique(self, a):
        r"""
        Finds unique elements of given tensor.

        This function follows the api from :any:`numpy.unique`

        See: https://numpy.org/doc/stable/reference/generated/numpy.unique.html
        """
        raise NotImplementedError()

    def logsumexp(self, a, axis=None):
        r"""
        Computes the log of the sum of exponentials of input elements.

        This function follows the api from :any:`scipy.special.logsumexp`

        See: https://docs.scipy.org/doc/scipy/reference/generated/scipy.special.logsumexp.html
        """
        raise NotImplementedError()

    def stack(self, arrays, axis=0):
        r"""
        Joins a sequence of tensors along a new dimension.

        This function follows the api from :any:`numpy.stack`

        See: https://numpy.org/doc/stable/reference/generated/numpy.stack.html
        """
        raise NotImplementedError()

    def outer(self, a, b):
        r"""
        Computes the outer product between two vectors.

        This function follows the api from :any:`numpy.outer`

        See: https://numpy.org/doc/stable/reference/generated/numpy.outer.html
        """
        raise NotImplementedError()

    def reshape(self, a, shape):
        r"""
        Gives a new shape to a tensor without changing its data.

        This function follows the api from :any:`numpy.reshape`

        See: https://numpy.org/doc/stable/reference/generated/numpy.reshape.html
        """
        raise NotImplementedError()

    def seed(self, seed=None):
        r"""
        Sets the seed for the random generator.

        This function follows the api from :any:`numpy.random.seed`

        See: https://numpy.org/doc/stable/reference/generated/numpy.random.seed.html
        """
        raise NotImplementedError()

    def rand(self, *size, type_as=None):
        r"""
        Generate uniform random numbers.

        This function follows the api from :any:`numpy.random.rand`

        See: https://numpy.org/doc/stable/reference/generated/numpy.random.rand.html
        """
        raise NotImplementedError()

    def randn(self, *size, type_as=None):
        r"""
        Generate normal Gaussian random numbers.

        This function follows the api from :any:`numpy.random.rand`

        See: https://numpy.org/doc/stable/reference/generated/numpy.random.rand.html
        """
        raise NotImplementedError()

    def coo_matrix(self, data, rows, cols, shape=None, type_as=None):
        r"""
        Creates a sparse tensor in COOrdinate format.

        This function follows the api from :any:`scipy.sparse.coo_matrix`

        See: https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.coo_matrix.html
        """
        raise NotImplementedError()

    def issparse(self, a):
        r"""
        Checks whether or not the input tensor is a sparse tensor.

        This function follows the api from :any:`scipy.sparse.issparse`

        See: https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.issparse.html
        """
        raise NotImplementedError()

    def tocsr(self, a):
        r"""
        Converts this matrix to Compressed Sparse Row format.

        This function follows the api from :any:`scipy.sparse.coo_matrix.tocsr`

        See: https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.coo_matrix.tocsr.html
        """
        raise NotImplementedError()

    def eliminate_zeros(self, a, threshold=0.):
        r"""
        Removes entries smaller than the given threshold from the sparse tensor.

        This function follows the api from :any:`scipy.sparse.csr_matrix.eliminate_zeros`

        See: https://docs.scipy.org/doc/scipy-0.14.0/reference/generated/scipy.sparse.csr_matrix.eliminate_zeros.html
        """
        raise NotImplementedError()

    def todense(self, a):
        r"""
        Converts a sparse tensor to a dense tensor.

        This function follows the api from :any:`scipy.sparse.csr_matrix.toarray`

        See: https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.csr_matrix.toarray.html
        """
        raise NotImplementedError()

    def where(self, condition, x, y):
        r"""
        Returns elements chosen from x or y depending on condition.

        This function follows the api from :any:`numpy.where`

        See: https://numpy.org/doc/stable/reference/generated/numpy.where.html
        """
        raise NotImplementedError()

    def copy(self, a):
        r"""
        Returns a copy of the given tensor.

        This function follows the api from :any:`numpy.copy`

        See: https://numpy.org/doc/stable/reference/generated/numpy.copy.html
        """
        raise NotImplementedError()

    def allclose(self, a, b, rtol=1e-05, atol=1e-08, equal_nan=False):
        r"""
        Returns True if two arrays are element-wise equal within a tolerance.

        This function follows the api from :any:`numpy.allclose`

        See: https://numpy.org/doc/stable/reference/generated/numpy.allclose.html
        """
        raise NotImplementedError()

    def dtype_device(self, a):
        r"""
        Returns the dtype and the device of the given tensor.
        """
        raise NotImplementedError()

    def assert_same_dtype_device(self, a, b):
        r"""
        Checks whether or not the two given inputs have the same dtype as well as the same device
        """
        raise NotImplementedError()

class NumpyBackend(Backend):
    """
    NumPy implementation of the backend

    - `__name__` is "numpy"
    - `__type__` is np.ndarray
    """

    __name__ = 'numpy'
    __type__ = np.ndarray
    __type_list__ = [np.array(1, dtype=np.float32),
                     np.array(1, dtype=np.float64)]

    rng_ = np.random.RandomState()

    def to_numpy(self, a):
        return a

    def from_numpy(self, a, type_as=None):
        if type_as is None:
            return a
        elif isinstance(a, float):
            return a
        else:
            return a.astype(type_as.dtype)

    def set_gradients(self, val, inputs, grads):
        # No gradients for numpy
        return val

    def zeros(self, shape, type_as=None):
        if type_as is None:
            return np.zeros(shape)
        else:
            return np.zeros(shape, dtype=type_as.dtype)

    def ones(self, shape, type_as=None):
        if type_as is None:
            return np.ones(shape)
        else:
            return np.ones(shape, dtype=type_as.dtype)

    def arange(self, stop, start=0, step=1, type_as=None):
        return np.arange(start, stop, step)

    def full(self, shape, fill_value, type_as=None):
        if type_as is None:
            return np.full(shape, fill_value)
        else:
            return np.full(shape, fill_value, dtype=type_as.dtype)

    def eye(self, N, M=None, type_as=None):
        if type_as is None:
            return np.eye(N, M)
        else:
            return np.eye(N, M, dtype=type_as.dtype)

    def sum(self, a, axis=None, keepdims=False):
        return np.sum(a, axis, keepdims=keepdims)

    def cumsum(self, a, axis=None):
        return np.cumsum(a, axis)

    def max(self, a, axis=None, keepdims=False):
        return np.max(a, axis, keepdims=keepdims)

    def min(self, a, axis=None, keepdims=False):
        return np.min(a, axis, keepdims=keepdims)

    def maximum(self, a, b):
        return np.maximum(a, b)

    def minimum(self, a, b):
        return np.minimum(a, b)

    def dot(self, a, b):
        return np.dot(a, b)

    def abs(self, a):
        return np.abs(a)

    def exp(self, a):
        return np.exp(a)

    def log(self, a):
        return np.log(a)

    def sqrt(self, a):
        return np.sqrt(a)

    def power(self, a, exponents):
        return np.power(a, exponents)

    def norm(self, a):
        return np.sqrt(np.sum(np.square(a)))

    def any(self, a):
        return np.any(a)

    def isnan(self, a):
        return np.isnan(a)

    def isinf(self, a):
        return np.isinf(a)

    def einsum(self, subscripts, *operands):
        return np.einsum(subscripts, *operands)

    def sort(self, a, axis=-1):
        return np.sort(a, axis)

    def argsort(self, a, axis=-1):
        return np.argsort(a, axis)

    def searchsorted(self, a, v, side='left'):
        if a.ndim == 1:
            return np.searchsorted(a, v, side)
        else:
            # this is a not very efficient way to make numpy
            # searchsorted work on 2d arrays
            ret = np.empty(v.shape, dtype=int)
            for i in range(a.shape[0]):
                ret[i, :] = np.searchsorted(a[i, :], v[i, :], side)
            return ret

    def flip(self, a, axis=None):
        return np.flip(a, axis)

    def outer(self, a, b):
        return np.outer(a, b)

    def clip(self, a, a_min, a_max):
        return np.clip(a, a_min, a_max)

    def repeat(self, a, repeats, axis=None):
        return np.repeat(a, repeats, axis)

    def take_along_axis(self, arr, indices, axis):
        return np.take_along_axis(arr, indices, axis)

    def concatenate(self, arrays, axis=0):
        return np.concatenate(arrays, axis)

    def zero_pad(self, a, pad_width):
        return np.pad(a, pad_width)

    def argmax(self, a, axis=None):
        return np.argmax(a, axis=axis)

    def mean(self, a, axis=None):
        return np.mean(a, axis=axis)

    def std(self, a, axis=None):
        return np.std(a, axis=axis)

    def linspace(self, start, stop, num):
        return np.linspace(start, stop, num)

    def meshgrid(self, a, b):
        return np.meshgrid(a, b)

    def diag(self, a, k=0):
        return np.diag(a, k)

    def unique(self, a):
        return np.unique(a)

    def logsumexp(self, a, axis=None):
        return scipy.logsumexp(a, axis=axis)

    def stack(self, arrays, axis=0):
        return np.stack(arrays, axis)

    def reshape(self, a, shape):
        return np.reshape(a, shape)

    def seed(self, seed=None):
        if seed is not None:
            self.rng_.seed(seed)

    def rand(self, *size, type_as=None):
        return self.rng_.rand(*size)

    def randn(self, *size, type_as=None):
        return self.rng_.randn(*size)

    def coo_matrix(self, data, rows, cols, shape=None, type_as=None):
        if type_as is None:
            return coo_matrix((data, (rows, cols)), shape=shape)
        else:
            return coo_matrix((data, (rows, cols)), shape=shape, dtype=type_as.dtype)

    def issparse(self, a):
        return issparse(a)

    def tocsr(self, a):
        if self.issparse(a):
            return a.tocsr()
        else:
            return csr_matrix(a)

    def eliminate_zeros(self, a, threshold=0.):
        if threshold > 0:
            if self.issparse(a):
                a.data[self.abs(a.data) <= threshold] = 0
            else:
                a[self.abs(a) <= threshold] = 0
        if self.issparse(a):
            a.eliminate_zeros()
        return a

    def todense(self, a):
        if self.issparse(a):
            return a.toarray()
        else:
            return a

    def where(self, condition, x, y):
        return np.where(condition, x, y)

    def copy(self, a):
        return a.copy()

    def allclose(self, a, b, rtol=1e-05, atol=1e-08, equal_nan=False):
        return np.allclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan)

    def dtype_device(self, a):
        if hasattr(a, "dtype"):
            return a.dtype, "cpu"
        else:
            return type(a), "cpu"

    def assert_same_dtype_device(self, a, b):
        # numpy has implicit type conversion so we automatically validate the test
        pass

class JaxBackend(Backend):
    """
    JAX implementation of the backend

    - `__name__` is "jax"
    - `__type__` is jax.numpy.ndarray
    """

    __name__ = 'jax'
    __type__ = jax_type
    __type_list__ = None

    rng_ = None

    def __init__(self):
        self.rng_ = jax.random.PRNGKey(42)

        for d in jax.devices():
            self.__type_list__ = [jax.device_put(jnp.array(1, dtype=jnp.float32), d),
                                  jax.device_put(jnp.array(1, dtype=jnp.float64), d)]

    def to_numpy(self, a):
        return np.array(a)

    def _change_device(self, a, type_as):
        return jax.device_put(a, type_as.device_buffer.device())

    def from_numpy(self, a, type_as=None):
        if type_as is None:
            return jnp.array(a)
        else:
            return self._change_device(jnp.array(a).astype(type_as.dtype), type_as)

    def set_gradients(self, val, inputs, grads):
        from jax.flatten_util import ravel_pytree
        val, = jax.lax.stop_gradient((val,))

        ravelled_inputs, _ = ravel_pytree(inputs)
        ravelled_grads, _ = ravel_pytree(grads)

        aux = jnp.sum(ravelled_inputs * ravelled_grads) / 2
        aux = aux - jax.lax.stop_gradient(aux)

        val, = jax.tree_map(lambda z: z + aux, (val,))
        return val

    def zeros(self, shape, type_as=None):
        if type_as is None:
            return jnp.zeros(shape)
        else:
            return self._change_device(jnp.zeros(shape, dtype=type_as.dtype), type_as)

    def ones(self, shape, type_as=None):
        if type_as is None:
            return jnp.ones(shape)
        else:
            return self._change_device(jnp.ones(shape, dtype=type_as.dtype), type_as)

    def arange(self, stop, start=0, step=1, type_as=None):
        return jnp.arange(start, stop, step)

    def full(self, shape, fill_value, type_as=None):
        if type_as is None:
            return jnp.full(shape, fill_value)
        else:
            return self._change_device(jnp.full(shape, fill_value, dtype=type_as.dtype), type_as)

    def eye(self, N, M=None, type_as=None):
        if type_as is None:
            return jnp.eye(N, M)
        else:
            return self._change_device(jnp.eye(N, M, dtype=type_as.dtype), type_as)

    def sum(self, a, axis=None, keepdims=False):
        return jnp.sum(a, axis, keepdims=keepdims)

    def cumsum(self, a, axis=None):
        return jnp.cumsum(a, axis)

    def max(self, a, axis=None, keepdims=False):
        return jnp.max(a, axis, keepdims=keepdims)

    def min(self, a, axis=None, keepdims=False):
        return jnp.min(a, axis, keepdims=keepdims)

    def maximum(self, a, b):
        return jnp.maximum(a, b)

    def minimum(self, a, b):
        return jnp.minimum(a, b)

    def dot(self, a, b):
        return jnp.dot(a, b)

    def abs(self, a):
        return jnp.abs(a)

    def exp(self, a):
        return jnp.exp(a)

    def log(self, a):
        return jnp.log(a)

    def sqrt(self, a):
        return jnp.sqrt(a)

    def power(self, a, exponents):
        return jnp.power(a, exponents)

    def norm(self, a):
        return jnp.sqrt(jnp.sum(jnp.square(a)))

    def any(self, a):
        return jnp.any(a)

    def isnan(self, a):
        return jnp.isnan(a)

    def isinf(self, a):
        return jnp.isinf(a)

    def einsum(self, subscripts, *operands):
        return jnp.einsum(subscripts, *operands)

    def sort(self, a, axis=-1):
        return jnp.sort(a, axis)

    def argsort(self, a, axis=-1):
        return jnp.argsort(a, axis)

    def searchsorted(self, a, v, side='left'):
        if a.ndim == 1:
            return jnp.searchsorted(a, v, side)
        else:
            # this is a not very efficient way to make jax numpy
            # searchsorted work on 2d arrays
            return jnp.array([jnp.searchsorted(a[i, :], v[i, :], side) for i in range(a.shape[0])])

    def flip(self, a, axis=None):
        return jnp.flip(a, axis)

    def outer(self, a, b):
        return jnp.outer(a, b)

    def clip(self, a, a_min, a_max):
        return jnp.clip(a, a_min, a_max)

    def repeat(self, a, repeats, axis=None):
        return jnp.repeat(a, repeats, axis)

    def take_along_axis(self, arr, indices, axis):
        return jnp.take_along_axis(arr, indices, axis)

    def concatenate(self, arrays, axis=0):
        return jnp.concatenate(arrays, axis)

    def zero_pad(self, a, pad_width):
        return jnp.pad(a, pad_width)

    def argmax(self, a, axis=None):
        return jnp.argmax(a, axis=axis)

    def mean(self, a, axis=None):
        return jnp.mean(a, axis=axis)

    def std(self, a, axis=None):
        return jnp.std(a, axis=axis)

    def linspace(self, start, stop, num):
        return jnp.linspace(start, stop, num)

    def meshgrid(self, a, b):
        return jnp.meshgrid(a, b)

    def diag(self, a, k=0):
        return jnp.diag(a, k)

    def unique(self, a):
        return jnp.unique(a)

    def logsumexp(self, a, axis=None):
        return jscipy.logsumexp(a, axis=axis)

    def stack(self, arrays, axis=0):
        return jnp.stack(arrays, axis)

    def reshape(self, a, shape):
        return jnp.reshape(a, shape)

    def seed(self, seed=None):
        if seed is not None:
            self.rng_ = jax.random.PRNGKey(seed)

    def rand(self, *size, type_as=None):
        self.rng_, subkey = jax.random.split(self.rng_)
        if type_as is not None:
            return jax.random.uniform(subkey, shape=size, dtype=type_as.dtype)
        else:
            return jax.random.uniform(subkey, shape=size)

    def randn(self, *size, type_as=None):
        self.rng_, subkey = jax.random.split(self.rng_)
        if type_as is not None:
            return jax.random.normal(subkey, shape=size, dtype=type_as.dtype)
        else:
            return jax.random.normal(subkey, shape=size)

    def coo_matrix(self, data, rows, cols, shape=None, type_as=None):
        # Currently, JAX does not support sparse matrices
        data = self.to_numpy(data)
        rows = self.to_numpy(rows)
        cols = self.to_numpy(cols)
        nx = NumpyBackend()
        coo_matrix = nx.coo_matrix(data, rows, cols, shape=shape, type_as=type_as)
        matrix = nx.todense(coo_matrix)
        return self.from_numpy(matrix)

    def issparse(self, a):
        # Currently, JAX does not support sparse matrices
        return False

    def tocsr(self, a):
        # Currently, JAX does not support sparse matrices
        return a

    def eliminate_zeros(self, a, threshold=0.):
        # Currently, JAX does not support sparse matrices
        if threshold > 0:
            return self.where(
                self.abs(a) <= threshold,
                self.zeros((1,), type_as=a),
                a
            )
        return a

    def todense(self, a):
        # Currently, JAX does not support sparse matrices
        return a

    def where(self, condition, x, y):
        return jnp.where(condition, x, y)

    def copy(self, a):
        # No need to copy, JAX arrays are immutable
        return a

    def allclose(self, a, b, rtol=1e-05, atol=1e-08, equal_nan=False):
        return jnp.allclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan)

    def dtype_device(self, a):
        return a.dtype, a.device_buffer.device()

    def assert_same_dtype_device(self, a, b):
        a_dtype, a_device = self.dtype_device(a)
        b_dtype, b_device = self.dtype_device(b)

        assert a_dtype == b_dtype, "Dtype discrepancy"
        assert a_device == b_device, f"Device discrepancy. First input is on {str(a_device)}, whereas second input is on {str(b_device)}"

class TorchBackend(Backend):
    """
    PyTorch implementation of the backend

    - `__name__` is "torch"
    - `__type__` is torch.Tensor
    """

    __name__ = 'torch'
    __type__ = torch_type
    __type_list__ = None

    rng_ = None

    def __init__(self):

        self.rng_ = torch.Generator()
        self.rng_.seed()

        self.__type_list__ = [torch.tensor(1, dtype=torch.float32),
                              torch.tensor(1, dtype=torch.float64)]

        if torch.cuda.is_available():
            self.__type_list__.append(torch.tensor(1, dtype=torch.float32, device='cuda'))
            self.__type_list__.append(torch.tensor(1, dtype=torch.float64, device='cuda'))

        from torch.autograd import Function

        # define a function that takes inputs val and grads
        # ad returns a val tensor with proper gradients
        class ValFunction(Function):

            @staticmethod
            def forward(ctx, val, grads, *inputs):
                ctx.grads = grads
                return val

            @staticmethod
            def backward(ctx, grad_output):
                # the gradients are grad
                return (None, None) + tuple(g * grad_output for g in ctx.grads)

        self.ValFunction = ValFunction

    def to_numpy(self, a):
        return a.cpu().detach().numpy()

    def from_numpy(self, a, type_as=None):
        if isinstance(a, float):
            a = np.array(a)
        if type_as is None:
            return torch.from_numpy(a)
        else:
            return torch.as_tensor(a, dtype=type_as.dtype, device=type_as.device)

    def set_gradients(self, val, inputs, grads):

        Func = self.ValFunction()

        res = Func.apply(val, grads, *inputs)

        return res

    def zeros(self, shape, type_as=None):
        if isinstance(shape, int):
            shape = (shape,)
        if type_as is None:
            return torch.zeros(shape)
        else:
            return torch.zeros(shape, dtype=type_as.dtype, device=type_as.device)

    def ones(self, shape, type_as=None):
        if isinstance(shape, int):
            shape = (shape,)
        if type_as is None:
            return torch.ones(shape)
        else:
            return torch.ones(shape, dtype=type_as.dtype, device=type_as.device)

    def arange(self, stop, start=0, step=1, type_as=None):
        if type_as is None:
            return torch.arange(start, stop, step)
        else:
            return torch.arange(start, stop, step, device=type_as.device)

    def full(self, shape, fill_value, type_as=None):
        if isinstance(shape, int):
            shape = (shape,)
        if type_as is None:
            return torch.full(shape, fill_value)
        else:
            return torch.full(shape, fill_value, dtype=type_as.dtype, device=type_as.device)

    def eye(self, N, M=None, type_as=None):
        if M is None:
            M = N
        if type_as is None:
            return torch.eye(N, m=M)
        else:
            return torch.eye(N, m=M, dtype=type_as.dtype, device=type_as.device)

    def sum(self, a, axis=None, keepdims=False):
        if axis is None:
            return torch.sum(a)
        else:
            return torch.sum(a, axis, keepdim=keepdims)

    def cumsum(self, a, axis=None):
        if axis is None:
            return torch.cumsum(a.flatten(), 0)
        else:
            return torch.cumsum(a, axis)

    def max(self, a, axis=None, keepdims=False):
        if axis is None:
            return torch.max(a)
        else:
            return torch.max(a, axis, keepdim=keepdims)[0]

    def min(self, a, axis=None, keepdims=False):
        if axis is None:
            return torch.min(a)
        else:
            return torch.min(a, axis, keepdim=keepdims)[0]

    def maximum(self, a, b):
        if isinstance(a, int) or isinstance(a, float):
            a = torch.tensor([float(a)], dtype=b.dtype, device=b.device)
        if isinstance(b, int) or isinstance(b, float):
            b = torch.tensor([float(b)], dtype=a.dtype, device=a.device)
        if hasattr(torch, "maximum"):
            return torch.maximum(a, b)
        else:
            return torch.max(torch.stack(torch.broadcast_tensors(a, b)), axis=0)[0]

    def minimum(self, a, b):
        if isinstance(a, int) or isinstance(a, float):
            a = torch.tensor([float(a)], dtype=b.dtype, device=b.device)
        if isinstance(b, int) or isinstance(b, float):
            b = torch.tensor([float(b)], dtype=a.dtype, device=a.device)
        if hasattr(torch, "minimum"):
            return torch.minimum(a, b)
        else:
            return torch.min(torch.stack(torch.broadcast_tensors(a, b)), axis=0)[0]

    def dot(self, a, b):
        return torch.matmul(a, b)

    def abs(self, a):
        return torch.abs(a)

    def exp(self, a):
        return torch.exp(a)

    def log(self, a):
        return torch.log(a)

    def sqrt(self, a):
        return torch.sqrt(a)

    def power(self, a, exponents):
        return torch.pow(a, exponents)

    def norm(self, a):
        return torch.sqrt(torch.sum(torch.square(a)))

    def any(self, a):
        return torch.any(a)

    def isnan(self, a):
        return torch.isnan(a)

    def isinf(self, a):
        return torch.isinf(a)

    def einsum(self, subscripts, *operands):
        return torch.einsum(subscripts, *operands)

    def sort(self, a, axis=-1):
        sorted0, indices = torch.sort(a, dim=axis)
        return sorted0

    def argsort(self, a, axis=-1):
        sorted, indices = torch.sort(a, dim=axis)
        return indices

    def searchsorted(self, a, v, side='left'):
        right = (side != 'left')
        return torch.searchsorted(a, v, right=right)

    def flip(self, a, axis=None):
        if axis is None:
            return torch.flip(a, tuple(i for i in range(len(a.shape))))
        if isinstance(axis, int):
            return torch.flip(a, (axis,))
        else:
            return torch.flip(a, dims=axis)

    def outer(self, a, b):
        return torch.outer(a, b)

    def clip(self, a, a_min, a_max):
        return torch.clamp(a, a_min, a_max)

    def repeat(self, a, repeats, axis=None):
        return torch.repeat_interleave(a, repeats, dim=axis)

    def take_along_axis(self, arr, indices, axis):
        return torch.gather(arr, axis, indices)

    def concatenate(self, arrays, axis=0):
        return torch.cat(arrays, dim=axis)

    def zero_pad(self, a, pad_width):
        from torch.nn.functional import pad
        # pad_width is an array of ndim tuples indicating how many 0 before and after
        # we need to add. We first need to make it compliant with torch syntax, that
        # starts with the last dim, then second last, etc.
        how_pad = tuple(element for tupl in pad_width[::-1] for element in tupl)
        return pad(a, how_pad)

    def argmax(self, a, axis=None):
        return torch.argmax(a, dim=axis)

    def mean(self, a, axis=None):
        if axis is not None:
            return torch.mean(a, dim=axis)
        else:
            return torch.mean(a)

    def std(self, a, axis=None):
        if axis is not None:
            return torch.std(a, dim=axis, unbiased=False)
        else:
            return torch.std(a, unbiased=False)

    def linspace(self, start, stop, num):
        return torch.linspace(start, stop, num, dtype=torch.float64)

    def meshgrid(self, a, b):
        X, Y = torch.meshgrid(a, b)
        return X.T, Y.T

    def diag(self, a, k=0):
        return torch.diag(a, diagonal=k)

    def unique(self, a):
        return torch.unique(a)

    def logsumexp(self, a, axis=None):
        if axis is not None:
            return torch.logsumexp(a, dim=axis)
        else:
            return torch.logsumexp(a, dim=tuple(range(len(a.shape))))

    def stack(self, arrays, axis=0):
        return torch.stack(arrays, dim=axis)

    def reshape(self, a, shape):
        return torch.reshape(a, shape)

    def seed(self, seed=None):
        if isinstance(seed, int):
            self.rng_.manual_seed(seed)
        elif isinstance(seed, torch.Generator):
            self.rng_ = seed
        else:
            raise ValueError("Non compatible seed : {}".format(seed))

    def rand(self, *size, type_as=None):
        if type_as is not None:
            return torch.rand(size=size, generator=self.rng_, dtype=type_as.dtype, device=type_as.device)
        else:
            return torch.rand(size=size, generator=self.rng_)

    def randn(self, *size, type_as=None):
        if type_as is not None:
            return torch.randn(size=size, dtype=type_as.dtype, generator=self.rng_, device=type_as.device)
        else:
            return torch.randn(size=size, generator=self.rng_)

    def coo_matrix(self, data, rows, cols, shape=None, type_as=None):
        if type_as is None:
            return torch.sparse_coo_tensor(torch.stack([rows, cols]), data, size=shape)
        else:
            return torch.sparse_coo_tensor(
                torch.stack([rows, cols]), data, size=shape,
                dtype=type_as.dtype, device=type_as.device
            )

    def issparse(self, a):
        return getattr(a, "is_sparse", False) or getattr(a, "is_sparse_csr", False)

    def tocsr(self, a):
        # Versions older than 1.9 do not support CSR tensors. PyTorch 1.9 and 1.10 offer a very limited support
        return self.todense(a)

    def eliminate_zeros(self, a, threshold=0.):
        if self.issparse(a):
            if threshold > 0:
                mask = self.abs(a) <= threshold
                mask = ~mask
                mask = mask.nonzero()
            else:
                mask = a._values().nonzero()
            nv = a._values().index_select(0, mask.view(-1))
            ni = a._indices().index_select(1, mask.view(-1))
            return self.coo_matrix(nv, ni[0], ni[1], shape=a.shape, type_as=a)
        else:
            if threshold > 0:
                a[self.abs(a) <= threshold] = 0
            return a

    def todense(self, a):
        if self.issparse(a):
            return a.to_dense()
        else:
            return a

    def where(self, condition, x, y):
        return torch.where(condition, x, y)

    def copy(self, a):
        return torch.clone(a)

    def allclose(self, a, b, rtol=1e-05, atol=1e-08, equal_nan=False):
        return torch.allclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan)

    def dtype_device(self, a):
        return a.dtype, a.device

    def assert_same_dtype_device(self, a, b):
        a_dtype, a_device = self.dtype_device(a)
        b_dtype, b_device = self.dtype_device(b)

        assert a_dtype == b_dtype, "Dtype discrepancy"
        assert a_device == b_device, f"Device discrepancy. First input is on {str(a_device)}, whereas second input is on {str(b_device)}"

import numpy,from numpy,import scipy,from scipy,import jax,from jax


"""
Part of code is from Kai Li "kailigo". The gitub website is https://github.com/kailigo/cvcZSL.
We add virtual classes and IAS in the code.
"""
from scipy import io
import numpy as np
import torch
import torch.nn as nn
from torch.autograd import Variable
from torch.nn import functional as F
from torch.optim import lr_scheduler
import torch.utils.data as data
from sklearn.metrics import accuracy_score
from tensorboardX import SummaryWriter

from utils import ReDirectSTD
from unseen_data_loader import data_loader_virtualCls
from unseen_option import Options

import os
import random
import pickle
# from test_embeded import test_while_training_simple

TMP = 10

args = Options().parse()
model_file_name = './chk/' + args.model_file
summaryFolder = './summary/' + args.log_file
if not os.path.exists(summaryFolder):
    os.mkdir(summaryFolder)
writer = SummaryWriter(summaryFolder)
print(args)

def init_seeds(seed=0):
    random.seed(seed)
    np.random.seed(seed)

    # torch cuda
    torch.cuda.empty_cache()
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

init_seeds(520)
def calc_accuracy(test_visual, test_label, attM, test_id, test_id_seen_unseen,cossim=False):
        outpred = [0] * test_visual.shape[0]
        end = 0
        outpred_list = []
        zslpred_list = []
        score_n_list = []
        for j in range(0,len(test_visual),64):
            if j+64>len(test_visual):
                end = len(test_visual)
            else:
                end = j+64
            all_cls_weights = forward(attM,test_visual[j:end],tmp=TMP)
            score,score_n=apply_classification_weights(test_visual[j:end].cuda(),
                    all_cls_weights,norm=True)
            score = score.squeeze(0)
            score_n = score_n.squeeze(0)

            _, pred = score.max(dim=1)
            pred = pred.view(-1)
            select_test_label = test_label[j:end].view(-1)

            outpred_list.extend(test_id[pred.cpu().detach().numpy()])
            zslpred_list.extend([ i in test_id_seen_unseen for i in test_id[pred.cpu().detach().numpy()]])
            score_n_list.extend(score_n.cpu().detach().numpy())

        seen_unseen_acc = accuracy_score(np.ones(len(zslpred_list)),zslpred_list)

        outpred = np.array(outpred_list, dtype='int')
        score_n = np.array(score_n_list)
        test_label = test_label.numpy()
        unique_labels = np.unique(test_label)
        acc = 0
        acc_cls = {}
        preds_cls = {}
        cos_cls = {}
        for l in unique_labels:
                idx = np.nonzero(test_label == l)[0]
                acc_cls[l] = accuracy_score(test_label[idx],outpred[idx])
                try:
                    preds_cls[l].extend(list(outpred[idx]))
                except(KeyError):
                    preds_cls[l] = list(outpred[idx])

                acc += acc_cls[l]
                loc = test_label == l
                outpred_l = outpred[idx]
                score_n_l = score_n[idx]
                loc = outpred_l == l

                if cossim:
                    cos_c = score_n_l[loc,l]
                    try:
                        cos_cls[l].extend(cos_c)
                    except(KeyError):
                        cos_cls[l] = [cos_c]
        acc = acc / unique_labels.shape[0]

        return acc,seen_unseen_acc,acc_cls,cos_cls,preds_cls

def compute_accuracy_all(test_att, att_all, test_visual_unseen, test_id_unseen, test_label_unseen,
                test_visual_seen, test_id_all, test_label_seen,train_id):

        acc_zsl,_,unseenacc_cls,_,unseenpred_cls = calc_accuracy(test_visual_unseen, test_label_unseen, test_att, test_id_unseen,test_id_unseen)
        acc_seenAcc,_,seenacc_cls,_,_ = calc_accuracy(test_visual_seen, test_label_seen, att_all, train_id, train_id)

        att_all_cls = torch.cat((att_all, test_att))

        acc_gzsl_unseen,Ru,unseengeneralAcc_cls,unseenCos_cls,_ = calc_accuracy(test_visual_unseen, test_label_unseen, att_all_cls, test_id_all,test_id_unseen,cossim=True)

        acc_gzsl_seen,Rs,seengeneralAcc_cls,seenCos_cls,_ = calc_accuracy(test_visual_seen, test_label_seen, att_all_cls, test_id_all,train_id,cossim=True)
        acc_cls = {**unseenacc_cls,**seenacc_cls}
        generalAcc_cls = {**unseengeneralAcc_cls,**seengeneralAcc_cls}
        H = 2 * acc_gzsl_seen * acc_gzsl_unseen / (acc_gzsl_seen + acc_gzsl_unseen)

        return acc_zsl, acc_seenAcc, acc_gzsl_unseen, acc_gzsl_seen, H, Rs, Ru, acc_cls, generalAcc_cls,unseenCos_cls,seenCos_cls,unseenpred_cls

def apply_classification_weights(features, cls_weights,norm=False):

        features = F.normalize(features,dim=-1)
        cls_weights = F.normalize(cls_weights, p=2, dim=-1, eps=1e-12)

        cls_scores = scale_cls * (torch.matmul(cls_weights,features.t()))
        cls_scores = cls_scores.permute(0,2,1)
        cls_scores = torch.diagonal(cls_scores,offset=0,dim1=0,dim2=1)
        cls_scores = cls_scores.t()
        if norm:
            return cls_scores,cls_scores/scale_cls
        else:
            return cls_scores

def IASatt(features,AttM,tmp):
        cls_num = len(AttM)
        attdims = AttM.shape[1]
        atten = torch.mm(features,w_IAS)+b_IAS
        atten = F.softmax(atten/tmp,dim=1).reshape(-1,1,attdims)
        atten = atten.reshape(-1,1,attdims)
        atten = atten + torch.ones_like(atten)
        atten = atten.repeat(1,cls_num,1)
        AttM = AttM.unsqueeze(0)
        AttM = AttM.repeat(len(features),1,1)
        AttM = atten*AttM
        return AttM

def forward(att,features,tmp):
        features = features.squeeze()
        att = IASatt(features,att,tmp)

        a1 = F.relu(torch.matmul(att, w1) + b1)
        a2 = F.relu(torch.matmul(a1, w2) + b2)

        return a2

dataroot = './dataset/xlsa/'
image_embedding = 'res101'
class_embedding = 'att'
dataset = args.dataset
matcontent = io.loadmat(dataroot + "/" + dataset + "/" + image_embedding + ".mat")

feature = matcontent['features'].T
label = matcontent['labels'].astype(int).squeeze() - 1
matcontent = io.loadmat(dataroot + "/" + dataset + "/" + class_embedding + "_splits.mat")

trainvalloc = matcontent['trainval_loc'].squeeze() - 1
test_seen_loc = matcontent['test_seen_loc'].squeeze() - 1
test_unseen_loc = matcontent['test_unseen_loc'].squeeze() - 1

att_name = 'att'
attribute = matcontent[att_name].T

chkFile = './dataset/finetune/'+args.dataset+'/dvbeExtracted.pkl'
with open(chkFile,'rb') as f:
    Feature_target = pickle.load(f)

clsname = [ matcontent['allclasses_names'][i][0][0] for i in range(len(matcontent['allclasses_names']))]

train_x = Feature_target['train_seen']['features']
train_label = Feature_target['train_seen']['labels']
train_att = attribute[train_label]
train_id, idx = np.unique(train_label, return_inverse=True)
train_att_unique = attribute[train_id]

test_x_unseen = Feature_target['test_unseen']['features']
test_label_unseen = Feature_target['test_unseen']['labels']
test_id, idx = np.unique(test_label_unseen, return_inverse=True)
att_pro = attribute[test_id]
train_test_id = np.concatenate((train_id, test_id))

test_x_seen = Feature_target['test_seen']['features']
test_label_seen = Feature_target['test_seen']['labels']
_, idx = np.unique(test_label_seen, return_inverse=True)

att_dim = train_att.shape[1]
feat_dim = train_x.shape[1]

att_pro = torch.from_numpy(att_pro).float().cuda()
test_x_seen = torch.from_numpy(test_x_seen).float().cuda()
test_x_seen = F.normalize(test_x_seen, p=2, dim=test_x_seen.dim()-1, eps=1e-12)
test_x_unseen = torch.from_numpy(test_x_unseen).float().cuda()
test_x_unseen = F.normalize(test_x_unseen, p=2, dim=test_x_unseen.dim()-1, eps=1e-12)
test_label_seen = torch.tensor(test_label_seen)
test_label_unseen = torch.tensor(test_label_unseen)

att_all = torch.from_numpy(train_att_unique).float().cuda()

bias = nn.Parameter(torch.FloatTensor(1).fill_(0).cuda(), requires_grad=True)
scale_cls = nn.Parameter(torch.FloatTensor(1).fill_(10).cuda(), requires_grad=True)
w1 = Variable(torch.FloatTensor(att_dim, args.hidden_dim).cuda(), requires_grad=True)
b1 = Variable(torch.FloatTensor(args.hidden_dim).cuda(), requires_grad=True)
w2 = Variable(torch.FloatTensor(args.hidden_dim, 2048).cuda(), requires_grad=True)
b2 = Variable(torch.FloatTensor(2048).cuda(), requires_grad=True)
w_IAS = Variable(torch.FloatTensor(2048, att_dim).cuda(), requires_grad=True)
b_IAS = Variable(torch.FloatTensor(att_dim).cuda(), requires_grad=True)

w1.data.normal_(0, 0.02)
w2.data.normal_(0, 0.02)
b1.data.fill_(0)
b2.data.fill_(0)
w_IAS.data.normal_(0,0.02)
b_IAS.data.fill_(0)

optimizer = torch.optim.Adam([w_IAS,b_IAS,w1, b1, w2, b2, bias, scale_cls], lr=args.lr, weight_decay=args.opt_decay)

# breakpoint()
step_size = args.step_size
gamma = args.gamma
lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=step_size, gamma=gamma)
criterion = nn.CrossEntropyLoss()

ways = args.ways
shots = args.shots

dataset = data_loader_virtualCls(train_x, train_att, train_label, ways=ways, shots=shots)

# breakpoint()
best_acc_zsl = 0.0
best_acc_gzsl_seen = 0.0
best_acc_gzsl_unseen = 0.0
best_H = 0.0
best_epoch = 0
best_unseenAcc = 0.0

for epoch in range(args.num_epochs):
        epoch_loss = 0
        lr_scheduler.step()

        for i in range(1000):
                batch_visual, batch_att, batch_label = dataset.__getitem__(i)
                batch_visual = batch_visual.cuda()
                batch_visual_norm = F.normalize(batch_visual, p=2, dim=batch_visual.dim()-1, eps=1e-12)

                indx = torch.tensor(list(range(0, ways*shots, shots)))
                unique_batch_att = torch.index_select(batch_att, 0, indx).float().cuda()

                batch_weights = forward(unique_batch_att,batch_visual_norm,tmp=TMP)
                all_cls_weights = batch_weights

                score = apply_classification_weights(batch_visual_norm, all_cls_weights)
                score = score.squeeze(0)
                loss = criterion(score, Variable(batch_label.cuda()))

                optimizer.zero_grad()
                loss.backward()
                torch.nn.utils.clip_grad_norm_([w_IAS,b_IAS,w1, b1, w2, b2, scale_cls, bias], 1)
                optimizer.step()
                epoch_loss = epoch_loss + loss

        epoch_loss = epoch_loss / 1000.0
        epoch_loss = epoch_loss.data.cpu().numpy()

        acc_zsl, seenAcc, acc_unseen_gzsl, acc_seen_gzsl, H, Rs,Ru,accs_cls,generalAccs_cls,unseenCos_cls,seenCos_cls,unseenpred_cls = compute_accuracy_all(att_pro, att_all, test_x_unseen,
                test_id, test_label_unseen, test_x_seen, train_test_id,  test_label_seen, train_id)

        H = 2 * acc_seen_gzsl * acc_unseen_gzsl / (acc_seen_gzsl + acc_unseen_gzsl)
        writer.add_scalar('general/acc_seen_gzsl',acc_seen_gzsl,epoch)
        writer.add_scalar('general/acc_unseen_gzsl',acc_unseen_gzsl,epoch)
        writer.add_scalar('general/H',H,epoch)
        writer.add_scalar('split/unseenAcc',acc_zsl,epoch)
        writer.add_scalar('split/seenAcc',seenAcc,epoch)
        writer.add_scalar('split/Rs',Rs,epoch)
        writer.add_scalar('split/Ru',Ru,epoch)
        writer.add_scalar('loss/loss',epoch_loss,epoch)

        if acc_zsl > best_unseenAcc:
                print('save best acc')
                best_unseenAcc = acc_zsl
                best_epoch = epoch
                best_acc_zsl = acc_zsl
                best_seen_acc = seenAcc
                best_acc_gzsl_seen = acc_seen_gzsl
                best_acc_gzsl_unseen = acc_unseen_gzsl
                best_H = H
                best_Ru = Ru
                best_Rs = Rs

                best_w1 = w1.data.clone()
                best_b1 = b1.data.clone()
                best_w2 = w2.data.clone()
                best_b2 = b2.data.clone()
                best_scale_cls = scale_cls.data.clone()
                best_bias = bias.data.clone()

                torch.save({'w1': best_w1, 'b1': best_b1, 'w2': best_w2, 'b2': best_b2,
                    'scale_cls': best_scale_cls, 'bias': best_bias,'w_IAS':w_IAS,'b_IAS':b_IAS}, model_file_name.replace('.pt','bestunseenAcc.pt'))

        for param_group in optimizer.param_groups:
                print('ep: %d,  lr: %lf, loss: %.4f,  zsl: %.4f, seenAcc: %.4f  gzsl: seen=%.4f, unseen=%.4f, h=%.4f, Rs=%.4f, Ru=%.4f ' %
                        (epoch, param_group['lr'],  epoch_loss, acc_zsl, seenAcc, acc_seen_gzsl, acc_unseen_gzsl, H, Rs, Ru,))
print(model_file_name)
print('best_ep: %d, zsl: %.4f, seenAcc: %.4f  gzsl: seen=%.4f, unseen=%.4f, h=%.4f, Rs=%.4f, Ru=%.4f' %
        (best_epoch, best_acc_zsl,best_seen_acc, best_acc_gzsl_seen, best_acc_gzsl_unseen, best_H, best_Rs, best_Ru))


import numpy as np
import array
import os, sys
import re
import time
import multiprocessing, logging
# mpl = multiprocessing.log_to_stderr()
# mpl.setLevel(logging.INFO)

import h5py
import logging
from astropy.table import Table, Column
from astropy import units as u
from scipy.interpolate import griddata

import argparse
parser = argparse.ArgumentParser()
parser.add_argument("-p","--params", type=str,
                    help = "Parameter file")
parser.add_argument("-q", "--quiet", help = "Suppress extra outputs",
                    action = "store_true")
args = parser.parse_args()
quiet = args.quiet

params_root = re.split(".py", args.params)[0]
if os.path.isfile(params_root+".pyc"):
    os.remove(params_root+".pyc")

import importlib
try:
    params = importlib.import_module(params_root)
    print('Successfully loaded "{}" as params'.format(args.params))
    #reload(params)
except:
    print('Failed to load "{}" as params'.format(args.params))
    raise

if quiet:
    quietprint = lambda *a: None
else:
    def quietprint(*args):
        for arg in args:
            print (arg)

# Fitting function definition for later use by Processess

def galaxyFit(inputQueue, printQueue, printlock):
    for gal in iter(inputQueue.get, 'STOP'):
        j = np.argmin(np.abs(z-zobs[gal])) # Find closest model redshift

        flux_obs = obs[gal,:]
        flux_err = obs_err[gal,:]

        #flux_obs[fo <= 0.] = 0.       # Set negative fluxes to zero
        I = np.where(flux_err > 0.)[0] # Find bands with no observation

        if len(I) == 0:
            if include_rest:
                M_scaled = np.ones(len(fo)) * -99.
                restframe_output = ' '.join(M_scaled.astype('str'))
                output_string = '{} {} {} {} {} {} {} {}' \
                                ' {} {} {} {} {} {} {} {} {}'.format(gal+1,ID[gal],zobs[gal],-99,-99,-99,-99,-99,-99, -99, -99,-99,len(I),-99,z[j],restframe_output,'\n')
            else:
                output_string = '{} {} {} {} {} {} {} {} {} {} {} {} {} {} {}'.format(gal+1,ID[gal],zobs[gal],-99,-99,-99,-99,-99,-99,-99, -99,-99,len(I),-99,'\n')
            printQueue.put(output_string)
            continue

        flux_obs = flux_obs[I]                    # and exclude from fit
        flux_err = flux_err[I]
        flux_models = f[j,I,:]

        tot_err = np.sqrt(flux_err**2 + (0.1*flux_obs)**2)

        top = 0.
        bottom = 0.

        for i in range(len(flux_obs)):
            top += (flux_models[i,:]*flux_obs[i])/(tot_err[i]**2)
            bottom += (flux_models[i,:]**2)/(tot_err[i]**2)

        scale = top/bottom
        scale = np.reshape(scale, (n_metal, n_tg, n_tau, n_tauv, n_fesc))

        chisq = 0.
        for i in range(len(flux_obs)):
            chisq += ((np.abs(scale*flux_models[i,:]-flux_obs[i])**2)/(flux_err[i])**2)

        chimin, minind = np.nanmin(chisq), np.nanargmin(chisq)

        if np.isinf(chimin) or np.isnan(minind):
            if include_rest:
                M_scaled = np.ones(len(flux_obs)) * -99.
                restframe_output = ' '.join(M_scaled.astype('str'))
                output_string = '{} {} {} {} {} {} {} {} {} {} {} {} {} {} {} {} {}'.format(gal+1,ID[gal],zobs[gal],-99,-99,-99,-99,-99,-99, -99, -99,-99,len(I),-99,z[j],restframe_output,'\n')
            else:
                output_string = '{} {} {} {} {} {} {} {} {} {} {} {} {} {} {}'.format(gal+1,ID[gal],zobs[gal],-99,-99,-99,-99,-99,-99,-99, -99,-99,len(I),-99,'\n')
            printQueue.put(output_string)
            continue

        #Find the coordinate of the model with the bestfit mass
        mi, tgi, ti, tvi, fi = np.unravel_index(minind,
                                                   (n_metal, n_tg,
                                                   n_tau, n_tauv, n_fesc))

        Bestfit_Mass = np.log10(scale[mi, tgi, ti, tvi, fi]*flux_corr)
        Bestfit_SFR = (scale[mi, tgi, ti, tvi, fi] *
                       SFR[mi, tgi, ti, tvi, fi]*flux_corr)
        #Bestfit_Beta = beta[tgi,tvi,ti,mi]
        Bestfit_Beta = -99.

        #Scale the observed tot_mag band of the template to be the same as the observed tot_mag band of the galaxy
        #Convert the templates so they are no longer units of per stellar mass

        F_rest = f[0,:]*scale[mi, tgi, ti, tvi, fi]*flux_corr
        restframeMags = 23.9 - 2.5*np.log10(F_rest)

        #UV_rest = UV_flux[0]*scale[tgi,tvi,ti,mi]*flux_corr
        #restframeMUV = 23.9 - 2.5*np.log10(UV_rest)

        M_scaled = restframeMags[:, mi, tgi, ti, tvi, fi]
        #MUV_scaled = restframeMUV[tgi,tvi,ti,mi]
        MUV_scaled = -99.

        if np.isnan(Bestfit_Mass) or np.isinf(chimin):
            Bestfit_Mass = -99
            #M_scaled[:] = -99
            tgs = -99
            tvs = -99
            taus = -99
            mis = -99
            escape_fraction = -99

        else:
            tgs = tg[tgi]/1e9
            tvs = tv[tvi]
            taus = tau[ti]
            mis = metallicities[mi]
            escape_fraction = fesc[fi]

        printlock.acquire()

        print('{:6d} {:8d} {:>5.2f} {:>7.2f} {:>8.1f} {:>8.3f} {:>5.1f} {:>8.2f} {:>4.2f} {:>5.2f}'.format(gal+1,ID[gal], zobs[gal],Bestfit_Mass,chimin,tgs,tvs,taus,mis,np.log10(Bestfit_SFR)))

        if include_rest:
            restframe_output = ' '.join(M_scaled.astype('str'))
            output_string = '{} {} {} {} {} {} {} {} {} {} {} {} {} {} {} {} {}'.format(gal+1,ID[gal],zobs[gal],Bestfit_Mass,chimin,tgs,tvs,taus,mis, MUV_scaled, minind,Bestfit_SFR,len(I),Bestfit_Beta,z[j],restframe_output,'\n')
        else:
            output_string = '{} {} {} {} {} {} {} {} {} {} {} {} {} {} {}'.format(gal+1,ID[gal],zobs[gal],Bestfit_Mass,chimin,tgs,tvs,taus,mis, MUV_scaled, minind,Bestfit_SFR,len(I),Bestfit_Beta,'\n')

        printlock.release()
        printQueue.put(output_string)

def galaxyFit2(inputQueue, printQueue, printlock):
    for gal in iter(inputQueue.get, 'STOP'):

        output_string = '{0[0]} {0[1]} {0[2]} {0[3]} {0[4]} {0[5]} ' + \
                        '{0[6]} {0[7]} {0[8]} {0[9]} {0[10]} {0[11]} ' + \
                        '{0[12]} {0[13]} {0[14]}'

        j = np.argmin(np.abs(z-zobs[gal])) # Find closest model redshift

        log_mass_min, log_mass_max = 7, 13
        log_sfr_min, log_sfr_max = -3, 4

        flux_obs = obs[gal,:]
        flux_err = obs_err[gal,:]

        #flux_obs[fo <= 0.] = 0.       # Set negative fluxes to zero
        #I = np.where(flux_err > 0.)[0] # Find bands with no observation
        I = (flux_err > 0.) * ((filt_lambda / (1+z[j])) < 3e5)

        flux_obs = flux_obs[I] * zp_offsets[I]                    # and exclude from fit
        flux_err = flux_err[I] * zp_offsets[I]
        flux_models = f[j,I,:,j,:]

        if params.temp_err != None:
            terr = griddata(terr_wl, terr_sigma, filt_lambda[I] / (1+z[j]))
            tot_err = np.sqrt(flux_err**2 + (terr*flux_obs)**2 + (params.flux_err*flux_obs)**2)
        else:
            tot_err = np.sqrt(flux_err**2 + (params.flux_err*flux_obs)**2)

        top = 0.
        bottom = 0.

        for i in range(len(flux_obs)):
            top += (flux_models[i,:]*flux_obs[i])/(tot_err[i]**2)
            bottom += (flux_models[i,:]**2)/(tot_err[i]**2)

        scale = top/bottom
        scale = np.reshape(scale, (n_metal, n_tau, n_tauv, n_fesc))

        chisq = 0.
        for i in range(len(flux_obs)):
            chisq += ((np.abs(scale*flux_models[i,:]-flux_obs[i])**2)/(tot_err[i])**2)

        chimin, minind = np.nanmin(chisq), np.nanargmin(chisq)
        likelihood = np.reshape(np.exp(-0.5*chisq),
                                (n_metal, n_tau, n_tauv, n_fesc))
        likelihood[np.isnan(likelihood)] = 0.
        likelihood = np.abs(likelihood/likelihood.sum())

        if np.isinf(chimin) or np.isnan(minind):
            output_string = '{n} {id} {zobs} {ztemp} {mass_best} {sfr_best} '+ \
                            '{chi_best} {tvs} {taus} {mis} {fesc} '+ \
                            '{mass_med} {mass_l68} {mass_u68} ' + \
                            '{sfr_med} {sfr_l68} {sfr_u68} ' + \
                            '{nfilts} '

            output_values = {'n': gal+1,
                             'id': ID[gal],
                             'zobs': zobs[gal], 'ztemp':z[j],
                             'mass_best': -99.,
                             'sfr_best': -99,
                             'chi_best': -99,
                             'tvs': -99, 'taus': -99,
                             'mis': -99, 'fesc': -99,
                             'mass_med': -99, 'mass_l68': -99, 'mass_u68': -99,
                             'sfr_med': -99, 'sfr_l68': -99, 'sfr_u68': -99,
                             'nfilts': np.sum(I)}

            output_array = [gal+1, ID[gal], zobs[gal],
                            Bestfit_Mass, chimin, tvs, taus, mis,
                            MUV_scaled, minind, Bestfit_SFR, np.sum(I), -99., '\n']
            output = output_string.format(**output_values)

            printlock.acquire()
            print_string = "{0[0]:6d} {0[1]:8d} {0[2]:>5.2f} " + \
                           "{0[3]:>7.2f} {0[4]:>8.3f} " + \
                           "{0[5]:>5.1f} {0[6]:>8.2f} {0[7]:>4.2f} " + \
                           "{0[8]:>5.2f}"

            print_array = [gal+1, ID[gal], zobs[gal],
                           -99, -99,
                           -99, -99, -99,
                           -99]
            print(print_string.format(print_array))
            printlock.release()

        else:
            #Find the coordinate of the model with the bestfit mass
            mi, ti, tvi, fi = np.unravel_index(minind,
                                                       (n_metal,
                                                       n_tau, n_tauv, n_fesc))

            Masses = np.abs(np.log10(scale*flux_corr))
            SFRs = np.log10(scale * SFR[:,j,:] * flux_corr)

            mass_hist = np.histogram(Masses.flatten(),
                                     range = (log_mass_min, log_mass_max),
                                     bins = 120,
                                     weights = likelihood.flatten(),
                                     density = True)

            sfr_hist = np.histogram(SFRs.flatten(),
                                     range = (log_sfr_min, log_sfr_max),
                                     bins = 140,
                                     weights = likelihood.flatten(),
                                     density = True)

            Bestfit_Mass = np.abs(np.log10(scale[mi, ti, tvi, fi]*flux_corr))
            Bestfit_SFR = np.abs(np.log10(scale[mi, ti, tvi, fi] *
                                   SFR[mi, j, ti, tvi, fi]*flux_corr))

            Bestfit_fluxes = (scale[mi, ti, tvi, fi] *
                              f[j,:, mi, j, ti, tvi, fi] *
                              flux_corr)

            tgs = tg[j]/1e9
            tvs = tv[tvi]
            taus = tau[ti]
            mis = metallicities[mi]
            escape_fraction = fesc[fi]

            m16, m50, m84 = weighted_quantile(Masses.flatten(),
                                              [0.16, 0.5, 0.84],
                                              sample_weight=likelihood.flatten(),
                                              values_sorted=False)
            s16, s50, s84 = weighted_quantile(SFRs.flatten(),
                                              [0.16, 0.5, 0.84],
                                              sample_weight=likelihood.flatten(),
                                              values_sorted=False)

            printlock.acquire()

            MUV_scaled = -99.
            Bestfit_Beta = -99.

            print_string = "{0[0]:6d} {0[1]:8d} {0[2]:>5.2f} " + \
                           "{0[3]:>7.2f} {0[4]:>8.3f} " + \
                           "{0[5]:>5.1f} {0[6]:>8.2f} {0[7]:>4.2f} " + \
                           "{0[8]:>5.2f}"

            print_array = [gal+1, ID[gal], zobs[gal],
                           Bestfit_Mass, chimin,
                           tvs, taus, mis,
                           Bestfit_SFR]
            print(print_string.format(print_array))
            printlock.release()
            output_string = '{n} {id} {zobs} {ztemp} {mass_best} {sfr_best} '+ \
                            '{chi_best} {tvs} {taus} {mis} {fesc} '+ \
                            '{mass_med} {mass_l68} {mass_u68} ' + \
                            '{sfr_med} {sfr_l68} {sfr_u68} ' + \
                            '{nfilts} '

            output_values = {'n': gal+1,
                             'id': ID[gal],
                             'zobs': zobs[gal], 'ztemp':z[j],
                             'mass_best': Bestfit_Mass,
                             'sfr_best': Bestfit_SFR,
                             'chi_best': chimin,
                             'tvs': tvs, 'taus': taus,
                             'mis': mis, 'fesc': escape_fraction,
                             'mass_med': m50, 'mass_l68': m16, 'mass_u68': m84,
                             'sfr_med': s50, 'sfr_l68': s16, 'sfr_u68': s84,
                             'nfilts': np.sum(I)}

            output_array = [gal+1, ID[gal], zobs[gal],
                            Bestfit_Mass, chimin, tvs, taus, mis,
                            MUV_scaled, minind, Bestfit_SFR, np.sum(I), -99., '\n']
            output = output_string.format(**output_values)

        if include_rest:
            if np.isinf(chimin) or np.isnan(minind):
                M_scaled = np.ones(len(flux_obs)) * -99.
                restframe_output = ' '.join(M_scaled.astype('str'))
                output = output + restframe_output + ' \n'

            else:
                F_rest = np.array(f[0, :, mi, j, ti, tvi, fi] *
                                  scale[mi, ti, tvi, fi] * flux_corr)
                restframeMags = 23.9 - 2.5*np.log10(F_rest)
                restframe_output = ' '.join(restframeMags.astype('str'))
                output = output + restframe_output + ' \n'
        else:
            output = output + ' \n'

        printQueue.put([gal, output, mass_hist, sfr_hist, Bestfit_fluxes,
                        [obs[gal, :], obs_err[gal, :]]])

def galaxyFitMz(inputQueue, printQueue, printlock):
    for gal in iter(inputQueue.get, 'STOP'):

        output_string = '{0[0]} {0[1]} {0[2]} {0[3]} {0[4]} {0[5]} ' + \
                        '{0[6]} {0[7]} {0[8]} {0[9]} {0[10]} {0[11]} ' + \
                        '{0[12]} {0[13]} {0[14]}'

        log_mass_min, log_mass_max = 7, 13
        log_sfr_min, log_sfr_max = -3, 4

        # Set up output arrays
        chi_z_best = np.zeros(len(z))

        m_z_best = np.zeros(len(z))
        m_z_median = np.zeros(len(z))
        m_z_u68 = np.zeros(len(z))
        m_z_l68 = np.zeros(len(z))

        sfr_z_best = np.zeros(len(z))
        sfr_z_median = np.zeros(len(z))
        sfr_z_u68 = np.zeros(len(z))
        sfr_z_l68 = np.zeros(len(z))

        #m_z_hist = np.zeros((len(z), 120))
        #sfr_z_hist = np.zeros((len(z), 140))

        flux_obs = obs[gal,:]
        flux_err = obs_err[gal,:]

        #flux_obs[fo <= 0.] = 0.       # Set negative fluxes to zero
        I = np.where(flux_err > 0.)[0] # Find bands with no observation

        if len(I) == 0:
            output_array = [gal+1, ID[gal], zobs[gal], z[j],
                            -99, -99, -99, -99, -99, -99, -99,
                            -99,-99,len(I),-99,'\n']
            output = output_string.format(output_array)

            if include_rest:
                M_scaled = np.ones(len(flux_obs)) * -99.
                restframe_output = ' '.join(M_scaled.astype('str'))
                output = output + restframe_output + ' \n'

            else:
                output = output + ' \n'
            printQueue.put(output_string)
            continue

        flux_obs = flux_obs[I]                    # and exclude from fit
        flux_err = flux_err[I]
        tot_err = np.sqrt(flux_err**2 + (params.flux_err*flux_obs)**2)

        for j, jz in enumerate(z):
            #j = np.argmin(np.abs(z-zobs[gal])) # Find closest model redshift
            flux_models = f[j,I,:,j]

            top = 0.
            bottom = 0.

            for i in range(len(flux_obs)):
                top += (flux_models[i,:]*flux_obs[i])/(tot_err[i]**2)
                bottom += (flux_models[i,:]**2)/(tot_err[i]**2)

            scale = top/bottom
            scale = np.reshape(scale, (n_metal, n_tau, n_tauv, n_fesc))

            chisq = 0.
            for i in range(len(flux_obs)):
                chisq += ((np.abs(scale*flux_models[i,:]-flux_obs[i])**2)/(tot_err[i])**2)

            chimin, minind = np.nanmin(chisq), np.nanargmin(chisq)
            likelihood = np.reshape(np.exp(-0.5*chisq),
                                    (n_metal, n_tau, n_tauv, n_fesc))
            likelihood[np.isnan(likelihood)] = 0.
            likelihood = np.abs(likelihood/likelihood.sum())

            if np.isinf(chimin) or np.isnan(minind):
                output_array = [gal+1, ID[gal], zobs[gal], z[j],
                                -99, -99, -99, -99, -99, -99,
                                -99,-99,len(I),-99,'\n']
                output = output_string.format(output_array)

            else:
                #Find the coordinate of the model with the bestfit mass
                mi, ti, tvi, fi = np.unravel_index(minind,
                                                   (n_metal,
                                                   n_tau, n_tauv, n_fesc))

                Masses = np.log10(np.abs(scale * flux_corr))
                SFRs = np.log10(np.abs(scale * SFR[:,j] * flux_corr))

                """
                mass_hist = np.histogram(Masses.flatten(),
                                         range = (log_mass_min, log_mass_max),
                                         bins = 120,
                                         weights = likelihood.flatten(),
                                         density = True)

                sfr_hist = np.histogram(SFRs.flatten(),
                                         range = (log_sfr_min, log_sfr_max),
                                         bins = 140,
                                         weights = likelihood.flatten(),
                                         density = True)
                """

                Bestfit_Mass = np.log10(np.abs(scale[mi, ti, tvi, fi]*flux_corr))
                Bestfit_SFR = np.log10(np.abs(scale[mi, ti, tvi, fi]) *
                                       SFR[mi, j, ti, tvi, fi]*flux_corr)

            if np.isnan(Bestfit_Mass) or np.isinf(chimin):
                Bestfit_Mass = -99
                #M_scaled[:] = -99
                tvs = -99
                taus = -99
                mis = -99
                escape_fraction = -99

            else:
                tvs = tv[tvi]
                taus = tau[ti]
                mis = metallicities[mi]
                escape_fraction = fesc[fi]

            m16, m50, m84 = weighted_quantile(Masses.flatten(),
                                              [0.16, 0.5, 0.84],
                                              sample_weight=likelihood.flatten(),
                                              values_sorted=False)
            s16, s50, s84 = weighted_quantile(SFRs.flatten(),
                                              [0.16, 0.5, 0.84],
                                              sample_weight=likelihood.flatten(),
                                              values_sorted=False)

            chi_z_best[j] = chimin
            m_z_best[j] = Bestfit_Mass
            sfr_z_best[j] = Bestfit_SFR

            m_z_l68[j], m_z_median[j], m_z_u68[j] = m16, m50, m84
            sfr_z_l68[j], sfr_z_median[j], sfr_z_u68[j] = s16, s50, s84

            MUV_scaled = -99.
            Bestfit_Beta = -99.

        printlock.acquire()
        j = np.argmin(np.abs(z-zobs[gal])) # Find closest model redshift
        print_string = "{0[0]:6d} {0[1]:8d} {0[2]:>5.2f} " + \
                       "{0[3]:>7.2f} {0[4]:>8.1f} {0[5]:>8.3f}"

        print_array = [gal+1, ID[gal], zobs[gal],
                       m_z_best[j], chi_z_best[j],
                       sfr_z_best[j]]
        print(print_string.format(print_array))

        output_string = '{n} {id} {zobs} {ztemp} {mass_best} {sfr_best} '+ \
                        '{chi_best} ' + \
                        '{mass_med} {mass_l68} {mass_u68} ' + \
                        '{sfr_med} {sfr_l68} {sfr_u68} ' + \
                        '{nfilts} '

        output_values = {'n': gal+1,
                         'id': ID[gal],
                         'zobs': zobs[gal], 'ztemp':z[j],
                         'mass_best': Bestfit_Mass,
                         'sfr_best': Bestfit_SFR,
                         'chi_best': chimin,
                         'mass_med': m50, 'mass_l68': m16, 'mass_u68': m84,
                         'sfr_med': s50, 'sfr_l68': s16, 'sfr_u68': s84,
                         'nfilts': len(I)}

        output = output_string.format(**output_values) + ' \n'

        printlock.release()
        printQueue.put([gal, output,
                        [m_z_best, sfr_z_best, chi_z_best],
                        [m_z_l68, m_z_median, m_z_u68]
                        [sfr_z_l68, sfr_z_median, sfr_z_u68]])

def getObservations(inputpath):
    input_data = Table.read(inputpath,format=input_format).filled(-99.)

    column_names = list(input_data.columns.keys())

    ID = input_data[ID_col]
    zobs = input_data[z_col]

    filter_names = []

    k,l = 0,0
    for ii in range(len(column_names)):
        if column_names[ii].lower().endswith(flux_col_end.lower()):
            if k == 0:
                fluxes = input_data[column_names[ii]]
            else:
                fluxes = np.column_stack((fluxes,input_data[column_names[ii]]))
            k+=1
            filter_names.append(column_names[ii])

        if column_names[ii].lower().endswith(fluxerr_col_end.lower()):
            if l == 0:
                fluxerrs = input_data[column_names[ii]]
            else:
                fluxerrs = np.column_stack((fluxerrs,input_data[column_names[ii]]))
            l+=1
    """
    if filts_used != None:
        try:
            fluxes = fluxes[:,filts_used]
            fluxerrs = fluxerrs[:,filts_used]
        except:r
            print('Filter mismatch 1')
            # Array slicing fail
    """
    return ID, zobs, fluxes, fluxerrs, k, filter_names

class _function_wrapper:
    """
    This is a hack to make the likelihood function pickleable when ``args``
    or ``kwargs`` are also included.

    Stolen from emcee
    """
    def __init__(self, f, args, kwargs):
        self.f = f
        self.args = args
        self.kwargs = kwargs

    def __call__(self, x):
        try:
            return self.f(x, *self.args, **self.kwargs)
        except:
            import traceback
            print("emcee: Exception while calling your likelihood function:")
            print("  params:", x)
            print("  args:", self.args)
            print("  kwargs:", self.kwargs)
            print("  exception:")
            traceback.print_exc()
            raise

def weighted_quantile(values, quantiles, sample_weight=None, values_sorted=False, old_style=False):
    """ Very close to np.percentile, but supports weights.
    NOTE: quantiles should be in [0, 1]!
    :param values: np.array with data
    :param quantiles: array-like with many quantiles needed
    :param sample_weight: array-like of the same length as `array`
    :param values_sorted: bool, if True, then will avoid sorting of initial array
    :param old_style: if True, will correct output to be consistent with np.percentile.
    :return: np.array with computed quantiles.
    """
    values = np.array(values)
    quantiles = np.array(quantiles)
    if sample_weight is None:
        sample_weight = np.ones(len(values))
    sample_weight = np.array(sample_weight)
    assert np.all(quantiles >= 0) and np.all(quantiles <= 1), 'quantiles should be in [0, 1]'

    if not values_sorted:
        sorter = np.argsort(values)
        values = values[sorter]
        sample_weight = sample_weight[sorter]

    weighted_quantiles = np.cumsum(sample_weight) - 0.5 * sample_weight
    if old_style:
        # To be convenient with np.percentile
        weighted_quantiles -= weighted_quantiles[0]
        weighted_quantiles /= weighted_quantiles[-1]
    else:
        weighted_quantiles /= np.sum(sample_weight)
    return np.interp(quantiles, weighted_quantiles, values)

if __name__ == '__main__':

    logfile = open("error.log", "w")
    original_stderr = sys.stderr
    sys.stderr = logfile

    start = time.time()

    """
    SECTION 1

    """
    model_path = params.model_path

    input_catalog = params.input_catalog
    input_format = params.input_format
    z_col = params.z_col
    ID_col = params.ID_col
    flux_col_end = params.flux_col_end
    fluxerr_col_end = params.fluxerr_col_end

    ncpus = params.ncpus
    filts_used = params.filts_used
    include_rest = params.include_rest

    output_path = params.output_catalog_path
    output_format = params.output_format
    output_hdf_path = params.output_hdf_path

    calc_mode = params.fitting_mode
    flux_corr = params.flux_corr

    ID, zobs, obs, obs_err, filters_found, filter_names = getObservations(input_catalog)

    """
    Section 2

    """

    print("Loading synthetic mags and mass array:")
    models = h5py.File(model_path, 'r')
    tg = models['ages'][()]
    tv = models['dust'][()]
    tau = models['sfh'][()]
    metallicities = models['metallicities'][()]
    fesc = models['fesc'][()]

    Mshape = models['fluxes'].shape
    z = models['z']
    nfilts = Mshape[1]
    n_metal = Mshape[2]
    n_tg = Mshape[3]
    n_tau = Mshape[4]
    n_tauv = Mshape[5]
    n_fesc = Mshape[6]

    filt_lambda = models['wl'][()]
    #UV_flux = synmags['UV_flux']
    SFR = models['SFR']
    Ms = models['Ms']

    if params.zp_offsets != None:
        zp_offsets = Table.read(params.zp_offsets, format='ascii.no_header')['col1']

    if params.temp_err != None:
        terr_wl, terr_sigma = np.loadtxt(params.temp_err).T

    if (nfilts == filters_found) and (filts_used == None):
        f = models['fluxes']

    elif (nfilts != filters_found) and (filts_used == None):
        raise Exception('Mis-match between model and observed filter numbers')

    elif filts_used != None:
        try:
            f = models['fluxes'][:,filts_used]
            obs = obs[:,filts_used]
            obs_err = obs_err[:,filts_used]
            filter_names = np.array(filter_names)[filts_used]
        except:
            print('Mis-match between model and observed filter numbers')
            raise
            # Slice fail

    print ("Done.")

    """
    SECTION 3
    """
    if os.path.isfile(output_path+".temp_output.txt"):
        os.remove(output_path+".temp_output.txt")
    temp_file = open(output_path+".temp_output.txt","w")

    """
    SECTION 4
    Chi-sq calculation

    """
    out_string = '{0:6s} {1:8s} {2:>5s} {3:>7s} {4:>8s}' + \
                 '{5:>5s} {6:>8s} {7:>4s} {8:>5s}'

    print(out_string.format('N','ID','zobs','Best', 'chimin',
                            'tauv','tau','met', 'sfr'))

    loop_start = time.time()
    ncpus = np.clip(ncpus, 1, multiprocessing.cpu_count())

    inputQueue = multiprocessing.Queue()
    printQueue = multiprocessing.Queue()
    printlock = multiprocessing.Lock()

    if calc_mode == 'hist':
        output_hdf = h5py.File(output_hdf_path, 'w')
        output_hdf.create_dataset("mass_pdf", (len(ID), 120), dtype="f")
        output_hdf.create_dataset("sfr_pdf", (len(ID), 140), dtype="f")
        output_hdf.create_dataset("fit_flux", (len(ID), f.shape[1]), dtype="f")
        output_hdf.create_dataset("obs_flux", (len(ID), f.shape[1]), dtype="f")
        output_hdf.create_dataset("obs_fluxerr", (len(ID), f.shape[1]),
                                  dtype="f")
        output_hdf.create_dataset("lambda_filt", data = models["wl"])
        output_hdf.create_dataset("fwhm_filt", data = models["fwhm"])

        fitFunction = galaxyFit2

    elif calc_mode == 'Mz':
        output_hdf = h5py.File(output_hdf_path, 'w')

        output_hdf.create_dataset("m_z_best", (len(ID), len(z)), dtype="f")
        output_hdf.create_dataset("sfr_z_best", (len(ID), len(z)), dtype="f")
        output_hdf.create_dataset("chi_z_best", (len(ID), len(z)), dtype="f")

        output_hdf.create_dataset("m_z_median", (len(ID), len(z)), dtype="f")
        output_hdf.create_dataset("m_z_l68", (len(ID), len(z)), dtype="f")
        output_hdf.create_dataset("m_z_u68", (len(ID), len(z)), dtype="f")

        output_hdf.create_dataset("sfr_z_median", (len(ID), len(z)), dtype="f")
        output_hdf.create_dataset("sfr_z_u68", (len(ID), len(z)), dtype="f")
        output_hdf.create_dataset("sfr_z_l68", (len(ID), len(z)), dtype="f")

        output_hdf.create_dataset("z", data=z)

        fitFunction = galaxyFitMz
    else:
        fitFunction = galaxyFit

    for i in range( ncpus ):
        multiprocessing.Process(target = fitFunction,
                                args = (inputQueue, printQueue,
                                        printlock)).start()

    # Put elements in the send queue for processing
    for gal in range( len(ID) ):
        inputQueue.put( gal )

    if calc_mode == 'hist':
        for i, gal in enumerate(ID):
            j, out, mass_hist, sfr_hist, fit_flux, obs_flux = printQueue.get()

            if i == 0:
                mass_centers = 0.5*(mass_hist[1][1:] + mass_hist[1][:-1])
                sfr_centers = 0.5*(sfr_hist[1][1:] + sfr_hist[1][:-1])

                output_hdf.create_dataset("mass_bins", data = mass_centers)
                output_hdf.create_dataset("sfr_bins", data = sfr_centers)

            output_hdf["mass_pdf"][j] = mass_hist[0]
            output_hdf["sfr_pdf"][j] = sfr_hist[0]
            output_hdf["fit_flux"][j] = fit_flux
            output_hdf["obs_flux"][j] = obs_flux[0]
            output_hdf["obs_fluxerr"][j] = obs_flux[1]

            temp_file.write( out )

    elif calc_mode == 'Mz':
        for i, gal in enumerate(ID):
            j, out, pz_best, mz_median, sfrz_median = printQueue.get()

            output_hdf["m_z_best"][j,:] = pz_best[0]
            output_hdf["sfr_z_best"][j,:] = pz_best[1]
            output_hdf["chi_z_best"][j,:] = pz_best[2]

            output_hdf["m_z_l68"][j,:] = mz_median[0]
            output_hdf["m_z_median"][j,:] = mz_median[1]
            output_hdf["m_z_u68"][j,:] = mz_median[2]

            output_hdf["sfr_z_l68"][j,:] = sfrz_median[0]
            output_hdf["sfr_z_median"][j,:] = sfrz_median[1]
            output_hdf["sfr_z_u68"][j,:] = sfrz_median[2]

            temp_file.write( out )
    else:
        for i, gal in enumerate(ID):
            printout = printQueue.get()
            temp_file.write( printout )
            #print len(mass_array), len(muv_array), len(beta_array)

    # Stop all the running processes
    for i in range( ncpus ):
        inputQueue.put( 'STOP' )

    # Close both send and receive queues
    inputQueue.close()
    printQueue.close()

    temp_file.close()
    models.close()
    output_hdf.close()
    print("Fitting time taken: {:.2f} {}".format(time.time()-loop_start,
                                                   '\n'))

    """
    Section 3
    Reload, format and save output table
    """
    while temp_file.closed == False:
        pause(0.1)

    data = np.loadtxt(output_path+".temp_output.txt")
    try:
        rows, cols = data.shape
    except:
        cols = len(data)

    output = Table()

    names = ['N', 'ID', 'z', 'zmodel',
             'Mass_best', 'SFR_best', 'chi_best',
             'Dust_best', 'SFH_best',
             'Metallicity_best', 'fesc_best',
             'Mass_median', 'Mass_l68', 'Mass_u68',
             'SFR_median', 'SFR_l68', 'SFR_u68',
             'Nfilts']

    units = [None, None, None, None,
             u.Msun, u.Msun/u.yr, None,
             None, None,
             None, None,
             u.Msun, u.Msun, u.Msun,
             u.Msun/u.yr, u.Msun/u.yr, u.Msun/u.yr,
             None]

    types = ['i4', 'i4', 'f4', 'f4',
             'f4', 'f4', 'f4',
             'f4', 'f4',
             'f4', 'f4',
             'f4', 'f4', 'f4',
             'f4', 'f4', 'f4',
             'i4']

    if include_rest:
        for name in filter_names:
            names.append(name[:-len(flux_col_end)]+'_rest')
            units.append(u.mag)
            types.append('f4')

    for col in range(cols):
        column = Column( data[:,col], name = names[col], unit=units[col], dtype=types[col])
        output.add_column(column)

    table_format = params.output_format
    output.sort('ID')
    if os.path.isfile(output_path):
        os.remove(output_path)
    output.write(output_path,format=table_format, overwrite=True)
    print('Catalog saved')

    os.remove(temp_file.name)

    print('\n')
    print("Total time taken: "+str(time.time()-start))

    sys.stderr = original_stderr
    logfile.close()