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

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

# %% auto 0
__all__ = ['Lbl', 'unlabel_dims', 'default_coordinate', 'dataarray_coordinate_data', 'validate_coord_data', 'coord_data_shape',
           'make_coordinates', 'XarrayVectorSpace', 'transpose_like', 'XarrayVectorArrayImpl', 'XarrayVectorArray',
           'plotly_dataarray']

# %% ../../nbs/api/pymor/vectorarrays.ipynb
from numbers import Number, Real
from functools import partial, cached_property, singledispatch
import inspect
from enum import StrEnum
import math

from fastcore.basics import patch, first
from fastcore.meta import delegates

import numpy as np
from numpy import array, ndarray
from sparse import SparseArray
import xarray as xr
from xarray import DataArray, Variable
from pymor.basic import NumpyVectorSpace
from pymor.core.base import classinstancemethod
from pymor.vectorarrays.interface import VectorArray, VectorArrayImpl, VectorSpace

from ..utilities.basic import get_item, filter_args, filter_out_args
from ..utilities.xarray import Coordinates
from ..utilities.formatting import prefix_format

# %% ../../nbs/api/pymor/vectorarrays.ipynb
import pandas as pd
import plotly.express as px
import plotly

# %% ../../nbs/api/pymor/vectorarrays.ipynb
class Lbl(StrEnum):
    SRC = ' (source)'
    RNG = ' (range)'

# %% ../../nbs/api/pymor/vectorarrays.ipynb
def unlabel_dims(obj):
    match obj:
        case str(): return obj.removesuffix(Lbl.SRC).removesuffix(Lbl.RNG)
        case dict(): return {unlabel_dims(k): v for k,v in obj.items()}
        case Coordinates(): return Coordinates({unlabel_dims(k): v.data for k,v in obj.items()})
        case DataArray(): return obj.rename({dim: unlabel_dims(dim) for dim in obj.dims})
        case XarrayVectorSpace(): return obj.with_(coord_data={unlabel_dims(k): v for k,v in obj.coord_data.items()})
    return obj

# %% ../../nbs/api/pymor/vectorarrays.ipynb
def default_coordinate(dim:str='len', size:int|list[int]|tuple[int]|None=None) -> tuple:
    if isinstance(size, list|tuple): size = size[0]
    if size: return (np.arange(size, dtype=int),)
    return ([dim],)

def dataarray_coordinate_data(da:DataArray, dim:str) -> tuple:
    if dim in da.coords:
        return (da[dim].data, da[dim].attrs)
    return default_coordinate(dim, da.sizes[dim])

# %% ../../nbs/api/pymor/vectorarrays.ipynb
@singledispatch
def validate_coord_data(obj, shape=None) -> dict:
    raise TypeError(f'Object {obj} not recognized as coordinates')

@validate_coord_data.register
def ____(obj:DataArray, shape=None):
    return {dim: dataarray_coordinate_data(obj, dim) for dim in obj.dims}

@validate_coord_data.register
def ____(obj:Coordinates, shape=None):
    return {k: (v.data, v.attrs) for k,v in obj.items()}

@validate_coord_data.register
def ____(obj:dict, shape=None):
    return {k: v if isinstance(v, tuple) else (v,) for k,v in obj.items()}

@validate_coord_data.register
def ____(obj:str, shape=None):
    return {obj: default_coordinate(obj, shape)}

@validate_coord_data.register
def ____(obj:int, shape=None): # Need int case because some internal pyMOR code calls .zeros() with integer argument
    return {obj: default_coordinate(None, obj)}

@validate_coord_data.register
def ____(obj:list, shape=None):
    if not isinstance(shape, list|tuple): raise ValueError(f'Shape must be supplied if list is used for coordinates.')
    return {dim: default_coordinate(dim, size) for dim,size in zip(obj, shape)}

@validate_coord_data.register
def ____(obj:None, shape=None):
    return {}

# %% ../../nbs/api/pymor/vectorarrays.ipynb
def coord_data_shape(coord_data):
    return [len(c[0]) for c in coord_data.values()]

# %% ../../nbs/api/pymor/vectorarrays.ipynb
def make_coordinates(coord_data:dict) -> Coordinates:
    return Coordinates({k:(k, *v) for k,v in coord_data.items()})

# %% ../../nbs/api/pymor/vectorarrays.ipynb
class XarrayVectorSpace(VectorSpace):
    """`VectorSpace` of `XarrayVectorArrays`."""

    def __init__(
        self, 
        coord_data:Coordinates|dict|DataArray|str|None=None,
        name:str|None=None, 
        attrs:dict|None=None, 
        id=None      # See `~pymor.vectorarrays.interface.VectorSpace.id`.
    ):
        if name is None and hasattr(coord_data, 'name'):
            name = coord_data.name
        coord_data = validate_coord_data(coord_data)
        self.__auto_init(locals())
        self.dims = list(self.coord_data)
        self.ndim = len(self.dims)
        self.shape = coord_data_shape(self.coord_data)
        self.size = math.prod(self.shape)

    @cached_property
    def _array(self):
        return DataArray(
            np.zeros(self.shape),
            dims=self.dims,
            coords=make_coordinates(self.coord_data), 
            name=self.name, 
            attrs=self.attrs
        ) 
    
    @cached_property
    def coords(self):
        return self._array.coords
    
    @property
    def dim(self): return self.size
    
    @classmethod
    def _array_factory(cls, data, space, data_dims=None, extended_coord_data=None):
        return XarrayVectorArray(space, XarrayVectorArrayImpl(data, space, data_dims, extended_coord_data))

    @classinstancemethod
    def make_array(cls, obj, id=None):
        vec = cls._array_factory(obj, id=id)
        # vec.name = obj.name
        return vec

    @make_array.instancemethod
    def make_array(self, data, data_dims=None, extended_coord_data=None):
        """:noindex:"""
        return self._array_factory(data, self, data_dims, extended_coord_data)
    
    def zeros(
        self, 
        coords:dict|None=None, # Coordinates to extend vector space over
        reserve=0
    ): 
        raise NotImplementedError # Defined below
    
    def ones(
        self, 
        coords:dict|None=None, # Coordinates to extend vector space over
        reserve=0
    ): 
        raise NotImplementedError # Defined below

    # def __str__(self):
    #     return _coords_str(self.coords)

    # def short_str(self):
    #     return _short_coords_str(self.coords)
    
    def _repr_html_(self, show_details=True):
        if not show_details:
            return ' ⛒ '.join(f'<b>{dim}</b>({length})' for dim,length in zip(self.dims, self.shape))
        return f'''
            <details closed>
                <summary>
                    {self._repr_html_(show_details=False)}
                </summary>
                {self._array._repr_html_()}
            </details>
        '''    
    
    def __hash__(self):
        return hash(self.dim)

    @property
    def is_scalar(self):
        return self.dim == 1

# %% ../../nbs/api/pymor/vectorarrays.ipynb
# @patch(as_prop=True)
# def ndim(self:XarrayVectorSpace): 
#     """Number of dimensions of the vector space."""
#     return self._array.ndim

# %% ../../nbs/api/pymor/vectorarrays.ipynb
@patch
def rename(
    self:XarrayVectorSpace,
    name,
):
    return self.with_(name=name) #XarrayVectorSpace(self._array.rename(new_name_or_name_dict, **names))

# %% ../../nbs/api/pymor/vectorarrays.ipynb
@patch
def __eq__(self:XarrayVectorSpace, other):
    return (
        type(other) is type(self)
        and all(k1 == k2 and np.array_equal(v1[0], v2[0]) for (k1,v1),(k2,v2) in zip(self.coord_data.items(), other.coord_data.items()))
        ## NEED TO FIX THIS
        # and all((self.coords_dict[k] == other.coords_dict[k]) for k in self.coords_dict)
    )

# %% ../../nbs/api/pymor/vectorarrays.ipynb
@patch
def __mul__(self:XarrayVectorSpace, other):
    if not isinstance(other, XarrayVectorSpace): raise NotImplementedError
    return XarrayVectorSpace(self.coord_data | other.coord_data, name=other.name)

# %% ../../nbs/api/pymor/vectorarrays.ipynb
def transpose_like(a:ndarray, dims, new_dims):
    return a.transpose([dims.index(d) for d in new_dims])

# %% ../../nbs/api/pymor/vectorarrays.ipynb
class XarrayVectorArrayImpl(VectorArrayImpl):

    def __init__(
        self, 
        _data:ndarray|DataArray, 
        _space, 
        _data_dims=None, 
        _extended_coord_data:Coordinates|dict|None=None
    ):
        if isinstance(_data, DataArray):
            _data_dims = _data.dims
            _extended_coord_data = _data
            _data = _data.data
        if isinstance(_data, SparseArray): 
            _data = _data.todense()
        if _data_dims is None:
            _data_dims = _space.dims
        _extended_coord_data = validate_coord_data(_extended_coord_data)
        _extended_coord_data = {k:v for k,v in _extended_coord_data.items() if k not in _space.dims}
        self._dims = list(_extended_coord_data)
        all_dims = self._dims + _space.dims
        if set(all_dims) != set(_data_dims): raise ValueError(f'Data dims {_data_dims} do not match specified dims {all_dims}')
        if len(all_dims) != _data.ndim: raise ValueError(f'Dimensions {all_dims} do not match number of axes of data of shape {_data.shape}.')
        _data = transpose_like(_data, _data_dims, all_dims)
        _data_dims = all_dims
        self.__auto_init(locals())
        self._shape = _data.shape
        self._sizes = dict(zip(_data_dims, self._shape))
        self._array = None
        self._extended_shape = coord_data_shape(self._extended_coord_data)
        self._extended_size = math.prod(self._extended_shape)

    @property
    def array(self): 
        if self._array is None:
            self._array = DataArray(self._data, dims=self._data_dims, coords=make_coordinates(self._extended_coord_data | self._space.coord_data))
        return self._array

    @property
    def _coords(self): return self.array.coords

    @property
    def _len(self): return self._data.size // self._space.size
    
    def __len__(self): return self._len

    def copy(self, deep, ind):
        _data = self._data if ind is None else self._data[ind].copy()
        return XarrayVectorArrayImpl(_data, self._space, self._data_dims, self._extended_coord_data)
    
    def to_numpy(self, ensure_copy, ind): raise NotImplementedError
    def real(self, ind): raise NotImplementedError
    def imag(self, ind): raise NotImplementedError
    def conj(self, ind): raise NotImplementedError
    def delete(self, ind): raise NotImplementedError
    def append(self, other, remove_from_other, oind): raise NotImplementedError
    def scal(self, alpha, ind): raise NotImplementedError
    def scal_copy(self, alpha, ind): raise NotImplementedError
    def axpy(self, alpha, x, ind, xind): raise NotImplementedError
    def axpy_copy(self, alpha, x, ind, xind): raise NotImplementedError
    def inner(self, other, ind, oind): raise NotImplementedError
    def pairwise_inner(self, other, ind, oind): raise NotImplementedError
    def lincomb(self, coefficients, ind): raise NotImplementedError
    def norm(self, ind): raise NotImplementedError
    def norm2(self, ind): raise NotImplementedError
    def dofs(self, dof_indices, ind): raise NotImplementedError
    def amax(self, ind): raise NotImplementedError

# %% ../../nbs/api/pymor/vectorarrays.ipynb
class XarrayVectorArray(VectorArray):
    """`VectorArray` implementation via xarray arrays."""

    impl_type = XarrayVectorArrayImpl

    @property
    def data(self): return self.impl._data
    
    @property
    def array(self): return self.impl.array

    @property
    def shape(self): return self.impl._shape
    
    @property
    def coords(self): return self.impl._coords

    @property
    def dims(self): return self.impl._space.dims
    
    @property
    def extended_dims(self): return self.impl._dims

    @property
    def extended_coord_data(self): return self.impl._extended_coord_data

    @property
    def extended_shape(self): return self.impl._extended_shape
    
    # def _coords_str(self, dims):
    #     if not dims: return '{1}'
    #     sizes = self.impl._sizes
    #     return '{' + ' ⨉ '.join(f'{dim}({sizes[dim]})' for dim in dims) + '}'        
    
    # def __str__(self):
    #     return self._coords_str(self.dims)[:-1] + ", " + self._coords_str(self.extended_dims)[1:]
    
    # def short_str(self):
    #     return _short_coords_str(self.space)[:-1] + ", " + _short_coords_str(self.coords, self.dims)[1:]
    
    def _repr_html_(self, show_details=True):
        space = self.space._repr_html_(show_details=False)
        extended = [f'{k}({len(v[0])})' for k,v in self.extended_coord_data.items()]
        if not extended: extended = ['∅']
        out = ' ⛒ '.join(extended + [space])
        if not show_details: return out
        return f'''
            <details closed>
                <summary>
                    {out}
                </summary>
                {self.array._repr_html_()}
            </details>
        '''

# %% ../../nbs/api/pymor/vectorarrays.ipynb
@patch
def from_data(self:XarrayVectorSpace, data, data_dims=None, extended_coord_data=None):
    return self.make_array(data, data_dims, extended_coord_data)

# %% ../../nbs/api/pymor/vectorarrays.ipynb
@patch
def from_numpy(
    self:XarrayVectorSpace, 
    data:ndarray, 
    extended_coord_data:dict|str|None=None, 
    id=None, 
    ensure_copy=False
)->XarrayVectorArray: # A vector array with data from the numpy array
    """Return an `XarrayVectorArray` in the vector space with data from the `ndarray` `data`."""
    if ensure_copy: data = data.copy()
    if extended_coord_data is None:
        extended_coord_data = {}
    elif isinstance(extended_coord_data, str):
        extended_coord_data = {extended_coord_data: np.arange(data.size // self.dim)}
    extended_coord_data = validate_coord_data(extended_coord_data)
    extended_shape = coord_data_shape(extended_coord_data)
    shape = extended_shape + self.shape
    return self.from_data(data.reshape(shape), list(extended_coord_data) + self.dims, extended_coord_data)

# %% ../../nbs/api/pymor/vectorarrays.ipynb
@patch
def from_xarray(
    self:XarrayVectorSpace, 
    data:DataArray
)->XarrayVectorArray:
    """Return an `XarrayVectorArray` containing data from `DataArray` `data`."""
    # extended_coord_data = {k:v for k,v in data.coords.items() if k not in self.dims}
    return self.from_data(data)

# %% ../../nbs/api/pymor/vectorarrays.ipynb
@patch
def concatenate(self:XarrayVectorSpace, arrays, dim):
    return self.from_xarray(xr.concat([array.array for array in arrays], Variable(first(dim), first(dim.values()))))

# %% ../../nbs/api/pymor/vectorarrays.ipynb
@patch
def array_from_fn(
    self:XarrayVectorSpace,
    fn,
    extended_coords:dict|str|None=None, # Coordinates to extend vector space over
    reserve=0,
)->XarrayVectorArray: # `XarrayVectorArray` with all elements equal to zero
    """Return `XarrayVectorArray` of null vectors in XarrayVectorSpace optionally extended to include supplied coordinates `coords`."""
    extended_coords = validate_coord_data(extended_coords)
    shape = coord_data_shape(extended_coords)
    data = fn(shape + self.shape)
    return self.from_data(data, list(extended_coords) + self.dims, extended_coords)

# %% ../../nbs/api/pymor/vectorarrays.ipynb
@patch
def zeros(
    self:XarrayVectorSpace, 
    extended_coords:dict|None=None, # Coordinates to extend vector space over
    reserve:int=0
)->XarrayVectorArray: # `XarrayVectorArray` with all elements equal to zero
    """Return `XarrayVectorArray` of null vectors in XarrayVectorSpace optionally extended to include supplied coordinates `coords`."""
    return self.array_from_fn(np.zeros, extended_coords, reserve)

# %% ../../nbs/api/pymor/vectorarrays.ipynb
@patch
def ones(
    self:XarrayVectorSpace, 
    extended_coords:dict|None=None, # Coordinates to extend vector space over
    reserve:int=0
)->XarrayVectorArray: # `XarrayVectorArray` with all elements equal to one
    """Return `XarrayVectorArray` of vectors with each element equal to one in XarrayVectorSpace optionally extended to include supplied coordinates `coords`."""
    return self.array_from_fn(np.ones, extended_coords, reserve)

# %% ../../nbs/api/pymor/vectorarrays.ipynb
from pymor.vectorarrays.interface import _create_random_values

# %% ../../nbs/api/pymor/vectorarrays.ipynb
@patch
def random(
    self:XarrayVectorSpace, 
    extended_coords=None, 
    distribution='uniform',
    name=None
)->XarrayVectorArray: # A random `XarrayVectorArray` in the vector space
    out = self.array_from_fn(lambda shape: _create_random_values(shape, distribution), extended_coords)
    if name is not None: out.name = name
    return out

# %% ../../nbs/api/pymor/vectorarrays.ipynb
@patch
def __contains__(self:XarrayVectorSpace, other):
    # This is reversed because when pymor calls `assert U in op.source`
    # we really want to assert that the source space of `op` is contained in the space of `U`, 
    # which more logically would be written "assert op.source in U"
    if isinstance(other, XarrayVectorSpace):
        theirs = other.coord_data
    elif isinstance(other, XarrayVectorArray):
        theirs = other.space.coord_data
    else:
        return False
    mine = self.coord_data
    if not all(k in theirs for k in mine): return False
    if any(not np.array_equal(mine[k][0], theirs[k][0]) for k in mine): return False
    return True
    
    # if isinstance(other, XarrayVectorSpace):
    #     return other._array.coords.contain(self._array.coords)
    # elif isinstance(other, XarrayVectorArray):
    #     return other.space._array.coords.contain(self._array.coords)
    # return False

# %% ../../nbs/api/pymor/vectorarrays.ipynb
@patch
def array_labeled_as_source(self:XarrayVectorArray):
    return self.array.rename({d: d + Lbl.SRC for d in self.dims})

# %% ../../nbs/api/pymor/vectorarrays.ipynb
@patch
def item(self:XarrayVectorArray):
    """Return the single value stored in the array if there is only one value."""
    return self.data.item()

# %% ../../nbs/api/pymor/vectorarrays.ipynb
@patch(as_prop=True)
def stacked_array(self:XarrayVectorArrayImpl):
    core = {'core': self._space.dims} if self._space.dims else {}
    extended = {'extended': self._dims} if self._dims else {}
    return self.array.stack(core | extended, create_index=False)

# %% ../../nbs/api/pymor/vectorarrays.ipynb
@patch
def stacked_data(self:XarrayVectorArrayImpl):
    return transpose_like(
        self._data, self._data_dims, self._space.dims + list(self._extended_coord_data)
    ).reshape(
        self._space.size, self._extended_size
    )

# %% ../../nbs/api/pymor/vectorarrays.ipynb
@patch
def to_numpy(self:XarrayVectorArrayImpl, ensure_copy, ind):
    out = self.stacked_data()
    if ind is not None: out = out[:, ind]
    if ensure_copy and not out.flags['OWNDATA']:
        return out.copy()
    else:
        return out

    # array = self.stacked_array.to_numpy()
    # if not self._space.dims: array = array[None, :]
    # if not self._dims: array = array[:, None]
    # if ensure_copy and not array.flags['OWNDATA']:
    #     return array.copy()
    # else:
    #     return array

# %% ../../nbs/api/pymor/vectorarrays.ipynb
@patch
def scal(self:XarrayVectorArrayImpl, alpha, ind):
    if ind is not None: raise NotImplementedError
    self._data *= alpha

# %% ../../nbs/api/pymor/vectorarrays.ipynb
@patch
def scal_copy(self:XarrayVectorArrayImpl, alpha, ind):
    if ind is not None: raise NotImplementedError
    return type(self)(self._data * alpha, self._space, self._data_dims, self._extended_coord_data)

# %% ../../nbs/api/pymor/vectorarrays.ipynb
@patch
def axpy(self:XarrayVectorArrayImpl, alpha, x, ind, xind):
    if ind is not None: raise NotImplementedError
    B = x._data
    if isinstance(alpha, Number):
        if alpha == 1:
            self._data += B
            return
        elif alpha == -1:
            self._data -= B
            return

    self._data += B * alpha

# %% ../../nbs/api/pymor/vectorarrays.ipynb
@patch
def axpy_copy(self:XarrayVectorArrayImpl, alpha, x, ind, xind):
    if ind is not None or xind is not None: raise NotImplementedError
    B = x._data
    # if issparse(B): B = B.todense()
    if isinstance(alpha, Number):
        if alpha == 1:
            return type(self)(self._data + B, self._space, self._data_dims, self._extended_coord_data)
        elif alpha == -1:
            return type(self)(self._data - B, self._space, self._data_dims, self._extended_coord_data)
    return type(self)(self._data + B * alpha, self._space, self._data_dims, self._extended_coord_data)

# %% ../../nbs/api/pymor/vectorarrays.ipynb
@patch
def __mul__(self:XarrayVectorArray, other):
    match other:
        case DataArray():
            other = unlabel_dims(other)
            if self.coords.contain(other.coords):
                out = self.space.from_xarray(self.array * other)
                out.name = self.name
                return out
            else:
                raise ValueError("DataArray coordinates are not a subset of VectorArray coordinates.")
        case XarrayVectorArray(): 
            out = (self.space * other.space).from_xarray(self.array * other.array)
            out.name = self.name
            return out
        case Number():
            return type(self)(self.space, self.impl.scal_copy(other, self.ind))
    raise NotImplementedError

# %% ../../nbs/api/pymor/vectorarrays.ipynb
def _reim(da:DataArray):
    return xr.concat([da.real, da.imag], dim=xr.Variable('Part', ['Re', 'Im'])).transpose(*da.dims, 'Part')

# %% ../../nbs/api/pymor/vectorarrays.ipynb
@patch
def real(self:XarrayVectorArrayImpl, ind):
    return XarrayVectorArrayImpl(self.copy(False, ind)._data.real, self._space, self._data_dims, self._extended_coord_data)
    
@patch
def imag(self:XarrayVectorArrayImpl, ind):
    return XarrayVectorArrayImpl(self.copy(False, ind)._data.imag, self._space, self._data_dims, self._extended_coord_data)
    
@patch
def conj(self:XarrayVectorArrayImpl, ind):
    if np.isrealobj(self._data):
        return self.copy(False, ind)
    return XarrayVectorArrayImpl(np.conj(self.copy(False, ind)._data), self._space, self._data_dims, self._extended_coord_data)
    
@patch
def reim(self:XarrayVectorArrayImpl):
    return XarrayVectorArrayImpl(
        np.array([self._data.real, self._data.imag]), 
        self._space, ['Part'] + self._data_dims, 
        {'Part': (['Re', 'Im'], {})} | self._extended_coord_data
    )
    
@patch(as_prop=True)
def reim(self:XarrayVectorArray):
    impl = self.impl.reim()
    if impl is self.impl:
        return self.copy()
    else:
        return type(self)(self.space, impl)

# %% ../../nbs/api/pymor/vectorarrays.ipynb
plotly.io.templates.default = "plotly_white"
plotly.io.templates['plotly_white'].layout.legend = plotly.graph_objects.layout.Legend(tracegroupgap=0)
plotly.io.templates["plotly_white"].layout.width = 700
plotly.io.templates["plotly_white"].layout.height = 400
plotly.io.templates["plotly_white"].layout.autosize = False
plotly.io.templates["plotly_white"].layout.colorway = array([
       '#2fa1da', '#fb4f2f', '#e4ae38', '#6d904f', '#8a8a8a', '#16bdcf',
       '#9367bc', '#d62628', '#1f77b3', '#e277c1', '#8c564b', '#bcbc21',
       '#3a0182', '#004200', '#0fffa8', '#5d003f', '#c6bcff', '#424f52',
       '#b80080', '#ffb6b3', '#7c0100', '#6026ff', '#ffff9a', '#aec8aa',
       '#00857c', '#543a00', '#93fbff', '#00bf00', '#7c00a0', '#aa7200',
       '#90ff00', '#01bd89', '#00447b', '#c8826e', '#ff1f82', '#dd00ff',
       '#057400', '#644460', '#878eff', '#ffb5f4', '#526236', '#cd85ff',
       '#676983', '#bdb3bd', '#a56089', '#95d3ff', '#0100f7', '#ff8001',
       '#8a2844', '#aca06d', '#52448a', '#c8ffd8', '#aa4600', '#ff798e',
       '#82d370', '#909ebf', '#9300f4', '#ebcf9a', '#ac8ab1', '#006249',
       '#ffdb00', '#877750', '#7eaaa3', '#000097', '#f400c6', '#643328',
       '#006677', '#03e2c8', '#a736ae', '#c4dbe1', '#4d6eff', '#9a9301',
       '#cd576b', '#efddfd', '#795900', '#5e879a', '#b3ff91', '#5d726b',
       '#520066', '#058750', '#831f6e', '#3b9505', '#647200', '#f0a06b',
       '#5e4f44', '#bc0049', '#cf6726', '#d695aa', '#895dff', '#826b75',
       '#2a54b8', '#6e7cba', '#e6d4d3', '#5d0018', '#7c3b01', '#80b17c',
       '#c8d87c', '#00e83b', '#7cb1ff', '#ff54ff', '#a32621', '#1ce4ff',
       '#7caf3b', '#7b4b90', '#dfff48', '#6b00c3', '#cda897', '#bd62c3',
       '#89cdcd', '#4603c8', '#5d9179', '#414901', '#05a79c', '#cf8c36',
       '#fff7cf', '#425470', '#b544ff', '#cf4993', '#cfa3df', '#93d400',
       '#a793da', '#2da557', '#8ce2b5', '#a3a89c', '#6b5bb6', '#ff7e5d',
       '#a78289', '#afbdd8', '#2ac3ff', '#a5673d', '#f690fd', '#874b64',
       '#ff0c4b', '#215d23', '#4291ff', '#87829c', '#672d44', '#b14f41',
       '#004d52', '#5e1a00', '#ac4167', '#4f3167', '#d6fffd', '#7eb5d1',
       '#a8b869', '#ff95ca', '#c87495', '#364f38', '#ffcf62', '#5d5762',
       '#879375', '#a877ff', '#03c862', '#e6bdd4', '#d4e2cf', '#876790',
       '#897c26', '#cddbff', '#aa676b', '#313474', '#ff5da8', '#009aaf',
       '#70ffdd', '#775b38', '#4f649a', '#cc00b3', '#567b54', '#506e7b',
       '#015e91', '#aabcbd', '#017e99', '#03dd97', '#873a2b', '#ef958e',
       '#75c6aa', '#70695d', '#ccdb08', '#af8556', '#d80075', '#9c3f80',
       '#d84400', '#dd6754', '#5eff79', '#d4b172', '#62265d', '#baa13d',
       '#d8f2b3', '#56018e', '#a19aaa', '#4d4926', '#a3a8ff', '#ace8db',
       '#995901', '#ac00e2', '#46822f', '#cac3ac', '#00c4b5', '#605277',
       '#336d67', '#a59180', '#8399a1', '#fd5664', '#7095d1', '#728c07',
       '#7e004b', '#152fa0', '#d1c1e2', '#c885cf', '#6b444b', '#7e0023',
       '#00a179', '#b1a8cf', '#f90000', '#afe8ff', '#939e4f', '#727982',
       '#d82d54', '#466001', '#0059ff', '#773fb5', '#ace460', '#674424',
       '#525d50', '#957267', '#a8e49a', '#a30057', '#d862f6', '#8e7ccf',
       '#ffbc93', '#a30091', '#9affb8', '#a7c1ff', '#f46200', '#e4efff',
       '#b89ca3', '#609593', '#ff9e34', '#8c2800', '#726b31', '#df824d',
       '#af7bd4', '#bc2d00', '#7b6ea3', '#484262', '#c6a3ff', '#004d28',
       '#c3c68e', '#df48d6', '#e6e864', '#e4c10a', '#00f4f0', '#9e5ba1',
       '#4b41b6', '#64338e', '#757e6b', '#a88936'])

# %% ../../nbs/api/pymor/vectorarrays.ipynb
def _categorical_dims(array:DataArray): return [dim for dim in array.dims if not isinstance(array[dim][0].item(), Number)]
def _numerical_dims(array:DataArray):   return [dim for dim in array.dims if     isinstance(array[dim][0].item(), Number)]

# %% ../../nbs/api/pymor/vectorarrays.ipynb
def _numerical_dim_sort_order(da, dim):
    """Assign priority for each dimension to be put on the x axis.
    Lower number is higher priority. Time is preferred on a scrubber, velocity preferred on the x-axis.
    """
    if "time" in dim.lower(): return 2
    if "velocity" in dim.lower(): return 0
    if isinstance(da[dim][0].item(), Real): return 1
    return 0

# %% ../../nbs/api/pymor/vectorarrays.ipynb
def _set_plotly_frame(fig, n):
    fig.update_layout(sliders=[dict(active=n)])
    for i in range(len(fig.data)):
        fig.data[i].y = fig.frames[n].data[i].y

# %% ../../nbs/api/pymor/vectorarrays.ipynb
def _format_slider_labels(fig, slider_label_precision=2):
    for step in fig.layout.sliders[0].steps:
        step.label = prefix_format(float(step.label), slider_label_precision)

# %% ../../nbs/api/pymor/vectorarrays.ipynb
def plotly_dataarray(da:DataArray, slider_label_precision=2, **kwargs):
    """Visualize the data contained in the `XarrayVectorArray`.
    Put the first dimension with numerical coordinates on the x axis, and include a scrubber if there is a second numerical dimension. 
    Put categorical dimensions in the legend. If the data are complex, add the real and imaginary parts as a categorical dimension.
    If there are no numerical dimensions, plot the categorical dimension(s) as a horizontal bar plot.
    """
    da = da.copy()
    if np.any(np.iscomplex(da)): da = _reim(da)
    if da.name is None: da = da.rename('Value')
    df = da.to_dataframe().reset_index()
    num = _numerical_dims(da)
    cat = _categorical_dims(da)
    if not num and da.ndim <= 2: 
        return px.bar(
            df, 
            y=cat[0], 
            x=da.name, 
            color=get_item(cat, 1), 
            orientation='h', 
            **filter_args(px.bar, **kwargs)
        ).update_layout(
            legend_title=None
        )
    if da.ndim > 4: raise NotImplementedError
    if len(num) > 2: num, cat = num[:2], cat + num[2:]
    if len(cat) > 2: cat, num = cat[:2], num + cat[2:]
    num = sorted(num, key=partial(_numerical_dim_sort_order, da))
    fig = px.line(
        df,
        x=get_item(num, 0),
        y=da.name,
        animation_frame=get_item(num, 1),
        color=get_item(cat, 0),
        line_dash=get_item(cat, 1),
        **filter_args(px.line, **kwargs)
    ).update_layout(
        legend_title=None
    )
    if len(num) == 2:
        padding = .05
        ymin, ymax = da.min(), da.max()
        yrange = (ymax - ymin)
        fig.update_layout(
            yaxis_range=(ymin - padding * yrange, ymax + padding * yrange)
        )
        _format_slider_labels(fig, slider_label_precision=slider_label_precision)
        _set_plotly_frame(fig, len(da[num[1]])//3)
        # Set total animation time to 3s, unless that would make individual frames shorter than 100ms, which seems to be performance limit
        fig.layout.updatemenus[0].buttons[0].args[1]['frame']['duration'] = max(3000/len(da[num[1]]), 100) 
        fig.layout.updatemenus[0].buttons[0].args[1]['transition']['duration'] = 0
    kwargs = filter_out_args(px.line, **kwargs)
    for k in kwargs:
        try: fig.update_layout({k: kwargs[k]})
        except ValueError: pass
        try: fig.update_traces({k: kwargs[k]})
        except ValueError: pass
    return fig

# %% ../../nbs/api/pymor/vectorarrays.ipynb
## Could put plotting hints in the attrs of the coordinates inside the model
@patch
@delegates(plotly_dataarray)
def visualize(self:XarrayVectorArray, **kwargs):
    """Visualize the data contained in the `XarrayVectorArray`.
    Put the first dimension with numerical coordinates on the x axis, and include a scrubber if there is a second numerical dimension.
    Atomic velocity is prioritized to be on the x axis, time is prioritized to be on the scrubber.
    Put categorical dimensions in the legend. If the data are complex, add the real and imaginary parts as a categorical dimension.
    If there are no numerical dimensions, plot the categorical dimension(s) as a horizontal bar plot.
    """
    return plotly_dataarray(self.array, **kwargs)
