"""Operators on density-matrix elements"""

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

# %% auto 0
__all__ = ['parse_dm_element', 'population_element', 'level_label', 'levels', 'population', 'total_population', 'level_mask_da',
           'level_population_mask_da', 'level_population', 'sum_over_dm']

# %% ../nbs/api/dmelements.ipynb
import importlib

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.timestepping import *
from .pymor.grids import *
from .pymor.models import *
from pymor.vectorarrays.interface import VectorArray
from pymor.operators.interface import Operator

# %% ../nbs/api/dmelements.ipynb
_dm_element_hfz = re.compile(r'ρ<sub>([^,]+), \(([^,]+), ([^,]+), ([^)]+)\), \(([^,]+), ([^,]+), ([^)]+)\)</sub>')
_dm_element_toy = re.compile(r'ρ<sub>([^,]+), ([^,]+), ([^,]+)</sub>')

# %% ../nbs/api/dmelements.ipynb
def parse_dm_element(s):
    if match := _dm_element_hfz.findall(s):
        return match[0][0], match[0][1:4], match[0][4:7] 
    if match := _dm_element_toy.findall(s): 
        return match[0][0], (match[0][1],), (match[0][2],)

# %% ../nbs/api/dmelements.ipynb
@dispatch
def population_element(s):
    parsed_dm = parse_dm_element(s)
    return parsed_dm[1] == parsed_dm[2]

@dispatch
def population_element(expr:DataArray):
    return xr.apply_ufunc(population_element, expr, vectorize=True)

# %% ../nbs/api/dmelements.ipynb
@dispatch
def level_label(expr:str):
    parsed = parse_dm_element(expr)
    if parsed[1][0] != parsed[2][0]: return ''
    return parsed[1][0]

@dispatch
def level_label(expr:DataArray):
    return xr.apply_ufunc(level_label, expr, vectorize=True)

# %% ../nbs/api/dmelements.ipynb
def levels(dm_elements):
    return xr.Variable(
        'Level', 
        list(set().union(filter(None, xr.apply_ufunc(level_label, dm_elements, vectorize=True).data)))
    )

# %% ../nbs/api/dmelements.ipynb
def population(dm_elements):
    elements = dm_elements.data
    mask = [population_element(e) for e in elements]
    pop_labels = [level_label(e) for e in elements[mask]]
    return XarrayMatrixOperator(
        DataArray(
            np.identity(len(elements))[mask],
            coords=[
                ("Population (range)", elements[mask]),
                ("Density matrix (source)", elements), 
            ],
        )
    )

# %% ../nbs/api/dmelements.ipynb
def total_population(dm_elements):
    return XarrayMatrixOperator(
        population(
            dm_elements
        ).matrix.sum(
            "Population (range)"
        ),
        name="Total population"
    )

# %% ../nbs/api/dmelements.ipynb
def level_mask_da(dm_elements):
    return xr.concat([level == level_label(dm_elements) for level in levels(dm_elements)], levels(dm_elements))

# %% ../nbs/api/dmelements.ipynb
def level_population_mask_da(dm_elements):
    return level_mask_da(dm_elements) * population_element(dm_elements)

# %% ../nbs/api/dmelements.ipynb
def level_population(dm_elements):
    return XarrayMatrixOperator(
            level_population_mask_da(dm_elements),
            source="Density matrix",
            range="Level",
            name="Population"
        )

# %% ../nbs/api/dmelements.ipynb
def sum_over_dm(dm_elements):
    return SumOperator({'Density matrix' + Lbl.SRC: dm_elements})
