"""
This package offers a generalisation of linear convolutions and max-poolings using
semifields and for an arbitrary number of spatial dimensions,
with many common choices for semifields being available.

Additionally, this package provides example implementations of 2D isotropic and
anisotropic quadratic kernels for use in e.g. dilation of 2-dimensional images.
These implementations are efficient when run under `torch.compile`.

As an example of anisotropic quadratic dilation, consider replacing:

    pooling = torch.nn.MaxPool2d(3, stride=2, padding=1)

with:

    import pytorch_nd_semiconv as semiconv

    dilation = semiconv.GenericConv(
        semiconv.QuadraticKernelSpectral2D(
            in_channels=5, out_channels=5, kernel_size=3
        ),
        semiconv.SelectSemifield.tropical_max().lazy_fixed(),
        padding="same",
        stride=2,
        groups=5,
    )
"""

from ._conv_modules import GenericClosing, GenericConv
from ._quadratic import (
    CovCholesky2D,
    CovSpectral2D,
    QuadraticKernelCholesky2D,
    QuadraticKernelIso2D,
    QuadraticKernelSpectral2D,
)
from ._sf_broadcast import BroadcastSemifield
from ._sf_select import SelectSemifield
from ._sf_subtract import SubtractSemifield
from ._utils import LearnedKernel2D, TorchLinearConv2D, TorchMaxPool2D

__all__ = [  # noqa: RUF022
    "GenericConv",
    "BroadcastSemifield",
    "SelectSemifield",
    "SubtractSemifield",
    "QuadraticKernelSpectral2D",
    "QuadraticKernelCholesky2D",
    "QuadraticKernelIso2D",
    "GenericClosing",
    "LearnedKernel2D",
    "TorchLinearConv2D",
    "TorchMaxPool2D",
    # These two are exported but not documented: the internal covariance matrices
    "CovCholesky2D",
    "CovSpectral2D",
]


# Modify autogenerated docs
__pdoc__ = {
    "QuadraticKernelSpectral2D.forward": False,
    "QuadraticKernelCholesky2D.forward": False,
    "QuadraticKernelIso2D.forward": False,
    "LearnedKernel2D.forward": False,
    "TorchLinearConv2D.forward": False,
    "TorchMaxPool2D.forward": False,
    "GenericClosing.forward": False,
    "QuadraticKernelIso2D.log_std": False,
}

for _tup in (BroadcastSemifield, SelectSemifield, SubtractSemifield):
    for _field in _tup._fields:
        __pdoc__[f"{_tup.__name__}.{_field}"] = False

for _export in __all__:
    _mod = eval(_export)
    __pdoc__[_export] = _mod.__doc__
    if "extra_repr" in _mod.__dict__:
        __pdoc__[f"{_export}.extra_repr"] = False

__pdoc__["CovCholesky2D"] = False
__pdoc__["CovSpectral2D"] = False
