import numpy as np

from pylops import LinearOperator

Zero-pad model in forward model and extract non-zero subsequence
in adjoint. Padding can be performed in one or multiple directions to any
multi-dimensional input arrays.

Parameters
----------
dims : :obj:int or :obj:tuple
Number of samples for each dimension
pad : :obj:tuple
Number of samples to pad. If dims is a scalar, pad is a single
tuple (pad_in, pad_end). If dims is a tuple,
pad is a tuple of tuples where each inner tuple contains
the number of samples to pad in each dimension
dtype : :obj:str, optional
Type of elements in input array.

Attributes
----------
shape : :obj:tuple
Operator shape
explicit : :obj:bool
Operator contains a matrix that can be solved explicitly (True) or
not (False)

Raises
------
ValueError
If any element of pad is negative.

Notes
-----
Given an array of size :math:N, the *Pad* operator simply adds
:math:\text{pad}_\text{in} at the start and :math:\text{pad}_\text{end} at the end in forward mode:

.. math::

and :math:y_i = 0 \quad \forall
i=0,\ldots,\text{pad}_\text{in}-1, \text{pad}_\text{in}+N-1,\ldots,N+\text{pad}_\text{in}+\text{pad}_\text{end}

In adjoint mode, values from :math:\text{pad}_\text{in} to :math:N-\text{pad}_\text{end} are
extracted from the data:

.. math::

"""

raise ValueError("Padding must be positive or zero")
self.dims = dims
self.reshape = False if isinstance(self.dims, int) else True
if self.reshape:
self.dimsd = [dim + p[0] + p[1] for dim, p in zip(dims, pad)]
else:
self.shape = (np.prod(np.array(self.dimsd)), np.prod(np.array(self.dims)))
self.dtype = np.dtype(dtype)
self.explicit = False

def _matvec(self, x):
if self.reshape:
y = x.reshape(self.dims)