# Copyright 2018-2021 Xanadu Quantum Technologies Inc.

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

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

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Function wrappers for the TensorBox API"""
# pylint:disable=abstract-class-instantiated,unexpected-keyword-arg
from collections.abc import Sequence
import itertools
import warnings

import numpy as np

from .tensorbox import TensorBox

def _get_multi_tensorbox(values):
    """Determines the correct framework to dispatch to given a
    sequence of tensor-like objects.

    Args:
        values (Sequence[tensor_like]): a sequence of tensor like objects

    Returns:
        .TensorBox: A TensorBox that will dispatch to the correct framework
        given the rules of precedence. This TensorBox will contain the *first*
        tensor-like object in ``values`` that corresponds to the highest-priority
        framework.

    To determine the framework to dispatch to, the following rules
    are applied:

    * Tensors that are incompatible (such as Torch and TensorFlow tensors)
      cannot both be present.

    * Autograd tensors *may* be present alongside Torch and TensorFlow tensors,
      but Torch and TensorFlow take precendence; the autograd arrays will
      be treated as non-differentiable NumPy arrays. A warning will be raised
      suggesting that vanilla NumPy be used instead.

    * Vanilla NumPy arrays can be used alongside other tensor objects; they will
      always be treated as non-differentiable constants.
    """
    interfaces = [get_interface(v) for v in values]

    if len(set(interfaces) - {"numpy", "autograd"}) > 1:
        # contains multiple non-autograd interfaces
        raise ValueError("Tensors contain mixed types; cannot determine dispatch library")

    non_numpy_interfaces = set(interfaces) - {"numpy"}

    if len(non_numpy_interfaces) > 1:
        # contains autograd and another interface
        warnings.warn(
            f"Contains tensors of types {non_numpy_interfaces}; dispatch will prioritize "
            "TensorFlow and PyTorch over autograd. Consider replacing Autograd with vanilla NumPy.",
            UserWarning,
        )

    if "tf" in interfaces:
        return TensorBox(values[interfaces.index("tf")])

    if "torch" in interfaces:
        return TensorBox(values[interfaces.index("torch")])

    if "autograd" in interfaces:
        return TensorBox(values[interfaces.index("autograd")])

    if "jax" in interfaces:
        return TensorBox(values[interfaces.index("jax")])

    return TensorBox(values[interfaces.index("numpy")])

def abs_(tensor):
    """Returns the element-wise absolute value.

    Args:
        tensor (tensor_like): input tensor

    Returns:
        tensor_like:

    **Example**

    >>> a = torch.tensor([1., -2.], requires_grad=True)
    >>> abs(a)
    tensor([1., 2.], grad_fn=<AbsBackward>)
    """
    return TensorBox(tensor).abs(wrap_output=False)

def allequal(tensor1, tensor2, **kwargs):
    """Returns True if two tensors are element-wise equal along a given axis.

    This function is equivalent to calling ``np.all(tensor1 == tensor2, **kwargs)``,
    but allows for ``tensor1`` and ``tensor2`` to differ in type.

    Args:
        tensor1 (tensor_like): tensor to compare
        tensor2 (tensor_like): tensor to compare
        **kwargs: Accepts any keyword argument that is accepted by ``np.all``,
            such as ``axis``, ``out``, and ``keepdims``. See the `NumPy documentation
            <https://numpy.org/doc/stable/reference/generated/numpy.all.html>`__ for
            more details.

    Returns:
        ndarray, bool: If ``axis=None``, a logical AND reduction is applied to all elements
        and a boolean will be returned, indicating if all elements evaluate to True. Otherwise,
        a boolean NumPy array will be returned.

    **Example**

    >>> a = torch.tensor([1, 2])
    >>> b = np.array([1, 2])
    >>> allequal(a, b)
    True
    """
    t1 = toarray(tensor1)
    t2 = toarray(tensor2)
    return np.all(t1 == t2, **kwargs)

def allclose(a, b, rtol=1e-05, atol=1e-08, **kwargs):
    """Wrapper around np.allclose, allowing tensors ``a`` and ``b``
    to differ in type"""
    t1 = toarray(a)
    t2 = toarray(b)
    return np.allclose(t1, t2, rtol=rtol, atol=atol, **kwargs)

allclose.__doc__ = np.allclose.__doc__

def angle(tensor):
    """Returns the element-wise angle of a complex tensor.

    Args:
        tensor (tensor_like): input tensor

    Returns:
        tensor_like:

    **Example**

    >>> a = torch.tensor([1.0, 1.0j, 1+1j], requires_grad=True)
    >>> angle(a)
    tensor([0.0000, 1.5708, 0.7854], grad_fn=<AngleBackward>)
    """
    return TensorBox(tensor).angle(wrap_output=False)

def arcsin(tensor):
    """Returns the element-wise inverse sine of the tensor"""
    return TensorBox(tensor).arcsin(wrap_output=False)

def block_diag(values):
    """Combine a sequence of 2D tensors to form a block diagonal tensor.

    Args:
        values (Sequence[tensor_like]): Sequence of 2D arrays/tensors to form
            the block diagonal tensor.

    Returns:
        tensor_like: the block diagonal tensor

    **Example**

    >>> t = [
    ...     np.array([[1, 2], [3, 4]]),
    ...     torch.tensor([[1, 2, 3], [-1, -6, -3]]),
    ...     torch.tensor(5)
    ... ]
    >>> qml.math.block_diag(t)
    tensor([[ 1,  2,  0,  0,  0,  0],
            [ 3,  4,  0,  0,  0,  0],
            [ 0,  0,  1,  2,  3,  0],
            [ 0,  0, -1, -6, -3,  0],
            [ 0,  0,  0,  0,  0,  5]])
    """
    return _get_multi_tensorbox(values).block_diag(values, wrap_output=False)

def cast(tensor, dtype):
    """Casts the given tensor to a new type.

    Args:
        tensor (tensor_like): tensor to cast
        dtype (str, np.dtype): Any supported NumPy dtype representation; this can be
            a string (``"float64"``), a ``np.dtype`` object (``np.dtype("float64")``), or
            a dtype class (``np.float64``). If ``tensor`` is not a NumPy array, the
            **equivalent** dtype in the dispatched framework is used.

    Returns:
        tensor_like: a tensor with the same shape and values as ``tensor`` and the
        same dtype as ``dtype``

    **Example**

    We can use NumPy dtype specifiers:

    >>> x = torch.tensor([1, 2])
    >>> cast(x, np.float64)
    tensor([1., 2.], dtype=torch.float64)

    We can also use strings:

    >>> x = tf.Variable([1, 2])
    >>> cast(x, "complex128")
    <tf.Tensor: shape=(2,), dtype=complex128, numpy=array([1.+0.j, 2.+0.j])>
    """
    return TensorBox(tensor).cast(dtype, wrap_output=False)

def cast_like(tensor1, tensor2):
    """Casts a tensor to the same dtype as another.

    Args:
        tensor1 (tensor_like): tensor to cast
        tensor2 (tensor_like): tensor with corresponding dtype to cast to

    Returns:
        tensor_like: a tensor with the same shape and values as ``tensor1`` and the
        same dtype as ``tensor2``

    **Example**

    >>> x = torch.tensor([1, 2])
    >>> y = torch.tensor([3., 4.])
    >>> cast(x, y)
    tensor([1., 2.])
    """
    dtype = toarray(tensor2).dtype.type
    return cast(tensor1, dtype)

def concatenate(values, axis=0):
    """Concatenate a sequence of tensors along the specified axis.

    .. warning::

        Tensors that are incompatible (such as Torch and TensorFlow tensors)
        cannot both be present.

    Args:
        values (Sequence[tensor_like]): Sequence of tensor-like objects to
            concatenate. The objects must have the same shape, except in the dimension corresponding
            to axis (the first, by default).
        axis (int): The axis along which the input tensors are concatenated. If axis is None,
            tensors are flattened before use. Default is 0.

    Returns:
        tensor_like: The concatenated tensor.

    **Example**

    >>> x = tf.constant([0.6, 0.1, 0.6])
    >>> y = tf.Variable([0.1, 0.2, 0.3])
    >>> z = np.array([5., 8., 101.])
    >>> concatenate([x, y, z])
    TensorBox: <tf.Tensor: shape=(3, 3), dtype=float32, numpy=
    array([6.00e-01, 1.00e-01, 6.00e-01, 1.00e-01, 2.00e-01, 3.00e-01, 5.00e+00, 8.00e+00, 1.01e+02], dtype=float32)>
    """
    return _get_multi_tensorbox(values).concatenate(values, axis=axis, wrap_output=False)

def conj(tensor):
    """Conjugate a tensor. Negate the imaginary part of a complex value.

    Args:
        tensor (tensor_like): A tensor-like object to conjugate.

    Returns:
        tensor_like: The conjugated tensor.

    **Example**

    >>> x = tf.constant([0.6 + 0.1j, 0.1 - 0.3j, 0.6])
    >>> conj(x)
    <tf.Tensor: shape=(3,), dtype=complex64, numpy=array([6.00e-01 + 1.00e-1j, 1.00e-01 + 3.00e-1j, 6.00e-01 + 0.00j], dtype=complex64)>
    """
    return TensorBox(tensor).conj(wrap_output=False)

def cov_matrix(prob, obs, wires=None, diag_approx=False):
    """Calculate the covariance matrix of a list of commuting observables, given
    the joint probability distribution of the system in the shared eigenbasis.

    .. note::
        This method only works for **commuting observables.**
        If the probability distribution is the result of a quantum circuit,
        the quantum state must be rotated into the shared
        eigenbasis of the list of observables before measurement.

    Args:
        prob (tensor_like): probability distribution
        obs (list[.Observable]): a list of observables for which
            to compute the covariance matrix for
        diag_approx (bool): if True, return the diagonal approximation
        wires (.Wires): The wire register of the system. If not provided,
            it is assumed that the wires are labelled with consecutive integers.

    Returns:
        tensor_like: the covariance matrix of size ``(len(obs), len(obs))``

    **Example**

    Consider the following ansatz and observable list:

    >>> obs_list = [qml.PauliX(0) @ qml.PauliZ(1), qml.PauliY(2)]
    >>> ansatz = qml.templates.StronglyEntanglingLayers

    We can construct a QNode to output the probability distribution in the shared eigenbasis of the
    observables:

    .. code-block:: python

        dev = qml.device("default.qubit", wires=3)

        @qml.qnode(dev, interface="autograd")
        def circuit(weights):
            ansatz(weights, wires=[0, 1, 2])
            # rotate into the basis of the observables
            for o in obs_list:
                o.diagonalizing_gates()
            return qml.probs(wires=[0, 1, 2])

    We can now compute the covariance matrix:

    >>> weights = qml.init.strong_ent_layers_normal(n_layers=2, n_wires=3)
    >>> cov = qml.math.cov_matrix(circuit(weights), obs_list)
    >>> cov
    array([[0.98707611, 0.03665537],
         [0.03665537, 0.99998377]])

    Autodifferentiation is fully supported using all interfaces.
    Here we use autograd:

    >>> cost_fn = lambda weights: qml.math.cov_matrix(circuit(weights), obs_list)[0, 1]
    >>> qml.grad(cost_fn)(weights)[0]
    array([[[ 4.94240914e-17, -2.33786398e-01, -1.54193959e-01],
            [-3.05414996e-17,  8.40072236e-04,  5.57884080e-04],
            [ 3.01859411e-17,  8.60411436e-03,  6.15745204e-04]],
           [[ 6.80309533e-04, -1.23162742e-03,  1.08729813e-03],
            [-1.53863193e-01, -1.38700657e-02, -1.36243323e-01],
            [-1.54665054e-01, -1.89018172e-02, -1.56415558e-01]]])
    """
    variances = []

    # diagonal variances
    for i, o in enumerate(obs):
        l = cast(o.eigvals, dtype=np.float64)
        w = o.wires.labels if wires is None else wires.indices(o.wires)
        p = marginal_prob(prob, w)

        res = dot(l ** 2, p) - (dot(l, p)) ** 2
        variances.append(res)

    cov = diag(variances)

    if diag_approx:
        return cov

    for i, j in itertools.combinations(range(len(obs)), r=2):
        o1 = obs[i]
        o2 = obs[j]

        o1wires = o1.wires.labels if wires is None else wires.indices(o1.wires)
        o2wires = o2.wires.labels if wires is None else wires.indices(o2.wires)
        shared_wires = set(o1wires + o2wires)

        l1 = cast(o1.eigvals, dtype=np.float64)
        l2 = cast(o2.eigvals, dtype=np.float64)
        l12 = cast(np.kron(l1, l2), dtype=np.float64)

        p1 = marginal_prob(prob, o1wires)
        p2 = marginal_prob(prob, o2wires)
        p12 = marginal_prob(prob, shared_wires)

        res = dot(l12, p12) - dot(l1, p1) * dot(l2, p2)

        cov = scatter_element_add(cov, [i, j], res)
        cov = scatter_element_add(cov, [j, i], res)

    return cov

def convert_like(tensor1, tensor2):
    """Convert a tensor to the same type as another.

    Args:
        tensor1 (tensor_like): tensor to convert
        tensor2 (tensor_like): tensor with corresponding type to convert to

    Returns:
        tensor_like: a tensor with the same shape, values, and dtype as ``tensor1`` and the
        same type as ``tensor2``.

    **Example**

    >>> x = np.array([1, 2])
    >>> y = tf.Variable([3, 4])
    >>> cast(x, y)
    <tf.Tensor: shape=(2,), dtype=int64, numpy=array([1, 2])>
    """
    return TensorBox(tensor2).astensor(tensor1)

def diag(values, k=0):
    """Construct a diagonal tensor from a list of scalars.

    Args:
        values (tensor_like or Sequence[scalar]): sequence of numeric values that
            make up the diagonal
        k (int): The diagonal in question. ``k=0`` corresponds to the main diagonal.
            Use ``k>0`` for diagonals above the main diagonal, and ``k<0`` for
            diagonals below the main diagonal.

    Returns:
        tensor_like: the 2D diagonal tensor

    **Example**

    >>> x = [1., 2., tf.Variable(3.)]
    >>> diag(x)
    <tf.Tensor: shape=(3, 3), dtype=float32, numpy=
    array([[1., 0., 0.],
           [0., 2., 0.],
           [0., 0., 3.]], dtype=float32)>
    >>> y = tf.Variable([0.65, 0.2, 0.1])
    >>> diag(y, k=-1)
    <tf.Tensor: shape=(4, 4), dtype=float32, numpy=
    array([[0.  , 0.  , 0.  , 0.  ],
           [0.65, 0.  , 0.  , 0.  ],
           [0.  , 0.2 , 0.  , 0.  ],
           [0.  , 0.  , 0.1 , 0.  ]], dtype=float32)>
    >>> z = torch.tensor([0.1, 0.2])
    >>> qml.diag(z, k=1)
    >>> qml.math.diag(z, k=1)
    tensor([[0.0000, 0.1000, 0.0000],
            [0.0000, 0.0000, 0.2000],
            [0.0000, 0.0000, 0.0000]])
    """
    if isinstance(values, Sequence):
        return _get_multi_tensorbox(values).diag(values, k=k, wrap_output=False)

    return TensorBox(values).diag(values, k=k, wrap_output=False)

def dot(tensor1, tensor2):
    """Returns the matrix or dot product of two tensors.

    * If both tensors are 0-dimensional, elementwise multiplication
      is performed and a 0-dimensional scalar returned.

    * If both tensors are 1-dimensional, the dot product is returned.

    * If the first array is 2-dimensional and the second array 1-dimensional,
      the matrix-vector product is returned.

    * If both tensors are 2-dimensional, the matrix product is returned.

    * Finally, if the the first array is N-dimensional and the second array
      M-dimensional, a sum product over the last dimension of the first array,
      and the second-to-last dimension of the second array is returned.

    Args:
        tensor1 (tensor_like): input tensor
        tensor2 (tensor_like): input tensor
    """
    return _get_multi_tensorbox([tensor1, tensor2]).dot(tensor1, tensor2, wrap_output=False)

def expand_dims(tensor, axis):
    """Expand the shape of an array by adding a new dimension of size 1
    at the specified axis location.

    .. warning::

        This function differs from ``np.expand_dims``.

    Args:
        tensor (tensor_like): tensor to expand
        axis (int): location in the axes to place the new dimension

    Returns:
        tensor_like: a tensor with the expanded shape

    **Example**

    >>> x = tf.Variable([3, 4])
    >>> expand_dims(x, axis=1)
    <tf.Tensor: shape=(2, 1), dtype=int32, numpy=
    array([[3],
           [4]], dtype=int32)>
    """
    return TensorBox(tensor).expand_dims(axis, wrap_output=False)

def flatten(tensor):
    """Flattens an N-dimensional tensor to a 1-dimensional tensor.

    Args:
        tensor (tensor_like): tensor to flatten

    Returns:
        tensor_like: the flattened tensor

    **Example**

    >>> x = tf.Variable([[1, 3], [2, 4]])
    >>> flatten(x)
    <tf.Tensor: shape=(4,), dtype=int32, numpy=array([1, 3, 2, 4], dtype=int32)>
    """
    return reshape(tensor, (-1,))

def gather(tensor, indices):
    """Gather tensor values given a tuple of indices.

    This is equivalent to the following NumPy fancy indexing:

    ..code-block:: python

        tensor[array(indices)]

    Args:
        tensor (tensor_like): tensor to gather from
        indices (Sequence[int]): the indices of the values to extract

    Returns:

        tensor_like: the gathered tensor values

    .. seealso::

        :func:`~.take`
    """
    return TensorBox(tensor).gather(np.array(indices), wrap_output=False)

def get_interface(tensor):
    """Returns the name of the package that any array/tensor manipulations
    will dispatch to. The returned strings correspond to those used for PennyLane
    :doc:`interfaces </introduction/interfaces>`.

    Args:
        tensor (tensor_like): tensor input

    Returns:
        str: name of the interface

    **Example**

    >>> x = torch.tensor([1., 2.])
    >>> get_interface(x)
    'torch'
    >>> from pennylane import numpy as np
    >>> x = np.array([4, 5], requires_grad=True)
    >>> get_interface(x)
    'autograd'
    """
    return TensorBox(tensor).interface

def marginal_prob(prob, axis):
    """Compute the marginal probability given a joint probability distribution expressed as a tensor.
    Each random variable corresponds to a dimension.

    If the distribution arises from a quantum circuit measured in computational basis, each dimension
    corresponds to a wire. For example, for a 2-qubit quantum circuit `prob[0, 1]` is the probability of measuring the
    first qubit in state 0 and the second in state 1.

    Args:
        prob (tensor_like): 1D tensor of probabilities. This tensor should of size
            ``(2**N,)`` for some integer value ``N``.
        axis (list[int]): the axis for which to calculate the marginal
            probability distribution

    Returns:
        tensor_like: the marginal probabilities, of
        size ``(2**len(axis),)``

    **Example**

    >>> x = tf.Variable([1, 0, 0, 1.], dtype=tf.float64) / np.sqrt(2)
    >>> marginal_prob(x, axis=[0, 1])
    <tf.Tensor: shape=(4,), dtype=float64, numpy=array([0.70710678, 0.        , 0.        , 0.70710678])>
    >>> marginal_prob(x, axis=[0])
    <tf.Tensor: shape=(2,), dtype=float64, numpy=array([0.70710678, 0.70710678])>
    """
    prob = flatten(prob)
    num_wires = int(np.log2(len(prob)))

    if num_wires == len(axis):
        return prob

    inactive_wires = tuple(set(range(num_wires)) - set(axis))
    prob = reshape(prob, [2] * num_wires)
    prob = sum_(prob, axis=inactive_wires)
    return flatten(prob)

def toarray(tensor):
    """Returns the tensor as a NumPy ``ndarray``. No copying
    is performed; the tensor and the returned array share the
    same storage.

    Args:
        tensor (tensor_like): input tensor

    Returns:
        array: a ``ndarray`` view into the same data

    **Example**

    >>> x = torch.tensor([1., 2.])
    >>> toarray(x)
    array([1, 2])
    """
    return TensorBox(tensor).numpy()

def ones_like(tensor, dtype=None):
    """Returns a tensor of all ones with the same shape and dtype
    as the input tensor.

    Args:
        tensor (tensor_like): input tensor
        dtype (str, np.dtype): The desired output datatype of the array. If not provided, the dtype of

            ``tensor`` is used. This argument can be any supported NumPy dtype representation, including
            a string (``"float64"``), a ``np.dtype`` object (``np.dtype("float64")``), or
            a dtype class (``np.float64``). If ``tensor`` is not a NumPy array, the
            **equivalent** dtype in the dispatched framework is used.

    Returns:
        tensor_like: an all-ones tensor with the same shape and
        size as ``tensor``

    **Example**

    >>> x = torch.tensor([1., 2.])
    >>> ones_like(x)
    tensor([1, 1])
    >>> y = tf.Variable([[0], [5]])
    >>> ones_like(y, dtype=np.complex128)
    <tf.Tensor: shape=(2, 1), dtype=complex128, numpy=
    array([[1.+0.j],
           [1.+0.j]])>
    """
    if dtype is not None:
        return TensorBox(tensor).ones_like().cast(dtype, wrap_output=False)

    return TensorBox(tensor).ones_like(wrap_output=False)

def reshape(tensor, shape):  # pylint: disable=redefined-outer-name
    """Gives a new shape to a tensor without changing its data.

    Args:
        tensor (tensor_like): input tensor
        shape (tuple[int]): The new shape. The special value of -1 indicates
            that the size of that dimension is computed so that the total size
            remains constant. A dimension of -1 can only be specified once.

    Returns:
        tensor_like: a new view into the input tensor with
        shape ``shape``

    **Example**

    >>> a = tf.range(4.)
    >>> reshape(a, (2, 2))
    <tf.Tensor: shape=(2, 2), dtype=float32, numpy=
    array([[0., 1.],
           [2., 3.]], dtype=float32)>
    >>> b = torch.tensor([[0, 1], [2, 3]])
    >>> torch.reshape(b, (-1,))
    tensor([0, 1, 2, 3])
    """
    return TensorBox(tensor).reshape(shape, wrap_output=False)

def requires_grad(tensor):
    """Returns True if the tensor is considered trainable.

    .. warning::

        The implemetation depends on the contained tensor type, and
        may be context dependent.

        For example, Torch tensors and PennyLane tensors track trainability
        as a property of the tensor itself. TensorFlow, on the other hand,

        only tracks trainability if being watched by a gradient tape.

    Args:
        tensor (tensor_like): input tensor

    **Example**

    Calling this function on a PennyLane NumPy array:

    >>> x = np.array([1., 5.], requires_grad=True)
    >>> requires_grad(x)
    True
    >>> x.requires_grad = False
    >>> requires_grad(x)
    False

    PyTorch has similar behaviour.

    With TensorFlow, the output is dependent on whether the tensor
    is currently being watched by a gradient tape:

    >>> x = tf.Variable([0.6, 0.1])
    >>> requires_grad(x)
    False
    >>> with tf.GradientTape() as tape:
    ...     print(requires_grad(x))
    True

    While TensorFlow constants are by default not trainable, they can be
    manually watched by the gradient tape:

    >>> x = tf.constant([0.6, 0.1])
    >>> with tf.GradientTape() as tape:
    ...     print(requires_grad(x))
    False
    >>> with tf.GradientTape() as tape:
    ...     tape.watch([x])
    ...     print(requires_grad(x))
    True
    """
    return TensorBox(tensor).requires_grad

def scatter_element_add(tensor, index, value):
    """Adds a scalar value to a specific index of a tensor.

    This is a pure equivalent of ``tensor[index] += value``.

    Args:
        tensor (tensor_like): the input tensor to be updated
        index (tuple[int]): the index of the input tensor to update
        value (scalar): the scalar value to add to the tensor element

    Returns:
        tensor_like: the output tensor

    **Example**

    >>> x = torch.ones((2, 3))
    >>> qml.math.scatter_element_add(x, [1, 2], 3)
    tensor([[1., 1., 1.],
            [1., 1., 4.]])
    """
    value = convert_like(value, tensor)
    return TensorBox(tensor).scatter_element_add(index, value, wrap_output=False)

def shape(tensor):
    """Returns the shape of the tensor.

    Args:
        tensor (tensor_like): input tensor

    Returns:
        tuple[int]: shape of the tensor

    **Example**

    >>> x = tf.constant([[0.6, 0.1, 0.6], [1., 2., 3.]])
    >>> shape(x)
    (2, 3)
    """
    return TensorBox(tensor).shape

def sqrt(tensor):
    """Returns the element-wise square root.

    Args:
        tensor (tensor_like): input tensor

    Returns:
        tensor_like:

    **Example**

    >>> a = torch.tensor([4., 9.], requires_grad=True)
    >>> sqrt(a)
    tensor([2., 3.], grad_fn=<SqrtBackward>)
    """
    return TensorBox(tensor).sqrt(wrap_output=False)

def stack(values, axis=0):
    """Stack a sequence of tensors along the specified axis.

    .. warning::

        Tensors that are incompatible (such as Torch and TensorFlow tensors)
        cannot both be present.

    Args:
        values (Sequence[tensor_like]): Sequence of tensor-like objects to
            stack. Each object in the sequence must have the same size in the given axis.
        axis (int): The axis along which the input tensors are stacked. ``axis=0`` corresponds
            to vertical stacking.

    Returns:
        tensor_like: The stacked array. The stacked array will have one additional dimension
        compared to the unstacked tensors.

    **Example**

    >>> x = tf.constant([0.6, 0.1, 0.6])
    >>> y = tf.Variable([0.1, 0.2, 0.3])
    >>> z = np.array([5., 8., 101.])
    >>> stack([x, y, z])
    <tf.Tensor: shape=(3, 3), dtype=float32, numpy=
    array([[6.00e-01, 1.00e-01, 6.00e-01],
           [1.00e-01, 2.00e-01, 3.00e-01],
           [5.00e+00, 8.00e+00, 1.01e+02]], dtype=float32)>
    """
    return _get_multi_tensorbox(values).stack(values, axis=axis, wrap_output=False)

def squeeze(tensor):
    """Remove single-dimensional entries from the shape of an array.

    Args:
        tensor (tensor_like): A tensor-like object.

    Returns:
        The input array, but with all or a subset of the dimensions of length 1 removed.
        This is always a itself or a view into a. Note that if all axes are squeezed,
        the result is a 0d array and not a scalar.

    **Example**

    >>> x = torch.ones((2, 1, 3, 4, 1))
    >>> y = squeeze(x)
    >>> y.shape
    (2, 3, 4)
    """
    return TensorBox(tensor).squeeze(wrap_output=False)

def sum_(tensor, axis=None, keepdims=False):
    """TensorBox: Returns the sum of the tensor elements across the specified dimensions.

    Args:
        tensor (tensor_like): input tensor
        axis (int or tuple[int]): The axis or axes along which to perform the sum.
            If not specified, all elements of the tensor across all dimensions
            will be summed, returning a tensor.
        keepdims (bool): If True, retains all summed dimensions.

    Returns:
        tensor_like: The tensor with specified dimensions summed over. Note that
        if all elements are summed, then a 0-dimensional tensor is returned, rather
        than a Python scalar.

    **Example**

    Summing over all dimensions:

    >>> x = tf.Variable([[1., 2.], [3., 4.]])
    >>> sum(x)
    <tf.Tensor: shape=(), dtype=float32, numpy=10.0>

    Summing over specified dimensions:

    >>> x = np.array([[[1, 1], [5, 3]], [[1, 4], [-6, -1]]])
    >>> x.shape
    (2, 2, 2)
    >>> sum(x, axis=(0, 2))
    tensor([7, 1], requires_grad=True)
    >>> sum(x, axis=(0, 2), keepdims=True)
    tensor([[[7],
             [1]]], requires_grad=True)
    """
    return TensorBox(tensor).sum(axis=axis, keepdims=keepdims, wrap_output=False)

def T(tensor):
    """Returns the transpose of the tensor by reversing the order
    of the axes. For a 2D tensor, this corresponds to the matrix transpose.

    Args:
        tensor (tensor_like): input tensor

    Returns:
        tensor_like: input tensor with axes reversed

    **Example**

    >>> x = tf.Variable([[1, 2], [3, 4]])
    >>> T(x)
    <tf.Tensor: shape=(2, 2), dtype=int32, numpy=
    array([[1, 3],
           [2, 4]], dtype=int32)>
    """
    return TensorBox(tensor).T(wrap_output=False)

def take(tensor, indices, axis=None):
    """Gather elements from a tensor.

    Note that ``take(indices, axis=3)`` is equivalent
    to ``tensor[:, :, :, indices, ...]`` for frameworks that support
    NumPy-like fancy indexing.

    This function is roughly equivalent to ``np.take`` and ``tf.gather``.
    In the case of a 1-dimensional set of indices, it is roughly equivalent
    to ``torch.index_select``, but deviates for multi-dimensional indices.

    Args:
        tensor (tensor_like): input tensor
        indices (Sequence[int]): the indices of the values to extract
        axis: The axis over which to select the values. If not provided,
            the tensor is flattened before value extraction.

    **Example**

    >>> x = torch.tensor([[1, 2], [3, 4]])
    >>> take(y, indices=[[0, 0], [1, 0]], axis=1)
    tensor([[[1, 1],
             [2, 1]],

            [[3, 3],
             [4, 3]]])
    """
    return TensorBox(tensor).take(indices, axis=axis, wrap_output=False)

def where(condition, x, y):
    """Returns elements chosen from x or y depending on a boolean tensor condition.

    The input tensors ``condition``, ``x``, and ``y`` must all be broadcastable to the same shape.

    Args:
        condition (tensor_like[bool]): A boolean tensor. Where True, elements from
            ``x`` will be chosen, otherwise ``y``.
        x (tensor_like): values from which to choose if the condition evaluates to True
        y (tensor_like): values from which to choose if the condition evaluates to False

    Returns:
        tensor_like: A tensor with elements from ``x`` where the condition is True, and
        ``y`` otherwise. The output tensor has the same shape as the input tensors.

    **Example**

    >>> a = torch.tensor([0.6, 0.23, 0.7, 1.5, 1.7], requires_grad=True)
    >>> b = torch.tensor([-1., -2., -3., -4., -5.], requires_grad=True)
    >>> math.where(a < 1, a, b)
    tensor([ 0.6000,  0.2300,  0.7000, -4.0000, -5.0000], grad_fn=<SWhereBackward>)
    """
    return _get_multi_tensorbox([x, y]).where(condition, x, y, wrap_output=False)


# -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4

#
# LICENSE
#
# Copyright (C) 2010-2018 GEM Foundation, G. Weatherill, M. Pagani,
# D. Monelli.
#
# The Hazard Modeller's Toolkit is free software: you can redistribute
# it and/or modify it under the terms of the GNU Affero General Public
# License as published by the Free Software Foundation, either version
# 3 of the License, or (at your option) any later version.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake. If not, see <http://www.gnu.org/licenses/>
#
# DISCLAIMER
#
# The software Hazard Modeller's Toolkit (openquake.hmtk) provided herein
# is released as a prototype implementation on behalf of
# scientists and engineers working within the GEM Foundation (Global
# Earthquake Model).
#
# It is distributed for the purpose of open collaboration and in the
# hope that it will be useful to the scientific, engineering, disaster
# risk and software design communities.
#
# The software is NOT distributed as part of GEM's OpenQuake suite
# (https://www.globalquakemodel.org/tools-products) and must be considered as a
# separate entity. The software provided herein is designed and implemented
# by scientific staff. It is not developed to the design standards, nor
# subject to same level of critical review by professional software
# developers, as GEM's OpenQuake software suite.
#
# Feedback and contribution to the software is welcome, and can be
# directed to the hazard scientific staff of the GEM Model Facility
# (hazard@globalquakemodel.org).
#
# The Hazard Modeller's Toolkit (openquake.hmtk) is therefore distributed
# WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
# FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
# for more details.
#
# The GEM Foundation, and the authors of the software, assume no
# liability for use of the software.

#!/usr/bin/env python

'''
Module openquake.hmtk.plotting.seismicity.completeness.simple_completeness is a graphical
function for estimating the completeness period of magnitude intervals
by plotting the cumulative rate of events with time in each interval
'''

import numpy as np
import pylab
import matplotlib.pyplot as plt
from openquake.hmtk.seismicity.completeness.base import (
    BaseCatalogueCompleteness)

class SimpleCumulativeRate(BaseCatalogueCompleteness):
    '''
    Class to define the temporal variation in completess using simple
    changes in cumulative rates in individual completeness bins
    '''

    def completeness(self, catalogue, config, saveplot=False, filetype='png',
                     timeout=120):
        '''
        :param catalogue:
            Earthquake catalogue as instance of
            :class:`openquake.hmtk.seismicity.catalogue.Catalogue`
        :param dict config:
            Configuration parameters of the algorithm, containing the
            following information:
            'magnitude_bin' Size of magnitude bin (non-negative float)
            'time_bin' Size (in dec. years) of the time window (non-negative
            float)
            'increment_lock' Boolean to indicate whether to ensure
            completeness magnitudes always decrease with more
            recent bins
        :returns:
            2-column table indicating year of completeness and corresponding
            magnitude numpy.ndarray
        '''
        if saveplot and not isinstance(saveplot, str):
            raise ValueError('To save the figures enter a filename: ')

        # Get magntitude bins
        magnitude_bins = self._get_magnitudes_from_spacing(
            catalogue.data['magnitude'],
            config['magnitude_bin'])
        dec_time = catalogue.get_decimal_time()
        completeness_table = np.zeros([len(magnitude_bins) - 1, 2],
                                      dtype=float)
        min_year = float(np.min(catalogue.data['year']))
        max_year = float(np.max(catalogue.data['year'])) + 1.0
        has_completeness = np.zeros(len(magnitude_bins) - 1, dtype=bool)
        for iloc in range(0, len(magnitude_bins) - 1):
            lower_mag = magnitude_bins[iloc]
            upper_mag = magnitude_bins[iloc + 1]
            idx = np.logical_and(catalogue.data['magnitude'] >= lower_mag,
                                 catalogue.data['magnitude'] < upper_mag)
            cumvals = np.cumsum(np.ones(np.sum(idx)))
            plt.plot(dec_time[idx], cumvals, '.')
            plt.xlim(min_year, max_year + 5)
            title_string = 'Magnitude %5.2f to %5.2f' % (lower_mag, upper_mag)
            plt.title(title_string)
            pts = pylab.ginput(1, timeout=timeout)[0]
            if pts[0] <= max_year:
                 # Magnitude bin has no completeness!
                has_completeness[iloc] = True
            completeness_table[iloc, 0] = np.floor(pts[0])
            completeness_table[iloc, 1] = magnitude_bins[iloc]
            print(completeness_table[iloc, :], has_completeness[iloc])
            if config['increment_lock'] and (iloc > 0) and \
                    (completeness_table[iloc, 0] > completeness_table[iloc - 1, 0]):
                completeness_table[iloc, 0] = \
                    completeness_table[iloc - 1, 0]
            # Add marker line to indicate completeness point
            marker_line = np.array([
                [0., completeness_table[iloc, 0]],
                [cumvals[-1], completeness_table[iloc, 0]]])
            plt.plot(marker_line[:, 0], marker_line[:, 1], 'r-')
            if saveplot:
                filename = saveplot + '_' + ('%5.2f' % lower_mag) + (
                    '%5.2f' % upper_mag) + '.' + filetype
                plt.savefig(filename, format=filetype)
            plt.close()
        return completeness_table[has_completeness, :]

    def _get_magnitudes_from_spacing(self, magnitudes, delta_m):
        '''If a single magnitude spacing is input then create the bins

        :param numpy.ndarray magnitudes:
            Vector of earthquake magnitudes

        :param float delta_m:
            Magnitude bin width

        :returns: Vector of magnitude bin edges (numpy.ndarray)
        '''
        min_mag = np.min(magnitudes)
        max_mag = np.max(magnitudes)
        if (max_mag - min_mag) < delta_m:
            raise ValueError('Bin width greater than magnitude range!')
        mag_bins = np.arange(np.floor(min_mag), np.ceil(max_mag), delta_m)
        # Check to see if there are magnitudes in lower and upper bins
        is_mag = np.logical_and(mag_bins - max_mag < delta_m,
                                min_mag - mag_bins < delta_m)
        mag_bins = mag_bins[is_mag]
        return mag_bins


"""The WaveBlocks Project

This file contains the basic interface for general wavepackets.

@author: R. Bourquin
@copyright: Copyright (C) 2011 R. Bourquin
@license: Modified BSD License
"""

from numpy import vstack, vsplit, cumsum, zeros, complexfloating

class Wavepacket:
    r"""
    This class is primarily an abstract interface to wavepackets in general.
    But it implemets some methods for both the homogeneous and inhomogeneous Hagedorn
    wavepackets.
    """

    def __init__(self, parameters):
        r"""
        Initialize the ``Wavepacket`` object that represents :math:`|\Psi\rangle`.

        :raise NotImplementedError: Abstract interface.
        """
        raise NotImplementedError("'Wavepacket' is an abstract interface.")

    def __str__(self):
        r"""
        :return: A string describing the wavepacket.
        :raise NotImplementedError: Abstract interface.
        """
        raise NotImplementedError("'Wavepacket' is an abstract interface.")

    def _resize_coefficient_vector(self, component):
        r"""
        Adapt the coefficient vector for a given component to a new size.
        """
        oldsize = self.coefficients[component].shape[0]
        newsize = self.basis_size[component]

        if oldsize == newsize:
            return
        elif oldsize < newsize:
            # Append some zeros
            z = zeros((newsize - oldsize, 1), dtype=complexfloating)
            self.coefficients[component] = vstack([self.coefficients[component], z])
        elif oldsize > newsize:
            # Cut off the last part
            self.coefficients[component] = self.coefficients[component][0:newsize]

    def clone(self):
        raise NotImplementedError("'Wavepacket' is an abstract interface.")

    def gen_id(self):
        r"""
        Generate an (unique) ID per wavepacket instance.
        """
        #TODO: Better id generating function!
        self._id = id(self)

    def get_id(self):
        r"""
        Return the packet ID of this wavepacket instance. The ID may be used for storing packets in associative lists.
        """
        if not hasattr(self, "_id"):
            self.gen_id()

        return self._id

    def set_id(self, anid):
        r"""
        Manually set an ID for the current wavepacket instance.
        """
        assert(type(anid) is int)
        self._id = anid

    def get_number_components(self):
        r"""
        :return: The number :math:`N` of components the wavepacket :math:`\Psi` has.
        """
        return self.number_components

    def get_basis_size(self, component=None):
        r"""
        :return: The size of the basis, i.e. the number :math:`K` of :math:`{\phi_k}_{k=1}^K`.
        """
        if component is not None:
            return self.basis_size[component]
        else:
            return tuple(self.basis_size)

    def set_basis_size(self, basis_size, component=None):
        r"""
        Set the size of the basis of a given component or all components.

        :param basis_size: An single positive integer or a list of :math:`N` positive integers.
        :param component: The component for which we want to set the basis size.
                          Default is ``None`` which means 'all'.
        """
        if component is not None:
            # Check for valid input basis size
            if not component in range(self.number_components):
                raise ValueError("Invalid component index " + str(component))

            if basis_size < 2:
                raise ValueError("Basis size has to be a positive integer >=2.")

            # Set the new basis size for the given component
            self.basis_size[component] = basis_size
            # And adapt the coefficient vectors
            self._resize_coefficient_vector(component)

        else:
            # Check for valid input basis size
            if any([bs < 2 for bs in basis_size]):
                raise ValueError("Basis size has to be a positive integer >=2.")

            if not len(basis_size) == self.number_components:
                raise ValueError("Number of value(s) for basis size(s) does not match.")

            # Set the new basis size for all components
            self.basis_size = [ bs for bs in basis_size ]
            # And adapt the coefficient vectors
            for index in xrange(self.number_components):
                self._resize_coefficient_vector(index)

    def set_coefficients(self, values, component=None):
        r"""
        Update the coefficients :math:`c` of :math:`\Psi`.

        :param values: The new values of the coefficients :math:`c^i` of :math:`\Phi_i`.
        :param component: The index :math:`i` of the component we want to update with new coefficients.
        :raise ValueError: For invalid indices :math:`i`.

        .. note:: This function can either set new coefficients for a single component :math:`\Phi_i`
                  only if the ``component`` attribute is set or for all components simultaneously if
                  ``values`` is a list of arrays.
        """
        if component is None:
            for index, value in enumerate(values):
                if index > self.number_components-1:
                    raise ValueError("There is no component with index "+str(index)+".")

                self.coefficients[index] = value.copy().reshape((self.basis_size[index],1))
        else:
            if component > self.number_components-1:
                raise ValueError("There is no component with index "+str(component)+".")

            self.coefficients[component] = values.copy().reshape((self.basis_size[component],1))

    def set_coefficient(self, component, index, value):
        r"""
        Set a single coefficient :math:`c^i_k` of the specified component :math:`\Phi_i` of :math:`|\Psi\rangle`.

        :param component: The index :math:`i` of the component :math:`\Phi_i` we want to update.
        :param index: The index :math:`k` of the coefficient :math:`c^i_k` we want to update.
        :param value: The new value of the coefficient :math:`c^i_k`.
        :raise ValueError: For invalid indices :math:`i` or :math:`k`.
        """
        if component > self.number_components-1:
            raise ValueError("There is no component with index "+str(component)+".")
        if index > self.basis_size[component]-1:
            raise ValueError("There is no basis function with index "+str(index)+".")

        self.coefficients[component][index] = value

    def get_coefficients(self, component=None):
        r"""
        Returns the coefficients :math:`c^i` for some components :math:`\Phi_i` of :math:`|\Psi\rangle`.

        :param component: The index :math:`i` of the coefficients :math:`c^i` we want to get.
        :return: The coefficients :math:`c^i` either for all components :math:`\Phi_i`
                 or for a specified one.
        """
        if component is None:
            return [ item.copy() for item in self.coefficients ]
        else:
            return self.coefficients[component].copy()

    def get_coefficient_vector(self):
        r"""
        :return: The coefficients :math:`c^i` of all components :math:`\Phi_i` as a single long column vector.
        """
        return vstack(self.coefficients)

    def set_coefficient_vector(self, vector):
        r"""
        Set the coefficients for all components :math:`\Phi_i` simultaneously.

        :param vector: The coefficients of all components as a single long column vector.

        .. note:: This function does *NOT* copy the input data! This is for efficiency as this
                  routine is used in the innermost loops.
        """
        # Compute the partition of the block-vector from the basis sizes
        partition = cumsum(self.basis_size)[:-1]

        # Split the block-vector with the given partition and assign
        self.coefficients = vsplit(vector, partition)

    def get_parameters(self, component=None, aslist=False):
        r"""
        Get the Hagedorn parameters :math:`{\Pi}` of the wavepacket :math:`\Psi`.

        :raise NotImplementedError: Abstract interface.
        """
        raise NotImplementedError("'Wavepacket' is an abstract interface.")

    def set_parameters(self, parameters, component=None):
        r"""
        Set the Hagedorn parameters :math:`{\Pi}` of the wavepacket :math:`\Psi`.

        :raise NotImplementedError: Abstract interface.
        """
        raise NotImplementedError("'Wavepacket' is an abstract interface.")

    def evaluate_basis_at(self, nodes, component, prefactor=False):
        r"""
        Evaluate the basis functions :math:`\phi_k` recursively at the given nodes :math:`\gamma`.

        :raise NotImplementedError: Abstract interface.
        """
        raise NotImplementedError("'Wavepacket' is an abstract interface.")

    def evaluate_at(self, nodes, component=None, prefactor=False):
        r"""
        Evaluete the wavepacket :math:`\Psi` at the given nodes :math:`\gamma`.

        :raise NotImplementedError: Abstract interface.
        """
        raise NotImplementedError("'Wavepacket' is an abstract interface.")

    def get_norm(self, component=None, summed=False):
        r"""
        Calculate the :math:`L^2` norm of the wavepacket :math:`|\Psi\rangle`.

        :raise NotImplementedError: Abstract interface.
        """
        raise NotImplementedError("'Wavepacket' is an abstract interface.")

    def potential_energy(self, potential, summed=False):
        r"""
        Calculate the potential energy :math:`\langle\Psi|V|\Psi\rangle ` of the wavepacket componentwise.

        :raise NotImplementedError: Abstract interface.
        """
        raise NotImplementedError("'Wavepacket' is an abstract interface.")

    def kinetic_energy(self, summed=False):
        r"""
        Calculate the kinetic energy :math:`\langle\Psi|T|\Psi\rangle ` of the wavepacket componentwise.

        :raise NotImplementedError: Abstract interface.
        """
        raise NotImplementedError("'Wavepacket' is an abstract interface.")

    def project_to_canonical(self, potential):
        r"""
        Project the wavepacket to the canonical basis.

        :raise NotImplementedError: Abstract interface.
        """
        raise NotImplementedError("'Wavepacket' is an abstract interface.")

    def project_to_eigen(self, potential):
        r"""
        Project the wavepacket to the eigenbasis of a given potential :math:`V`.

        :raise NotImplementedError: Abstract interface.
        """
        raise NotImplementedError("'Wavepacket' is an abstract interface.")

    def to_fourier_space(self, assign=True):
        r"""
        Transform the wavepacket to Fourier space.

        :raise NotImplementedError: Abstract interface.
        """
        raise NotImplementedError("'Wavepacket' is an abstract interface.")

    def to_real_space(self, assign=True):
        r"""
        Transform the wavepacket to real space.

        :raise NotImplementedError: Abstract interface.
        """
        raise NotImplementedError("'Wavepacket' is an abstract interface.")


"""
Everything to store electromagnetic material properties for the solver.
"""
# Copyright 2018-2021 The emsig community.
#
# This file is part of emg3d.
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not
# use this file except in compliance with the License.  You may obtain a copy
# of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  See the
# License for the specific language governing permissions and limitations under
# the License.

from copy import deepcopy

import numpy as np
from scipy.constants import epsilon_0

from emg3d import maps, meshes, utils

__all__ = ['Model', 'VolumeModel', 'expand_grid_model']

# MODEL
@utils._known_class
class Model:
    r"""A model containing the electromagnetic properties of the Earth.

    A model provides the required model parameters to the solver. The x-, y-,
    and z-directed electrical properties are by default resistivities. However,
    they can be defined as resistivities, :math:`\rho\ (\Omega\,\mathrm{m})`,
    or conductivities, :math:`\sigma\ (\mathrm{S/m})`, either on a linear or on
    a logarithmic scale (log or ln), by choosing the appropriate ``mapping``.
    Relative magnetic permeability :math:`\mu_\mathrm{r}` is by default set to
    one and electric permittivity :math:`\varepsilon_\mathrm{r}` is by default
    set to zero, but they can also be provided (isotropically). Keep in mind
    that the multigrid method as implemented in emg3d works for the diffusive
    approximation. When the displacement part in Maxwell's equations becomes
    too dominant it will start to fail (high frequencies or very high electric
    permittivity).

    Parameters
    ----------
    grid : TensorMesh
        The grid; a :class:`emg3d.meshes.TensorMesh` instance.

    property_{x;y;z} : {None, array_like}, default: 1 (x), None (y, z)
        Electrical material property in x-, y-, and z-directions. The
        properties are stored as Fortran-ordered arrays with the shape given by
        ``grid.shape``. The provided value must be broadcastable to that shape.

        By default, property refers to electrical resistivity. However, this
        can be changed with an appropriate ``mapping`` (the internals of
        emg3d work, irrelevant of the map, with electrical conductivities).

        The properties have to be of finite value, bigger than zero (on linear
        scale). The four supported anisotropy cases are:

        - ``x;y=None;z=None``: isotropic (``y=z=x``);
        - ``x;y=None;z``: vertical transverse isotropy VTI (``y=x``);
        - ``x;y;z=None``: horizontal transverse isotropy HTI (``z=x``);
        - ``x;y;z``: triaxial  anisotropy.

        If a property is not initiated it cannot be set later on (e.g., if a
        VTI model is created, it is not possible to set ``property_y`` later
        on, instead, a new model has to be initiated).

    mu_r, epsilon_r : {None, array_like}, default: None
        Relative magnetic permeability (-) and relative electric permittivity
        (-), respectively, both isotropic. The properties are stored as
        Fortran-ordered arrays with the shape given by ``grid.shape``. The
        provided value must be broadcastable to that shape.

        The properties have to be of finite value, bigger than zero.

        The relative magnetic permeability is assumed to be 1 if not provided.

        The relative electric permittivity is assumed to be 0 if not provided,
        which ignores the displacement part completely (diffusive
        approximation)

    mapping : str, default: 'Resistivity'
        Defines what type the electrical input ``property_{x;y;z}``-values
        correspond to. The implemented types are:

        - ``'Resistivity'``; ρ (Ω m);
        - ``'Conductivity'``; σ (S/m);
        - ``'LgResistivity'``; log_10(ρ);
        - ``'LgConductivity'``; log_10(σ);
        - ``'LnResistivity'``; log_e(ρ);
        - ``'LnConductivity'``; log_e(σ).

    """

    def __init__(self, grid, property_x=1., property_y=None, property_z=None,
                 mu_r=None, epsilon_r=None, mapping='Resistivity'):
        """Initiate a new model."""

        # Store grid.
        self.grid = grid

        # Alias shape_cells and n_cells to shape and size.
        self.shape = self.grid.shape_cells
        self.size = self.grid.n_cells

        # Get and store map.
        self.map = getattr(maps, 'Map'+mapping)()

        # Initiate and store all parameters.
        self._property_x = self._init_parameter(property_x, 'property_x')
        self._property_y = self._init_parameter(property_y, 'property_y')
        self._property_z = self._init_parameter(property_z, 'property_z')
        self._mu_r = self._init_parameter(mu_r, 'mu_r',)
        self._epsilon_r = self._init_parameter(epsilon_r, 'epsilon_r')
        self._properties = ['property_x', 'property_y', 'property_z',
                            'mu_r', 'epsilon_r']

        # Store case.
        if self._property_y is None and self._property_z is None:
            self.case = 'isotropic'
        elif self._property_z is None:
            self.case = 'HTI'
        elif self._property_y is None:
            self.case = 'VTI'
        else:
            self.case = 'triaxial'

    def __repr__(self):
        """Simple representation."""
        return (f"{self.__class__.__name__}: {self.map.description}; "
                f"{self.case}{'' if self.mu_r is None else '; mu_r'}"
                f"{'' if self.epsilon_r is None else '; epsilon_r'}"
                f"; {self.shape[0]} x {self.shape[1]} x {self.shape[2]} "
                f"({self.size:,})")

    def __add__(self, model):
        """Add two models."""

        # Ensure model is a Model instance.
        if model.__class__.__name__ != 'Model':
            return NotImplemented

        # Check input.
        self._operator_test(model)

        # Apply operator.
        kwargs = self._apply_operator(model, np.add)

        # Return new Model instance.
        return Model(grid=self.grid, mapping=self.map.name, **kwargs)

    def __sub__(self, model):
        """Subtract two models."""

        # Ensure model is a Model instance.
        if model.__class__.__name__ != 'Model':
            return NotImplemented

        # Check input.
        self._operator_test(model)

        # Apply operator.
        kwargs = self._apply_operator(model, np.subtract)

        # Return new Model instance.
        return Model(grid=self.grid, mapping=self.map.name, **kwargs)

    def __eq__(self, model):
        """Compare two models."""

        # Check if model is a Model instance.
        equal = model.__class__.__name__ == 'Model'

        # Check input.
        if equal:
            try:
                self._operator_test(model)
            except ValueError:
                equal = False

        # Compare values if not None.
        if equal:
            for prop in self._properties:
                val = getattr(self, prop)
                if val is not None:
                    equal *= np.allclose(val, getattr(model, prop))

        return bool(equal)

    def copy(self):
        """Return a copy of the Model."""
        return self.from_dict(self.to_dict(True))

    def to_dict(self, copy=False):
        """Store the necessary information in a dict for serialization.

        Parameters
        ----------
        copy : bool, default: False
            If True, returns a deep copy of the dict.

        Returns
        -------
        out : dict
            Dictionary containing all information to re-create the Model.

        """
        out = {
            '__class__': self.__class__.__name__,  # v ensure emg3d-TensorMesh
            'grid': meshes.TensorMesh(self.grid.h, self.grid.origin).to_dict(),
            **{prop: getattr(self, prop) for prop in self._properties},
            'mapping': self.map.name,
        }
        if copy:
            return deepcopy(out)
        else:
            return out

    @classmethod
    def from_dict(cls, inp):
        """Convert dictionary into :class:`emg3d.models.Model` instance.

        Parameters
        ----------
        inp : dict
            Dictionary as obtained from :func:`emg3d.models.Model.to_dict`. The
            dictionary needs the keys ``property_x``, ``property_y``,
            ``property_z``, ``mu_r``, ``epsilon_r``, ``grid``, and ``mapping``;
            ``grid`` itself is also a dict which needs the keys ``hx``, ``hy``,
            ``hz``, and ``origin``.

        Returns
        -------
        model : Model
            A :class:`emg3d.models.Model` instance.

        """
        inp = {k: v for k, v in inp.items() if k != '__class__'}
        MeshClass = getattr(meshes, inp['grid']['__class__'])
        return cls(grid=MeshClass.from_dict(inp.pop('grid')), **inp)

    # ELECTRICAL PROPERTIES
    @property
    def property_x(self):
        r"""Electrical property in x-direction."""
        return self._property_x

    @property_x.setter
    def property_x(self, property_x):
        r"""Update electrical property in x-direction."""
        self._check_positive_finite(property_x, 'property_x')
        self._property_x[:] = np.asfortranarray(property_x, dtype=np.float64)

    @property
    def property_y(self):
        r"""Electrical property in y-direction."""
        return self._property_y

    @property_y.setter
    def property_y(self, property_y):
        r"""Update electrical property in y-direction."""
        self._check_positive_finite(property_y, 'property_y')
        self._property_y[:] = np.asfortranarray(property_y, dtype=np.float64)

    @property
    def property_z(self):
        r"""Electrical property in z-direction."""
        return self._property_z

    @property_z.setter
    def property_z(self, property_z):
        r"""Update electrical property in z-direction."""
        self._check_positive_finite(property_z, 'property_z')
        self._property_z[:] = np.asfortranarray(property_z, dtype=np.float64)

    @property
    def mu_r(self):
        r"""Relative magnetic permeability."""
        return self._mu_r

    @mu_r.setter
    def mu_r(self, mu_r):
        r"""Update relative magnetic permeability."""
        self._check_positive_finite(mu_r, 'mu_r')
        self._mu_r[:] = np.asfortranarray(mu_r, dtype=np.float64)

    @property
    def epsilon_r(self):
        r"""Relative electric permittivity."""
        return self._epsilon_r

    @epsilon_r.setter
    def epsilon_r(self, epsilon_r):
        r"""Update relative electric permittivity."""
        self._check_positive_finite(epsilon_r, 'epsilon_r')
        self._epsilon_r[:] = np.asfortranarray(epsilon_r, dtype=np.float64)

    # INTERPOLATION
    def interpolate_to_grid(self, grid, **interpolate_opts):
        """Interpolate the model to a new grid.

        If the provided grid is identical to the grid of the model, it returns
        the actual model (not a copy).

        Parameters
        ----------
        grid : TensorMesh
            Grid of the new model; a :class:`emg3d.meshes.TensorMesh` instance.

        interpolate_opts : dict
            Passed through to :func:`emg3d.maps.interpolate`. Defaults are
            ``method='volume'``, ``log=True``, and ``extrapolate=True``.

        Returns
        -------
        obj : Model
            A new :class:`emg3d.models.Model` instance on ``grid``.

        """
        # If grids are identical, return model.
        if grid == self.grid:
            return self

        # Get solver options, set to defaults if not provided.
        g2g_inp = {
            'method': 'volume',
            'extrapolate': True,
            'log': not self.map.name.startswith('L'),
            **({} if interpolate_opts is None else interpolate_opts),
            'grid': self.grid,
            'xi': grid,
        }

        # Interpolate property_{x;y;z}; mu_r; and epsilon_r; add to dict.
        model_inp = {}
        for prop in self._properties:
            var = getattr(self, prop)
            if var is None:
                model_inp[prop] = None
            else:
                model_inp[prop] = maps.interpolate(values=var, **g2g_inp)

        # Assemble new model.
        return Model(grid, mapping=self.map.name, **model_inp)

    # INTERNAL UTILITIES
    def _init_parameter(self, values, name):
        """Initiate parameter by casting and broadcasting."""

        # If None, exit.
        if values is None:
            return None

        # Cast it to an array of floats, in Fortran order.
        values = np.asfortranarray(values, dtype=np.float64)

        # If 1D array of self.size, reshape it.
        if values.size == self.size:
            values = values.reshape(self.shape, order='F')

        # If not of shape self.shape, broadcast it.
        elif values.shape != self.shape:
            values = np.ones(self.shape, order='F')*values

        # Check >0 and finite.
        self._check_positive_finite(values, name)

        return values

    def _check_positive_finite(self, values, name):
        """Check parameter values are positive (on linear scale) and finite."""

        # If it is None, it cannot be set.
        if hasattr(self, '_'+name) and getattr(self, '_'+name) is None:
            raise ValueError(
                f"Model was initiated without `{name}`; cannot set values."
            )

        # Get mapped values; checks are carried out on conductivities.
        if 'property_' in name:
            mapped = self.map.backward(np.asarray(values))
        else:
            mapped = values

        # Check they are positive.
        if not np.all(np.real(mapped) > 0.0):
            raise ValueError(f"`{name}` must be all bigger than zero.")

        # Check |val| < inf.
        if not np.all(np.isfinite(mapped)):
            raise ValueError(f"`{name}` must be all finite.")

    def _operator_test(self, model):
        """Check if ``self`` and ``model`` are consistent for operations."""

        # Ensure the two instances have the same grid.
        if self.grid != model.grid:
            raise ValueError("Models have different grids.")

        # Ensure the two instances have the same case.
        if self.case != model.case:
            raise ValueError("Models have different anisotropy.")

        # Ensure both or none has mu_r:
        if (self.mu_r is None) != (model.mu_r is None):
            raise ValueError("One model has mu_r, the other not.")

        # Ensure both or none has epsilon_r:
        if (self.epsilon_r is None) != (model.epsilon_r is None):
            raise ValueError("One model has epsilon_r, the other not.")

        # Ensure the two instances have the same mapping:
        if self.map.name != model.map.name:
            raise ValueError("Models have different mappings.")

    def _apply_operator(self, model, operator):
        """Apply the provided operator to self and model."""

        # Apply operator to property_{x;y;z}; mu_r; and epsilon_r; add to dict.
        kwargs = {}
        for prop in self._properties:
            val = getattr(self, prop)
            if val is None:
                kwargs[prop] = None
            else:
                kwargs[prop] = operator(val, getattr(model, prop))

        return kwargs

class VolumeModel:
    r"""Return simplified model with volume-averaged eta_{x;y;z}; zeta.

    Takes a model and a source field and returns the volume-averaged eta and
    zeta values. This is used internally by the solver.

    .. math::

        \eta_{\{x,y,z\}} = -V\mathrm{i}\omega\mu_0
              \left(\sigma_{\{x,y,z\}} + \mathrm{i}\omega\varepsilon\right)

    .. math::

        \zeta = V\mu_\mathrm{r}^{-1}

    Parameters
    ----------
    model : Model
        Model to transform to volume-averaged values.

    sfield : Field
       A VolumeModel is frequency-dependent. The frequency-information is taken
       from the provided source field.

    """

    def __init__(self, model, sfield):
        """Initiate a new model with volume-averaged properties."""

        # Store case and minimal TensorMesh.
        self.case = model.case
        self.grid = meshes.BaseMesh(model.grid.h, model.grid.origin)

        # Get volume for volume-averaged values.
        vol = self.grid.cell_volumes.reshape(model.shape, order='F')

        # Compute and store eta.
        for name in model._properties[:3]:

            prop = getattr(model, name)

            if prop is None:
                eta = None

            else:
                # Get conductivities.
                cond = model.map.backward(prop)

                # Diffusive approximation.
                if model.epsilon_r is None:
                    eta = -sfield.smu0*vol*cond

                # Complete version.
                else:
                    eta = -sfield.smu0*vol*(
                            cond + sfield.sval*epsilon_0*model.epsilon_r)

            setattr(self, '_eta_' + name[-1], eta)

        # Compute and store zeta.
        zeta = vol
        if model.mu_r is not None:
            zeta /= model.mu_r
        self._zeta = zeta

    @property
    def eta_x(self):
        r"""Volume-averaged eta in x-direction."""
        return self._eta_x

    @property
    def eta_y(self):
        r"""Volume-averaged eta in y-direction."""
        if self.case in ['HTI', 'triaxial']:
            return self._eta_y
        else:
            return self._eta_x

    @property
    def eta_z(self):
        r"""Volume-averaged eta in z-direction."""
        if self.case in ['VTI', 'triaxial']:
            return self._eta_z
        else:
            return self._eta_x

    @property
    def zeta(self):
        r"""Volume-averaged, isotropic zeta."""
        return self._zeta

def expand_grid_model(model, expand, interface):
    """Expand model and grid according to provided parameters.

    Expand the grid and corresponding model in positive z-direction from the
    edge of the grid to the interface with property ``expand[0]``, and a 100 m
    thick layer above the interface with property ``expand[1]``.

    The provided properties are taken as isotropic (as is the case in water and
    air); ``mu_r`` and ``epsilon_r`` are expanded with ones, if necessary.

    The ``interface`` is usually the sea-surface, and ``expand`` is therefore
    ``[property_sea, property_air]``.

    Parameters
    ----------
    model : Model
        The model; a :class:`emg3d.models.Model` instance.

    expand : list
        The two properties below and above the interface:
        ``[below_interface, above_interface]``.

    interface : float
        Interface between the two properties in ``expand``.

    Returns
    -------
    exp_grid : TensorMesh
        Expanded grid; a :class:`emg3d.meshes.TensorMesh` instance.

    exp_model : Model
        The expanded model; a :class:`emg3d.models.Model` instance.

    """
    grid = model.grid

    def extend_property(prop, add_values, nadd):
        """Expand property `model.prop`, IF it is not None."""

        if getattr(model, prop) is None:
            prop_ext = None

        else:
            prop_ext = np.zeros((grid.shape_cells[0], grid.shape_cells[1],
                                 grid.shape_cells[2]+nadd))
            prop_ext[:, :, :-nadd] = getattr(model, prop)
            if nadd == 2:
                prop_ext[:, :, -2] = add_values[0]
            prop_ext[:, :, -1] = add_values[1]

        return prop_ext

    # Initiate.
    nzadd = 0
    hz_ext = grid.h[2]

    # Fill-up property_below.
    if grid.nodes_z[-1] < interface-0.05:  # At least 5 cm.
        hz_ext = np.r_[hz_ext, interface-grid.nodes_z[-1]]
        nzadd += 1

    # Add 100 m of property_above.
    if grid.nodes_z[-1] <= interface+0.001:  # +1mm
        hz_ext = np.r_[hz_ext, 100]
        nzadd += 1

    if nzadd > 0:
        # Extend properties.
        property_x = extend_property('property_x', expand, nzadd)
        property_y = extend_property('property_y', expand, nzadd)
        property_z = extend_property('property_z', expand, nzadd)
        mu_r = extend_property('mu_r', [1, 1], nzadd)
        epsilon_r = extend_property('epsilon_r', [1, 1], nzadd)

        # Create extended grid and model.
        grid = meshes.TensorMesh(
                [grid.h[0], grid.h[1], hz_ext], origin=grid.origin)
        model = Model(grid, property_x, property_y, property_z, mu_r,
                      epsilon_r, mapping=model.map.name)

    return model