# Source code for pylops.utils.tapers

import numpy as np

r"""1D Hanning taper

Create unitary mask of length nmask with Hanning tapering
at edges of size ntap

Parameters
----------
nmask : :obj:int
ntap : :obj:int
Number of samples of hanning tapering at edges

Returns
-------
taper : :obj:numpy.ndarray
taper

"""
if ntap > 0:
if (nmask // ntap) < 2:
ntap_min = nmask / 2 if nmask % 2 == 0 else (nmask - 1) / 2
raise ValueError(
"ntap=%d must be smaller or " "equal than %d" % (ntap, ntap_min)
)
han_win = np.hanning(ntap * 2 - 1)
st_tpr = han_win[
:ntap,
]
mid_tpr = np.ones(
[
]
)
end_tpr = np.flipud(st_tpr)
tpr_1d = np.concatenate([st_tpr, mid_tpr, end_tpr])
return tpr_1d

r"""1D Cosine or Cosine square taper

Create unitary mask of length nmask with Hanning tapering
at edges of size ntap

Parameters
----------
nmask : :obj:int
ntap : :obj:int
Number of samples of hanning tapering at edges
square : :obj:bool
Cosine square taper (True)or Cosine taper (False)

Returns
-------
taper : :obj:numpy.ndarray
taper

"""
exponent = 1 if not square else 2
cos_win = (
0.5
* (
np.cos(
(np.arange(ntap * 2 - 1) - (ntap * 2 - 2) / 2)
* np.pi
/ ((ntap * 2 - 2) / 2)
)
+ 1.0
)
) ** exponent
st_tpr = cos_win[
:ntap,
]
mid_tpr = np.ones(
[
]
)
end_tpr = np.flipud(st_tpr)
tpr_1d = np.concatenate([st_tpr, mid_tpr, end_tpr])
return tpr_1d

r"""1D taper

Create unitary mask of length nmask with tapering of choice
at edges of size ntap

Parameters
----------
nmask : :obj:int
ntap : :obj:int
Number of samples of hanning tapering at edges
tapertype : :obj:str, optional
Type of taper (hanning, cosine, cosinesquare or None)

Returns
-------
taper : :obj:numpy.ndarray
taper

"""
if tapertype == "hanning":
elif tapertype == "cosine":
elif tapertype == "cosinesquare":
else:
return tpr_1d

r"""2D taper

Create 2d mask of size :math:[n_\text{mask} \times n_t]
with tapering of size ntap along the first (and possibly
second) dimensions

Parameters
----------
nt : :obj:int
Number of samples along second dimension
nmask : :obj:int
Number of samples along first dimension
ntap : :obj:int or :obj:list
Number of samples of tapering at edges of first dimension (or
both dimensions).
tapertype : :obj:str, optional
Type of taper (hanning, cosine, cosinesquare or None)

Returns
-------
taper : :obj:numpy.ndarray
2d mask with tapering along first dimension
of size :math:[n_\text{mask} \times n_t]

"""
# create 1d window along first dimension
tpr_x = taper(
nmask, ntap[0] if isinstance(ntap, (list, tuple)) else ntap, tapertype
)

# create 1d window along second dimension
if isinstance(ntap, (list, tuple)):
tpr_t = taper(nt, ntap[1], tapertype)

# create 2d taper
if isinstance(ntap, (list, tuple)):
# replicate taper to second dimension
tpr_2d = np.outer(tpr_x, tpr_t)
else:
# replicate taper to second dimension
tpr_2d = np.tile(tpr_x[:, np.newaxis], (1, nt))

return tpr_2d

r"""3D taper

Create 2d mask of size :math:[n_\text{mask}[0] \times n_\text{mask}[1] \times n_t]
with tapering of size ntap along the first and second dimension

Parameters
----------
nt : :obj:int
Number of time samples of mask along third dimension
nmask : :obj:tuple
Number of space samples of mask along first dimension
ntap : :obj:tuple
Number of samples of tapering at edges of first dimension
tapertype : :obj:int
Type of taper (hanning, cosine,
cosinesquare or None)

Returns
-------
taper : :obj:numpy.ndarray
2d mask with tapering along first dimension
of size :math:[n_\text{mask,0} \times n_\text{mask,1} \times n_t]

"""
ntapy, ntapx = ntap[0], ntap[1]

# create 1d window
if tapertype == "hanning":
elif tapertype == "cosine":
elif tapertype == "cosinesquare":