"""Extended functionality for [pyMOR](https://pymor.org/) grids"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/api/pymor/grids.ipynb.

# %% auto 0
__all__ = ['OnedUnstructuredGrid']

# %% ../../nbs/api/pymor/grids.ipynb
from pymor.discretizers.builtin.grids.oned import OnedGrid
import numpy as np

# %% ../../nbs/api/pymor/grids.ipynb
class OnedUnstructuredGrid(OnedGrid):
    """One-dimensional |Grid| on an interval.

    Parameters
    ----------
    edges
        Sequence specifying the edges of the grid intervals.
    identify_left_right
        Whether the left and right edges of the grid should be identified.
    """

    def __init__(self, edges, identify_left_right=False):
        self.edges = np.array(edges)
        self.domain = (self.edges[0], self.edges[-1])
        assert self.edges.ndim == 1
        assert np.all(self.edges[:-1] <= self.edges[1:])
        self.__auto_init(locals())
        num_intervals = len(self.edges) - 1
        self._sizes = [num_intervals, num_intervals] if identify_left_right else [num_intervals, num_intervals + 1]
        self._widths = np.diff(self.edges)
        self._subentities = np.vstack((np.arange(num_intervals, dtype=np.int32),
                                        np.arange(num_intervals, dtype=np.int32) + 1))
        if identify_left_right:
            self._subentities[-1, -1] = 0
        self._A = self._widths[:, np.newaxis, np.newaxis]
        self._B = np.cumsum(np.pad(self._widths[:-1], (1, 0)))[:, np.newaxis] + self.domain[0]

    # The following methods would not need to be redefined from the parent class, 
    # except that name-mangling (double underscore) is used in the parent.
    # Ask about that and/or submit PR?
    
    def subentities(self, codim, subentity_codim):
        assert 0 <= codim <= 1, 'Invalid codimension'
        assert codim <= subentity_codim <= self.dim, 'Invalid subentity codimension'
        if codim == 0:
            if subentity_codim == 0:
                return np.arange(self.size(0), dtype='int32')[:, np.newaxis]
            else:
                return self._subentities.T
        else:
            return super().subentities(codim, subentity_codim)

    def embeddings(self, codim):
        if codim == 0:
            return self._A, self._B
        else:
            return super().embeddings(codim)
