"""Create and manipulate a list of bins of the atomic velocity distribution for use in the LGS model"""

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

# %% auto 0
__all__ = ['VelocityGroups', 'velocity_density_vector']

# %% ../nbs/api/velocitygroups.ipynb
from .imports import *
from .utilities.nbdev import DictTbl, AttributeTbl
from .utilities.testing import test_array
from .utilities.sparse import sparse_kronecker_matrix, sparse_toeplitz, sparse_identity, sparse_diag, sparse
from .utilities.numpy import sym_range
from .pymor.parameters import *
from .pymor.vectorarrays import *
from .pymor.operators import *
from .pymor.grids import *
from pymor.vectorarrays.interface import VectorArray
from pymor.operators.interface import Operator
from pymor.basic import IdentityOperator

# %% ../nbs/api/velocitygroups.ipynb
class VelocityGroups(dict):
    """`VelocityGroups` contains informations about the centers, edges, and relative densities of a list of velocity group bins."""
    def __init__(self, 
                 bins:Iterable|int=6, # Iterable of velocity group edges or int specifying number of evenly spaced velocity groups
                 range:Number|Iterable=3 # Default range (`min_vel`, `max_vel`) or half-width `max_vel` for velocity groups
                ):
        super().__init__()
        if isinstance(bins, VelocityGroups):
            self.edges = bins.edges
            for k, v in bins.items():
                self[k] = v
            return
        if isinstance(range, Number): range = (-range, range)
        self.edges = np.linspace(*range, bins + 1) if isinstance(bins, int) else bins
        self['VGCenter'] = np.mean(np.array(list(zip(self.edges[:-1], self.edges[1:]))), axis=1)
        # self['VGBins'] = np.array(list(zip(bins[:-1], bins[1:])))
        # self['VGCenter'] = np.mean(self['VGBins'], axis=1)
        self['VGDensity'] = np.diff(erf(self.edges))/2
        self['VGWidth'] = np.diff(self.edges)
        # self['VGEdges'] = bins
        self['VGInverseWidth'] = 1/self['VGWidth']
        self['VGNumber'] = len(self['VGCenter'])
        self['velocity_groups'] = np.ones(self['VGDensity'].shape)

# %% ../nbs/api/velocitygroups.ipynb
@patch
def subdivide(
    self:VelocityGroups, 
    indices:int|ndarray|slice=slice(None), # Indices, boolean array, or slice object specifying velocity groups to subdivide
    n:int=2 # Number of groups to subdivide into
):
    """Divide the velocity groups with indices `indices` in two."""
    new_edges = np.concat([np.linspace(c - w/2, c + w/2, n + 1)[1:-1] for c,w in zip(self["VGCenter"][indices], self["VGWidth"][indices])])
    return VelocityGroups(np.union1d(self.edges, new_edges))

# %% ../nbs/api/velocitygroups.ipynb
def _velocity_space(vg, ext=''):
    return XarrayVectorSpace({'Atomic velocity' + ext: vg['VGCenter']})

# %% ../nbs/api/velocitygroups.ipynb
def _vg_identity(vg):
    return XarrayMatrixOperator(
        DataArray(sparse_identity(vg['VGNumber']), coords=[('Atomic velocity' + Lbl.RNG, vg['VGCenter']), ('Atomic velocity', vg['VGCenter'])])
    )

# %% ../nbs/api/velocitygroups.ipynb
@patch
def identity(self:VelocityGroups)->Operator:
    """The velocity-space identity operator."""
    return IdentityOperator(_velocity_space(self))
    # return XarrayMatrixOperator(
    #         sparse_identity(self['VGNumber']), 
    #         range=_velocity_space(self), 
    #         source=_velocity_space(self)
    #     )

# %% ../nbs/api/velocitygroups.ipynb
def _vg_diagonal(vg, diags):
    return XarrayMatrixOperator(
        DataArray(sparse_diag(diags), coords=[('Atomic velocity' + Lbl.RNG, vg['VGCenter']), ('Atomic velocity' + Lbl.SRC, vg['VGCenter'])])
    )

# %% ../nbs/api/velocitygroups.ipynb
@patch
def velocity_diagonal(self:VelocityGroups)->Operator:
    """A velocity-space operator with the velocity on the diagonal."""
    return _vg_diagonal(self, self['VGCenter'])

# %% ../nbs/api/velocitygroups.ipynb
def velocity_density_vector(vg):
    return _velocity_space(vg).from_numpy(vg['VGDensity'])

# %% ../nbs/api/velocitygroups.ipynb
@patch
def n_times_1(self:VelocityGroups)->Operator:
    """Operator that sums over all velocity groups then scales by the Maxwell-Boltzmann distribution."""
    return XarrayMatrixOperator(
        sparse(self['VGDensity'][:, None] * np.ones((self['VGNumber']))),
        range=_velocity_space(self), 
        source=_velocity_space(self)
    )

# %% ../nbs/api/velocitygroups.ipynb
@patch
def drho_dv(self:VelocityGroups)->Operator:
    """Derivative with respect to velocity operator.$"""
    return XarrayMatrixOperator(
        self['VGInverseWidth'] * (sparse_kronecker_matrix(self['VGNumber'], 0) - sparse_kronecker_matrix(self['VGNumber'], -1)), 
        range=_velocity_space(self), 
        source=_velocity_space(self)
    )

# %% ../nbs/api/velocitygroups.ipynb
def _vg_da(a, vg):
    range = ('Atomic velocity' + Lbl.RNG, vg['VGCenter']) if a.shape[0] == vg['VGNumber'] else ("none", ["none"])
    source = ('Atomic velocity' + Lbl.SRC, vg['VGCenter']) if a.shape[1] == vg['VGNumber'] else ("none", ["none"])
    return DataArray(a, coords=[range, source])

# %% ../nbs/api/velocitygroups.ipynb
@patch
def normalize(self:VelocityGroups)->Operator:
    """Returns the operator that normalizes a vector by dividing each component by the width of the corresponding velocity group."""
    return ScaleOperator(DataArray(self['VGInverseWidth'], {'Atomic velocity': self['VGCenter']}))

# %% ../nbs/api/velocitygroups.ipynb
@patch
def sum(self:VelocityGroups)->Operator:
    """Returns the operator that sums a vector over velocity groups."""
    return SumOperator({'Atomic velocity': self['VGCenter']})
