###################################################################################
# Copyright 2021 National Technology & Engineering Solutions of Sandia,           #
# LLC (NTESS). Under the terms of Contract DE-NA0003525 with NTESS, the           #
# U.S. Government retains certain rights in this software.                        #
# If you want to use this code, please refer to the README.rst and LICENSE files. #
###################################################################################

import logging
import numpy as np
import mpi4py.rc
mpi4py.rc.initialize = False
from mpi4py import MPI
from PyNucleus_fem.DoFMaps import P1_DoFMap, P2_DoFMap
from PyNucleus_fem import str2DoFMap
from PyNucleus_base.myTypes import REAL
from PyNucleus_base.linear_operators import CSR_LinearOperator
from PyNucleus_base.linear_operators import SSS_LinearOperator
from . restrictionProlongation import buildRestrictionProlongation
from PyNucleus_base import TimerManager
from PyNucleus_base.ip_norm import ip_serial, norm_serial, ip_distributed, norm_distributed
from PyNucleus_fem import (assembleDrift,
                 assembleMatrix,
                 mass_0d_in_1d_sym_P1,
                 mass_1d_in_2d_sym_P1,
                 mass_1d_in_2d_sym_P2,
                 DistributedLinearOperator,
                 CSR_DistributedLinearOperator,
                 function,
                 DIRICHLET, NEUMANN, HOMOGENEOUS_DIRICHLET, HOMOGENEOUS_NEUMANN, boundaryConditions)
from PyNucleus_fem.mesh import (PHYSICAL, NO_BOUNDARY, INTERIOR_NONOVERLAPPING, INTERIOR)
LOGGER = logging.getLogger(__name__)

# what should be built
DOFMAPS = 1
RESTRICTION_PROLONGATION = 2
SPARSITY_PATTERN = 4
OVERLAPS = 8
ASSEMBLY = 16

NO_BUILD = 0
DOFMAPS_ONLY = DOFMAPS
RESTRICTION_PROLONGATION_ONLY = DOFMAPS + RESTRICTION_PROLONGATION
SPARSITY_ONLY = DOFMAPS + OVERLAPS + RESTRICTION_PROLONGATION + SPARSITY_PATTERN
SINGLE_LEVEL = DOFMAPS + OVERLAPS + ASSEMBLY
FULL_BUILD = DOFMAPS + OVERLAPS + RESTRICTION_PROLONGATION + ASSEMBLY

# What information is retained in meshLevels
DELETE_MESH = 0
KEEP_MESH = 1

class level:
    def __init__(self, params, previousLevel=None,
                 comm=None, label='', startLevelNo=0,
                 isLastLevel=False):
        self.params = params
        self.previousLevel = previousLevel
        if previousLevel is not None:
            assert not previousLevel.isLastLevel
        self.startLevelNo = startLevelNo
        self.nextLevel = None
        self.comm = comm
        self.label = label
        self.isLastLevel = isLastLevel

        label = '{}: '.format(self.levelID)
        self.Timer = TimerManager(LOGGER,
                                  comm=self.comm, prefix=label)

    def getLevelNo(self):
        if self.previousLevel is None:
            return self.startLevelNo
        else:
            return self.previousLevel.getLevelNo()+1

    levelNo = property(fget=getLevelNo)

    def getLevelID(self):
        if len(self.label) > 0:
            label = '{} {}'.format(self.label, self.levelNo)
        else:
            label = 'Level {}'.format(self.levelNo)
        return label

    levelID = property(fget=getLevelID)

    def __repr__(self):
        if len(self.label) > 0:
            label = '{} {}'.format(self.label, self.levelNo)
        else:
            label = '{}'.format(self.levelNo)
        s = '{} {}\n'.format(self.__class__.__name__, label)
        return s

######################################################################

class meshLevel(level):
    def __init__(self, mesh, params, previousLevel=None,
                 interfaces=None, meshOverlaps=None,
                 interiorBL=None,
                 comm=None,
                 label='', meshInformationPolicy=KEEP_MESH, startLevelNo=0,
                 isLastLevel=False):
        super(meshLevel, self).__init__(params, previousLevel, comm,
                                        label, startLevelNo, isLastLevel)
        self.mesh = mesh
        self.global_mesh = None
        self.interfaces = interfaces
        if self.interfaces is not None and self.params['debugOverlaps']:
            self.interfaces.validate(self.mesh, self.comm, label='Mesh interface \'{} {}\''.format(self.label, self.levelNo))
        self.meshOverlaps = meshOverlaps
        if self.meshOverlaps is not None and self.params['debugOverlaps']:
            self.meshOverlaps.check(self.mesh, self.comm, label='Mesh overlap \'{} {}\''.format(self.label, self.levelNo))
        self.interiorBL = interiorBL
        self.algebraicLevel = None
        self.meshInformationPolicy = meshInformationPolicy
        self._h = None
        self.algebraicLevelType = algebraicLevel

    def setAlgebraicLevelType(self, algLevelType):
        self.algebraicLevelType = algLevelType

    def refine(self, meshInformationPolicy):
        with self.Timer('Refined mesh'):
            newMesh, self.lookup = self.mesh.refine(returnLookup=True)
            if self.params['meshTransformation'] is not None:
                self.params['meshTransformation'](newMesh, self.lookup)
        if self.interfaces is not None:
            with self.Timer('Refined interfaces'):
                self.interfaces.refine(newMesh)
        if self.meshOverlaps is not None:
            with self.Timer('Refined mesh overlaps'):
                meshOverlaps = self.meshOverlaps.copy()
                meshOverlaps.refine(newMesh)
        if self.meshOverlaps is None:
            meshOverlaps = None
        if self.interiorBL is not None:
            with self.Timer('Refined boundary layers'):
                self.interiorBL.refine(newMesh)
        newMeshLevel = meshLevel(newMesh, self.params, self, self.interfaces, meshOverlaps, self.interiorBL, self.comm, self.label, meshInformationPolicy)
        if hasattr(self, 'numberCellsBeforeExtension'):
            newMeshLevel.numberCellsBeforeExtension = 2**self.mesh.dim * self.numberCellsBeforeExtension
        if hasattr(self, 'numberCellsLastLayer'):
            newMeshLevel.numberCellsLastLayer = 2**self.mesh.dim * self.numberCellsLastLayer
        newMeshLevel.setAlgebraicLevelType(self.algebraicLevelType)
        self.nextLevel = newMeshLevel
        return newMeshLevel

    def copy(self):
        newMeshLevel = meshLevel(self.mesh, self.params, self, self.interfaces, self.meshOverlaps, self.interiorBL, self.comm, self.label, self.meshInformationPolicy)
        return newMeshLevel

    def getIsDistributed(self):
        return self.interfaces is not None

    isDistributed = property(fget=getIsDistributed)

    def getAlgebraicLevel(self, buildType):
        self.algebraicLevel = self.algebraicLevelType(self, buildType)
        return self.algebraicLevel

    def clean(self):
        if self.meshInformationPolicy == DELETE_MESH:
            self.mesh = None
            self.meshOverlaps = None
        self.interfaces = None

    def getLevelDict(self):
        lvl = {}
        if self.mesh is not None:
            lvl['mesh'] = self.mesh
        if self.interfaces is not None:
            lvl['interfaces'] = self.interfaces
        if self.meshOverlaps is not None:
            lvl['meshOverlaps'] = self.meshOverlaps
        return lvl

    @staticmethod
    def fromLevelDict(lvl, params={}, previousLevel=None, comm=None, startLevelNo=0, label=''):
        alvl = meshLevel(None, params, previousLevel, comm=comm, startLevelNo=startLevelNo, label=label)
        if 'mesh' in lvl:
            alvl.mesh = lvl['mesh']
        if 'interfaces' in lvl:
            alvl.interfaces = lvl['interfaces']
        if 'meshOverlaps' in lvl:
            alvl.meshOverlaps = lvl['meshOverlaps']
        return alvl

    def __repr__(self):
        s = super(meshLevel, self).__repr__()
        if self.mesh is not None:
            s += ' mesh: '+self.mesh.__repr__()
        if self.interfaces is not None:
            s += self.interfaces.__repr__()
        return s

    def getH(self):
        if self._h is None:
            h = self.mesh.h
            if self.comm is not None:
                self._h = self.comm.allreduce(h, op=MPI.MAX)
        return self._h

    h = property(fget=getH)

######################################################################

class algebraicLevelBase(level):
    def __init__(self, meshLevel, buildType):
        if meshLevel.previousLevel is not None:
            previousLevel = meshLevel.previousLevel.algebraicLevel
        else:
            previousLevel = None
        super(algebraicLevelBase, self).__init__(meshLevel.params, previousLevel, meshLevel.comm, meshLevel.label, meshLevel.levelNo, meshLevel.isLastLevel)
        self.meshLevel = meshLevel
        self.P = None
        self.R = None
        self.DoFMap = None
        self.algebraicOverlaps = None
        self.build(buildType)

    def build(self, buildType):

        buildNeumann = self.params.get('buildNeumann', False)
        element = self.params['element']
        reorder = self.params['reorder']
        commType = self.params['commType']
        DoFMap_type = str2DoFMap(element)

        # Set DoFMap
        if buildType & DOFMAPS:
            if 'tag' in self.params:
                self.DoFMap = DoFMap_type(self.meshLevel.mesh, self.params['tag'])
            elif 'boundaryCondition' in self.params:
                if self.params['boundaryCondition'] in (HOMOGENEOUS_NEUMANN, DIRICHLET, NEUMANN):
                    self.DoFMap = DoFMap_type(self.meshLevel.mesh, NO_BOUNDARY)
                elif self.params['boundaryCondition'] == HOMOGENEOUS_DIRICHLET:
                    self.DoFMap = DoFMap_type(self.meshLevel.mesh, PHYSICAL)
                else:
                    raise NotImplementedError(boundaryConditions[self.params['boundaryCondition']])
            else:
                if self.isLastLevel and self.params['interiorBC'] == 'homogeneousDirichlet' and hasattr(self.meshLevel, 'numberCellsLastLayer'):
                    self.DoFMap = DoFMap_type(self.meshLevel.mesh, [PHYSICAL,
                                                                    INTERIOR],
                                              skipCellsAfter=self.meshLevel.mesh.num_cells-self.meshLevel.numberCellsLastLayer)
                elif not hasattr(self.meshLevel, 'numberCellsLastLayer') or not self.isLastLevel or self.params['interiorBC'] == 'homogeneousNeumann':
                    self.DoFMap = DoFMap_type(self.meshLevel.mesh, [PHYSICAL])
                else:
                    raise NotImplementedError()
            if buildNeumann:
                self.DoFMapNeumann = DoFMap_type(self.meshLevel.mesh, [PHYSICAL])

        if not reorder:
            if buildType & OVERLAPS:
                # build algebraic overlaps
                if self.meshLevel.meshOverlaps is not None:
                    with self.Timer('Build algebraic overlaps of type \'{}\''.format(commType)):
                        self.algebraicOverlaps = self.meshLevel.meshOverlaps.getDoFs(self.meshLevel.mesh, self.DoFMap, commType,
                                                                                allowInteriorBoundary=self.params['interiorBC'] == 'homogeneousNeumann' or not self.isLastLevel)
                    if self.params['debugOverlaps']:
                        self.algebraicOverlaps.check(mesh=self.meshLevel.mesh,
                                                     dm=self.DoFMap,
                                                     label='algebraicOverlaps in \'{} {}\''.format(self.label, self.levelNo),
                                                     interfaces=self.meshLevel.meshOverlaps)
                elif self.meshLevel.interfaces is not None:
                    with self.Timer('Build algebraic overlaps of type \'{}\''.format(commType)):
                        self.algebraicOverlaps = self.meshLevel.interfaces.getDoFs(self.meshLevel.mesh, self.DoFMap, commType)
                    if self.params['debugOverlaps']:
                        self.algebraicOverlaps.check(mesh=self.meshLevel.mesh,
                                                     dm=self.DoFMap,
                                                     label='algebraicOverlaps in \'{} {}\''.format(self.label, self.levelNo),
                                                     interfaces=self.meshLevel.interfaces)

            if self.algebraicOverlaps is not None:
                self.inner = ip_distributed(self.algebraicOverlaps, 0)
                self.norm = norm_distributed(self.algebraicOverlaps, 0)
            else:
                self.inner = ip_serial()
                self.norm = norm_serial()
            if self.DoFMap is not None:
                self.DoFMap.set_ip_norm(self.inner, self.norm)

            if (buildType & RESTRICTION_PROLONGATION) and (self.previousLevel is not None):
                assert (self.previousLevel.DoFMap is not None) and (self.DoFMap is not None)
                # use reorder here, since reorder=False bugs out
                (self.R,
                 self.P) = buildRestrictionProlongation(self.previousLevel.DoFMap,
                                                        self.DoFMap)

    def buildCoarserMatrices(self):
        """
        Recursively build matrices on coarser levels
        """
        if self.previousLevel is not None:
            self.previousLevel.buildCoarserMatrices()

    def clean(self):
        if not self.params['keepAllDoFMaps'] and not self.previousLevel is None:
            self.DoFMap = None

    @classmethod
    def getKeys(cls):
        return ['P', 'R', 'DoFMap', 'algebraicOverlaps']

    def getLevelDict(self):
        lvl = {}
        for key in self.getKeys():
            if getattr(self, key) is not None:
                lvl[key] = getattr(self, key)
        return lvl

    @classmethod
    def fromLevelDict(cls, meshLevel, lvl):
        alvl = algebraicLevel(meshLevel, NO_BUILD)
        for key in cls.getKeys():
            if key in lvl:
                setattr(alvl, key, lvl[key])
        return alvl

    @property
    def accumulateOperator(self):
        if self.algebraicOverlaps is not None:
            return self.algebraicOverlaps.getAccumulateOperator()
        else:
            return None

class algebraicLevel(algebraicLevelBase):
    def __init__(self, meshLevel, buildType):
        self.A = None
        self.S = None
        self.D = None
        self.M = None
        self.surface_mass = None
        self.surface_stiffness = None
        super(algebraicLevel, self).__init__(meshLevel, buildType)

    def build(self, buildType):
        super(algebraicLevel, self).build(buildType)

        diffusivity = self.params['diffusivity']
        reaction = self.params['reaction']
        symmetric = self.params['symmetric']
        element = self.params['element']
        reorder = self.params['reorder']
        commType = self.params['commType']
        buildMass = self.params['buildMass'] or reaction is not None
        driftCoeff = self.params.get('driftCoeff', None)
        buildNeumann = self.params.get('buildNeumann', False)

        if buildType & SPARSITY_PATTERN:
            # set up sparsity patterns only
            DoFMap = self.DoFMap
            mesh = self.meshLevel.mesh
            self.fullyAssembled = False
            with self.Timer('Prepared sparsity patterns'):
                self.S = DoFMap.buildSparsityPattern(mesh.cells,
                                                     symmetric=symmetric,
                                                     reorder=reorder)
                if driftCoeff is not None:
                    self.D = self.S.copy()
                if buildMass:
                    self.M = self.S.copy()

        if buildType & ASSEMBLY:
            # fully build matrices
            DoFMap = self.DoFMap
            mesh = self.meshLevel.mesh
            self.fullyAssembled = True
            with self.Timer('Assembled matrices'):
                self.S = DoFMap.assembleStiffness(sss_format=symmetric,
                                                  reorder=reorder,
                                                  diffusivity=diffusivity)
                if buildMass:
                    self.M = DoFMap.assembleMass(sss_format=symmetric,
                                                 reorder=reorder)
                if driftCoeff is not None:
                    self.D = assembleDrift(mesh,
                                           DoFMap,
                                           driftCoeff)
                if buildNeumann:
                    self.neumannA = self.DoFMapNeumann.assembleStiffness(sss_format=symmetric,
                                                                         reorder=reorder,
                                                                         diffusivity=diffusivity)
                if isinstance(reaction, (float, REAL)):
                    self.A = self.S.copy()
                    for j in range(self.A.data.shape[0]):
                        self.A.data[j] += reaction*self.M.data[j]
                        if isinstance(self.A, SSS_LinearOperator):
                            for j in range(self.A.num_rows):
                                self.A.diagonal[j] += reaction*self.M.diagonal[j]
                elif isinstance(reaction, function):
                    self.A = self.S.copy()
                    dm = self.DoFMap
                    c = dm.interpolate(reaction)
                    for k in range(dm.num_dofs):
                        for j in range(self.A.indptr[k], self.A.indptr[k+1]):
                            self.A.data[j] += c[k]*self.M.data[j]
                    if isinstance(self.A, SSS_LinearOperator):
                        for k in range(self.A.num_rows):
                            self.A.diagonal[k] += c[k]*self.M.diagonal[k]
                elif reaction is None:
                    self.A = self.S
                else:
                    raise NotImplementedError()

            # surface mass matrix
            if self.isLastLevel and self.params['buildSurfaceMass']:
                with self.Timer('Build surface mass matrix'):
                    if self.params['depth'] > 0:
                        surface = mesh.get_surface_mesh(INTERIOR)
                    else:
                        surface = mesh.get_surface_mesh(INTERIOR_NONOVERLAPPING)
                    from PyNucleus_fem import assembleSurfaceMass
                    self.surface_mass = assembleSurfaceMass(mesh, surface,
                                                            self.DoFMap,
                                                            sss_format=symmetric,
                                                            reorder=reorder)
                    # ToDo: Don't just copy the sparsity pattern, this is a big waste of memory
                    # data = np.zeros((self.A.nnz), dtype=REAL)
                    # if symmetric:
                    #     diagonal = np.zeros(self.A.shape[0], dtype=REAL)
                    #     M = SSS_LinearOperator(self.A.indices, self.A.indptr, data, diagonal)
                    # else:
                    #     M = CSR_LinearOperator(self.A.indices, self.A.indptr, data)
                    # if element == 'P1':
                    #     dmS = P1_DoFMap(mesh, [PHYSICAL])
                    #     dmS.cells = surface.cells
                    # elif element == 'P2':
                    #     assert False, "Surface mass matrix not implemented for P2."
                    #     dmS = P2_DoFMap(mesh, [PHYSICAL])
                    #     cellOrig = dmS.mesh.cells
                    #     dmS.mesh.cells = surface.cells
                    # if mesh.dim == 1 and element == 'P1':
                    #     dmS.dofs_per_element = 1
                    #     self.surface_mass = assembleMatrix(surface, dmS, mass_0d_in_1d_sym_P1(), A=M,
                    #                                       sss_format=symmetric, reorder=reorder)
                    # elif mesh.dim == 2 and element == 'P1':
                    #     dmS.dofs_per_element = 2
                    #     self.surface_mass = assembleMatrix(surface, dmS, mass_1d_in_2d_sym_P1(), A=M,
                    #                                        sss_format=symmetric, reorder=reorder)
                    # elif mesh.dim == 2 and element == 'P2':
                    #     dmS.dofs_per_element = 3
                    #     self.surface_mass = assembleMatrix(surface, dmS, mass_1d_in_2d_sym_P2(), A=M,
                    #                                       sss_format=symmetric, reorder=reorder)
                    #     dmS.mesh.cells = cellOrig
                    # else:
                    #     raise NotImplementedError()

            # surface stiffness matrix
            if self.isLastLevel and self.params['buildSurfaceStiffness']:
                with self.Timer('Build surface stiffness matrix'):
                    if self.params['depth'] > 0:
                        surface = mesh.get_surface_mesh(INTERIOR)
                    else:
                        surface = mesh.get_surface_mesh(INTERIOR_NONOVERLAPPING)
                    # ToDo: Don't just copy the sparsity pattern, this is a big waste of memory
                    data = np.zeros((self.A.nnz), dtype=REAL)
                    if symmetric:
                        diagonal = np.zeros(self.A.shape[0], dtype=REAL)
                        AS = SSS_LinearOperator(self.A.indices, self.A.indptr, data, diagonal)
                    else:
                        AS = CSR_LinearOperator(self.A.indices, self.A.indptr, data)
                    assert element == 'P1', "Surface stiffness matrix only implemented for P1"
                    dmS = P1_DoFMap(mesh, [PHYSICAL])
                    dmS.cells = surface.cells
                    if mesh.dim == 2:
                        dmS.dofs_per_element = 2
                        self.surfaceStiffness = assembleMatrix(surface, dmS, stiffness_1d_in_2d_sym(), A=AS,
                                                               sss_format=symmetric, reorder=reorder)
                    else:
                        raise NotImplementedError()

        if reorder and buildType & OVERLAPS:
            # build algebraic overlaps
            if self.meshLevel.meshOverlaps is not None:
                with self.Timer('Build algebraic overlaps of type \'{}\''.format(commType)):
                    self.algebraicOverlaps = self.meshLevel.meshOverlaps.getDoFs(self.meshLevel.mesh, self.DoFMap, commType,
                                                                                allowInteriorBoundary=self.params['interiorBC'] == 'homogeneousNeumann' or not self.isLastLevel)
                if self.params['debugOverlaps']:
                    self.algebraicOverlaps.check(mesh=self.meshLevel.mesh,
                                                 dm=self.DoFMap,
                                                 label='algebraicOverlaps in \'{} {}\''.format(self.label, self.levelNo),
                                                 interfaces=self.meshLevel.meshOverlaps)
            elif self.meshLevel.interfaces is not None:
                with self.Timer('Build algebraic overlaps of type \'{}\''.format(commType)):
                    self.algebraicOverlaps = self.meshLevel.interfaces.getDoFs(self.meshLevel.mesh, self.DoFMap, commType)
                if self.params['debugOverlaps']:
                    self.algebraicOverlaps.check(mesh=self.meshLevel.mesh,
                                                 dm=self.DoFMap,
                                                 label='algebraicOverlaps in \'{} {}\''.format(self.label, self.levelNo),
                                                 interfaces=self.meshLevel.interfaces)

        if reorder and (buildType & RESTRICTION_PROLONGATION) and (self.previousLevel is not None):
            assert (self.previousLevel.DoFMap is not None) and (self.DoFMap is not None)
            # use reorder here, since reorder=False bugs out
            (self.R,
             self.P) = buildRestrictionProlongation(self.previousLevel.DoFMap,
                                                    self.DoFMap)

    def buildCoarserMatrices(self):
        """
        Recursively build matrices on coarser levels
        """
        if self.previousLevel is None:
            return
        if self.S is not None and self.P is not None and self.previousLevel.S is not None and not self.previousLevel.fullyAssembled:
            assert self.P.shape[0] == self.S.shape[0], (self.R.shape[1], self.S.shape[0])
            assert self.P.shape[1] == self.previousLevel.S.shape[0]
            with self.Timer('Restrict stiffness matrix'):
                self.P.restrictMatrix(self.S, self.previousLevel.S)
            if self.previousLevel.A is None:
                self.previousLevel.A = self.previousLevel.S
        if self.D is not None and self.P is not None and self.previousLevel.D is not None and not self.previousLevel.fullyAssembled:
            assert self.P.shape[0] == self.D.shape[0]
            assert self.P.shape[1] == self.previousLevel.D.shape[0]
            with self.Timer('Restrict drift matrix'):
                self.P.restrictMatrix(self.D, self.previousLevel.D)
        if self.M is not None and self.P is not None and self.previousLevel.M is not None and not self.previousLevel.fullyAssembled:
            assert self.P.shape[0] == self.M.shape[0]
            assert self.P.shape[1] == self.previousLevel.M.shape[0]
            with self.Timer('Restrict mass matrix'):
                self.P.restrictMatrix(self.M, self.previousLevel.M)
        if self.M is not None and self.A is not None and self.R is not None and self.previousLevel.A is not None and self.previousLevel.M is not None:
            reaction = self.params['reaction']
            if isinstance(reaction, (float, REAL)):
                for j in range(self.previousLevel.A.data.shape[0]):
                    self.previousLevel.A.data[j] += reaction*self.previousLevel.M.data[j]
                    if isinstance(self.previousLevel.A, SSS_LinearOperator):
                        for j in range(self.previousLevel.A.num_rows):
                            self.previousLevel.A.diagonal[j] += reaction*self.previousLevel.M.diagonal[j]
            elif isinstance(reaction, function):
                dm = self.previousLevel.DoFMap
                c = dm.interpolate(reaction)
                for k in range(dm.num_dofs):
                    for j in range(self.previousLevel.A.indptr[k], self.previousLevel.A.indptr[k+1]):
                        self.previousLevel.A.data[j] += c[k]*self.previousLevel.M.data[j]
                if isinstance(self.previousLevel.A, SSS_LinearOperator):
                    for k in range(self.previousLevel.A.num_rows):
                        self.previousLevel.A.diagonal[k] += c[k]*self.previousLevel.M.diagonal[k]
            elif reaction is None:
                pass
            else:
                raise NotImplementedError()
        if self.previousLevel is not None:
            self.previousLevel.fullyAssembled = True
            self.previousLevel.buildCoarserMatrices()

    @classmethod
    def getKeys(cls):
        return algebraicLevelBase.getKeys() + ['A', 'S', 'D', 'M', 'surface_mass', 'surface_stiffness']

    def getLevelDict(self):
        lvl = super(algebraicLevel, self).getLevelDict()
        if hasattr(self, ' neumannA'):
            lvl['neumannA'] = self.neumannA
        return lvl

    def getGlobalA(self, doDistribute=False, keepDistributedResult=False):
        if self.A is not None:
            if self.algebraicOverlaps is not None:
                if isinstance(self.A, CSR_LinearOperator):
                    return CSR_DistributedLinearOperator(self.A, self.algebraicOverlaps,
                                                         doDistribute=doDistribute,
                                                         keepDistributedResult=keepDistributedResult)
                else:
                    return DistributedLinearOperator(self.A, self.algebraicOverlaps,
                                                     doDistribute=doDistribute,
                                                     keepDistributedResult=keepDistributedResult)
            else:
                return self.A
        else:
            return None


# Copyright 2018-2021 Xanadu Quantum Technologies Inc.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Multiple dispatch functions"""
# pylint: disable=import-outside-toplevel,too-many-return-statements
import warnings
from collections.abc import Sequence
import functools

from autograd.numpy.numpy_boxes import ArrayBox
from autoray import numpy as np
from numpy import ndarray

from . import single_dispatch  # pylint:disable=unused-import
from .utils import cast, get_interface, requires_grad

def _multi_dispatch(values):
    """Determines the correct framework to dispatch to given a
    sequence of tensor-like objects.

    Args:
        values (Sequence[tensor_like]): a sequence of tensor like objects

    Returns:
        str: the name of the interface

    To determine the framework to dispatch to, the following rules
    are applied:

    * Tensors that are incompatible (such as Torch and TensorFlow tensors)
      cannot both be present.

    * Autograd tensors *may* be present alongside Torch and TensorFlow tensors,
      but Torch and TensorFlow take precendence; the autograd arrays will
      be treated as non-differentiable NumPy arrays. A warning will be raised
      suggesting that vanilla NumPy be used instead.

    * Vanilla NumPy arrays and SciPy sparse matrices can be used alongside other tensor objects;
      they will always be treated as non-differentiable constants.
    """
    if "resource_variable" in getattr(values, "__module__", tuple()):
        values = np.asarray(values)

    interfaces = {get_interface(v) for v in values}

    if len(set(interfaces) - {"numpy", "scipy", "autograd"}) > 1:
        # contains multiple non-autograd interfaces
        raise ValueError("Tensors contain mixed types; cannot determine dispatch library")

    non_numpy_scipy_interfaces = set(interfaces) - {"numpy", "scipy"}

    if len(non_numpy_scipy_interfaces) > 1:
        # contains autograd and another interface
        warnings.warn(
            f"Contains tensors of types {non_numpy_scipy_interfaces}; dispatch will prioritize "
            "TensorFlow and PyTorch over autograd. Consider replacing Autograd with vanilla NumPy.",
            UserWarning,
        )

    if "tensorflow" in interfaces:
        return "tensorflow"

    if "torch" in interfaces:
        return "torch"

    if "autograd" in interfaces:
        return "autograd"

    if "jax" in interfaces:
        return "jax"

    return "numpy"

def multi_dispatch(argnum=None, tensor_list=None):
    r"""Decorater to dispatch arguments handled by the interface.

    This helps simplify definitions of new functions inside PennyLane. We can
    decorate the function, indicating the arguments that are tensors handled
    by the interface:

    >>> @qml.math.multi_dispatch(argnum=[0, 1])
    ... def some_function(tensor1, tensor2, option, like):
    ...     # the interface string is stored in `like`.
    ...     ...

    Args:
        argnum (list[int]): A list of integers indicating indicating the indices
            to dispatch (i.e., the arguments that are tensors handled by an interface).
            If ``None``, dispatch over all arguments.
        tensor_lists (list[int]): a list of integers indicating which indices
            in ``argnum`` are expected to be lists of tensors. If an argument
            marked as tensor list is not a ``tuple`` or ``list``, it is treated
            as if it was not marked as tensor list. If ``None``, this option is ignored.

    Returns:
        func: A wrapped version of the function, which will automatically attempt
        to dispatch to the correct autodifferentiation framework for the requested
        arguments. Note that the ``like`` argument will be optional, but can be provided
        if an explicit override is needed.

    .. seealso:: :func:`pennylane.math.multi_dispatch._multi_dispatch`

    .. note::
        This decorator makes the interface argument "like" optional as it utilizes
        the utility function `_multi_dispatch` to automatically detect the appropriate
        interface based on the tensor types.

    **Examples**

    We can redefine external functions to be suitable for PennyLane. Here, we
    redefine Autoray's ``stack`` function.

    >>> stack = multi_dispatch(argnum=0, tensor_list=0)(autoray.numpy.stack)

    We can also use the ``multi_dispatch`` decorator to dispatch
    arguments of more more elaborate custom functions. Here is an example
    of a ``custom_function`` that
    computes :math:`c \\sum_i (v_i)^T v_i`, where :math:`v_i` are vectors in ``values`` and
    :math:`c` is a fixed ``coefficient``. Note how ``argnum=0`` only points to the first argument ``values``,
    how ``tensor_list=0`` indicates that said first argument is a list of vectors, and that ``coefficient`` is not
    dispatched.

    >>> @math.multi_dispatch(argnum=0, tensor_list=0)
    >>> def custom_function(values, like, coefficient=10):
    >>>     # values is a list of vectors
    >>>     # like can force the interface (optional)
    >>>     if like == "tensorflow":
    >>>         # add interface-specific handling if necessary
    >>>     return coefficient * np.sum([math.dot(v,v) for v in values])

    We can then run

    >>> values = [np.array([1, 2, 3]) for _ in range(5)]
    >>> custom_function(values)
    700

    """

    def decorator(fn):
        @functools.wraps(fn)
        def wrapper(*args, **kwargs):
            argnums = argnum if argnum is not None else list(range(len(args)))
            tensor_lists = tensor_list if tensor_list is not None else []

            if not isinstance(argnums, Sequence):
                argnums = [argnums]
            if not isinstance(tensor_lists, Sequence):
                tensor_lists = [tensor_lists]

            dispatch_args = []

            for a in argnums:
                # Only use extend if the marked argument really
                # is a (native) python Sequence
                if a in tensor_lists and isinstance(args[a], (list, tuple)):
                    dispatch_args.extend(args[a])
                else:
                    dispatch_args.append(args[a])

            interface = kwargs.pop("like", None)
            interface = interface or _multi_dispatch(dispatch_args)
            kwargs["like"] = interface

            return fn(*args, **kwargs)

        return wrapper

    return decorator

@multi_dispatch(argnum=[0], tensor_list=[0])
def block_diag(values, like=None):
    """Combine a sequence of 2D tensors to form a block diagonal tensor.

    Args:
        values (Sequence[tensor_like]): Sequence of 2D arrays/tensors to form
            the block diagonal tensor.

    Returns:
        tensor_like: the block diagonal tensor

    **Example**

    >>> t = [
    ...     np.array([[1, 2], [3, 4]]),
    ...     torch.tensor([[1, 2, 3], [-1, -6, -3]]),
    ...     torch.tensor(5)
    ... ]
    >>> qml.math.block_diag(t)
    tensor([[ 1,  2,  0,  0,  0,  0],
            [ 3,  4,  0,  0,  0,  0],
            [ 0,  0,  1,  2,  3,  0],
            [ 0,  0, -1, -6, -3,  0],
            [ 0,  0,  0,  0,  0,  5]])
    """
    values = np.coerce(values, like=like)
    return np.block_diag(values, like=like)

@multi_dispatch(argnum=[0], tensor_list=[0])
def concatenate(values, axis=0, like=None):
    """Concatenate a sequence of tensors along the specified axis.

    .. warning::

        Tensors that are incompatible (such as Torch and TensorFlow tensors)
        cannot both be present.

    Args:
        values (Sequence[tensor_like]): Sequence of tensor-like objects to
            concatenate. The objects must have the same shape, except in the dimension corresponding
            to axis (the first, by default).
        axis (int): The axis along which the input tensors are concatenated. If axis is None,
            tensors are flattened before use. Default is 0.

    Returns:
        tensor_like: The concatenated tensor.

    **Example**

    >>> x = tf.constant([0.6, 0.1, 0.6])
    >>> y = tf.Variable([0.1, 0.2, 0.3])
    >>> z = np.array([5., 8., 101.])
    >>> concatenate([x, y, z])
    <tf.Tensor: shape=(3, 3), dtype=float32, numpy=
    array([6.00e-01, 1.00e-01, 6.00e-01, 1.00e-01, 2.00e-01, 3.00e-01, 5.00e+00, 8.00e+00, 1.01e+02], dtype=float32)>
    """

    if like == "torch":
        import torch

        device = (
            "cuda"
            if any(t.device.type == "cuda" for t in values if isinstance(t, torch.Tensor))
            else "cpu"
        )

        if axis is None:
            # flatten and then concatenate zero'th dimension
            # to reproduce numpy's behaviour
            values = [
                np.flatten(torch.as_tensor(t, device=torch.device(device)))  # pragma: no cover
                for t in values
            ]
            axis = 0
        else:
            values = [
                torch.as_tensor(t, device=torch.device(device)) for t in values  # pragma: no cover
            ]

    if like == "tensorflow" and axis is None:
        # flatten and then concatenate zero'th dimension
        # to reproduce numpy's behaviour
        values = [np.flatten(np.array(t)) for t in values]
        axis = 0

    return np.concatenate(values, axis=axis, like=like)

@multi_dispatch(argnum=[0], tensor_list=[0])
def diag(values, k=0, like=None):
    """Construct a diagonal tensor from a list of scalars.

    Args:
        values (tensor_like or Sequence[scalar]): sequence of numeric values that
            make up the diagonal
        k (int): The diagonal in question. ``k=0`` corresponds to the main diagonal.
            Use ``k>0`` for diagonals above the main diagonal, and ``k<0`` for
            diagonals below the main diagonal.

    Returns:
        tensor_like: the 2D diagonal tensor

    **Example**

    >>> x = [1., 2., tf.Variable(3.)]
    >>> qml.math.diag(x)
    <tf.Tensor: shape=(3, 3), dtype=float32, numpy=
    array([[1., 0., 0.],
           [0., 2., 0.],
           [0., 0., 3.]], dtype=float32)>
    >>> y = tf.Variable([0.65, 0.2, 0.1])
    >>> qml.math.diag(y, k=-1)
    <tf.Tensor: shape=(4, 4), dtype=float32, numpy=
    array([[0.  , 0.  , 0.  , 0.  ],
           [0.65, 0.  , 0.  , 0.  ],
           [0.  , 0.2 , 0.  , 0.  ],
           [0.  , 0.  , 0.1 , 0.  ]], dtype=float32)>
    >>> z = torch.tensor([0.1, 0.2])
    >>> qml.math.diag(z, k=1)
    tensor([[0.0000, 0.1000, 0.0000],
            [0.0000, 0.0000, 0.2000],
            [0.0000, 0.0000, 0.0000]])
    """
    if isinstance(values, (list, tuple)):
        values = np.stack(np.coerce(values, like=like), like=like)

    return np.diag(values, k=k, like=like)

@multi_dispatch(argnum=[0, 1])
def dot(tensor1, tensor2, like=None):
    """Returns the matrix or dot product of two tensors.

    * If both tensors are 0-dimensional, elementwise multiplication
      is performed and a 0-dimensional scalar returned.

    * If both tensors are 1-dimensional, the dot product is returned.

    * If the first array is 2-dimensional and the second array 1-dimensional,
      the matrix-vector product is returned.

    * If both tensors are 2-dimensional, the matrix product is returned.

    * Finally, if the the first array is N-dimensional and the second array
      M-dimensional, a sum product over the last dimension of the first array,
      and the second-to-last dimension of the second array is returned.

    Args:
        tensor1 (tensor_like): input tensor
        tensor2 (tensor_like): input tensor

    Returns:
        tensor_like: the matrix or dot product of two tensors
    """
    x, y = np.coerce([tensor1, tensor2], like=like)

    if like == "torch":
        if x.ndim == 0 and y.ndim == 0:
            return x * y

        if x.ndim <= 2 and y.ndim <= 2:
            return x @ y

        return np.tensordot(x, y, axes=[[-1], [-2]], like=like)

    if like == "tensorflow":
        if len(np.shape(x)) == 0 and len(np.shape(y)) == 0:
            return x * y

        if len(np.shape(y)) == 1:
            return np.tensordot(x, y, axes=[[-1], [0]], like=like)

        if len(np.shape(x)) == 2 and len(np.shape(y)) == 2:
            return x @ y

        return np.tensordot(x, y, axes=[[-1], [-2]], like=like)

    return np.dot(x, y, like=like)

@multi_dispatch(argnum=[0, 1])
def tensordot(tensor1, tensor2, axes=None, like=None):
    """Returns the tensor product of two tensors.
    In general ``axes`` specifies either the set of axes for both
    tensors that are contracted (with the first/second entry of ``axes``
    giving all axis indices for the first/second tensor) or --- if it is
    an integer --- the number of last/first axes of the first/second
    tensor to contract over.
    There are some non-obvious special cases:

    * If both tensors are 0-dimensional, ``axes`` must be 0.
      and a 0-dimensional scalar is returned containing the simple product.

    * If both tensors are 1-dimensional and ``axes=0``, the outer product
      is returned.

    * Products between a non-0-dimensional and a 0-dimensional tensor are not
      supported in all interfaces.

    Args:
        tensor1 (tensor_like): input tensor
        tensor2 (tensor_like): input tensor
        axes (int or list[list[int]]): Axes to contract over, see detail description.

    Returns:
        tensor_like: the tensor product of the two input tensors
    """
    tensor1, tensor2 = np.coerce([tensor1, tensor2], like=like)
    return np.tensordot(tensor1, tensor2, axes=axes, like=like)

@multi_dispatch(argnum=[0], tensor_list=[0])
def get_trainable_indices(values, like=None):
    """Returns a set containing the trainable indices of a sequence of
    values.

    Args:
        values (Iterable[tensor_like]): Sequence of tensor-like objects to inspect

    Returns:
        set[int]: Set containing the indices of the trainable tensor-like objects
        within the input sequence.

    **Example**

    >>> def cost_fn(params):
    ...     print("Trainable:", qml.math.get_trainable_indices(params))
    ...     return np.sum(np.sin(params[0] * params[1]))
    >>> values = [np.array([0.1, 0.2], requires_grad=True),
    ... np.array([0.5, 0.2], requires_grad=False)]
    >>> cost_fn(values)
    Trainable: {0}
    tensor(0.0899685, requires_grad=True)
    """
    trainable = requires_grad
    trainable_params = set()

    if like == "jax":
        import jax

        if not any(isinstance(v, jax.core.Tracer) for v in values):
            # No JAX tracing is occuring; treat all `DeviceArray` objects as trainable.
            trainable = lambda p, **kwargs: isinstance(p, jax.numpy.DeviceArray)
        else:
            # JAX tracing is occuring; use the default behaviour (only traced arrays
            # are treated as trainable). This is required to ensure that `jax.grad(func, argnums=...)
            # works correctly, as the argnums argnument determines which parameters are
            # traced arrays.
            trainable = requires_grad

    for idx, p in enumerate(values):
        if trainable(p, interface=like):
            trainable_params.add(idx)

    return trainable_params

def ones_like(tensor, dtype=None):
    """Returns a tensor of all ones with the same shape and dtype
    as the input tensor.

    Args:
        tensor (tensor_like): input tensor
        dtype (str, np.dtype, None): The desired output datatype of the array. If not provided, the dtype of
            ``tensor`` is used. This argument can be any supported NumPy dtype representation, including
            a string (``"float64"``), a ``np.dtype`` object (``np.dtype("float64")``), or
            a dtype class (``np.float64``). If ``tensor`` is not a NumPy array, the
            **equivalent** dtype in the dispatched framework is used.

    Returns:
        tensor_like: an all-ones tensor with the same shape and
        size as ``tensor``

    **Example**

    >>> x = torch.tensor([1., 2.])
    >>> ones_like(x)
    tensor([1, 1])
    >>> y = tf.Variable([[0], [5]])
    >>> ones_like(y, dtype=np.complex128)
    <tf.Tensor: shape=(2, 1), dtype=complex128, numpy=
    array([[1.+0.j],
           [1.+0.j]])>
    """
    if dtype is not None:
        return cast(np.ones_like(tensor), dtype)

    return np.ones_like(tensor)

@multi_dispatch(argnum=[0], tensor_list=[0])
def stack(values, axis=0, like=None):
    """Stack a sequence of tensors along the specified axis.

    .. warning::

        Tensors that are incompatible (such as Torch and TensorFlow tensors)
        cannot both be present.

    Args:
        values (Sequence[tensor_like]): Sequence of tensor-like objects to
            stack. Each object in the sequence must have the same size in the given axis.
        axis (int): The axis along which the input tensors are stacked. ``axis=0`` corresponds
            to vertical stacking.

    Returns:
        tensor_like: The stacked array. The stacked array will have one additional dimension
        compared to the unstacked tensors.

    **Example**

    >>> x = tf.constant([0.6, 0.1, 0.6])
    >>> y = tf.Variable([0.1, 0.2, 0.3])
    >>> z = np.array([5., 8., 101.])
    >>> stack([x, y, z])
    <tf.Tensor: shape=(3, 3), dtype=float32, numpy=
    array([[6.00e-01, 1.00e-01, 6.00e-01],
           [1.00e-01, 2.00e-01, 3.00e-01],
           [5.00e+00, 8.00e+00, 1.01e+02]], dtype=float32)>
    """
    values = np.coerce(values, like=like)
    return np.stack(values, axis=axis, like=like)

def where(condition, x=None, y=None):
    """Returns elements chosen from x or y depending on a boolean tensor condition,
    or the indices of entries satisfying the condition.

    The input tensors ``condition``, ``x``, and ``y`` must all be broadcastable to the same shape.

    Args:
        condition (tensor_like[bool]): A boolean tensor. Where ``True`` , elements from
            ``x`` will be chosen, otherwise ``y``. If ``x`` and ``y`` are ``None`` the
            indices where ``condition==True`` holds will be returned.
        x (tensor_like): values from which to choose if the condition evaluates to ``True``
        y (tensor_like): values from which to choose if the condition evaluates to ``False``

    Returns:
        tensor_like or tuple[tensor_like]: If ``x is None`` and ``y is None``, a tensor
        or tuple of tensors with the indices where ``condition`` is ``True`` .
        Else, a tensor with elements from ``x`` where the ``condition`` is ``True``,
        and ``y`` otherwise. In this case, the output tensor has the same shape as
        the input tensors.

    **Example with three arguments**

    >>> a = torch.tensor([0.6, 0.23, 0.7, 1.5, 1.7], requires_grad=True)
    >>> b = torch.tensor([-1., -2., -3., -4., -5.], requires_grad=True)
    >>> math.where(a < 1, a, b)
    tensor([ 0.6000,  0.2300,  0.7000, -4.0000, -5.0000], grad_fn=<SWhereBackward>)

    .. warning::

        The output format for ``x=None`` and ``y=None`` follows the respective
        interface and differs between TensorFlow and all other interfaces:
        For TensorFlow, the output is a tensor with shape
        ``(num_true, len(condition.shape))`` where ``num_true`` is the number
        of entries in ``condition`` that are ``True`` .
        The entry at position ``(i, j)`` is the ``j`` th entry of the ``i`` th
        index.
        For all other interfaces, the output is a tuple of tensor-like objects,
        with the ``j`` th object indicating the ``j`` th entries of all indices.
        Also see the examples below.

    **Example with single argument**

    For Torch, Autograd, JAX and NumPy, the output formatting is as follows:

    >>> a = [[0.6, 0.23, 1.7],[1.5, 0.7, -0.2]]
    >>> math.where(torch.tensor(a) < 1)
    (tensor([0, 0, 1, 1]), tensor([0, 1, 1, 2]))

    This is not a single tensor-like object but corresponds to the shape
    ``(2, 4)`` . For TensorFlow, on the other hand:

    >>> math.where(tf.constant(a) < 1)
    tf.Tensor(
    [[0 0]
     [0 1]
     [1 1]
     [1 2]], shape=(4, 2), dtype=int64)

    As we can see, the dimensions are swapped and the output is a single Tensor.
    Note that the number of dimensions of the output does *not* depend on the input
    shape, it is always two-dimensional.

    """
    if x is None and y is None:
        interface = _multi_dispatch([condition])
        return np.where(condition, like=interface)

    interface = _multi_dispatch([condition, x, y])
    res = np.where(condition, x, y, like=interface)

    if interface == "tensorflow":
        return np.transpose(np.stack(res))

    return res

@multi_dispatch(argnum=[0, 1])
def frobenius_inner_product(A, B, normalize=False, like=None):
    r"""Frobenius inner product between two matrices.

    .. math::

        \langle A, B \rangle_F = \sum_{i,j=1}^n A_{ij} B_{ij} = \operatorname{tr} (A^T B)

    The Frobenius inner product is equivalent to the Hilbert-Schmidt inner product for
    matrices with real-valued entries.

    Args:
        A (tensor_like[float]): First matrix, assumed to be a square array.
        B (tensor_like[float]): Second matrix, assumed to be a square array.
        normalize (bool): If True, divide the inner product by the Frobenius norms of A and B.

    Returns:
        float: Frobenius inner product of A and B

    **Example**

    >>> A = np.random.random((3,3))
    >>> B = np.random.random((3,3))
    >>> qml.math.frobenius_inner_product(A, B)
    3.091948202943376
    """
    A, B = np.coerce([A, B], like=like)

    inner_product = np.sum(A * B)

    if normalize:
        norm = np.sqrt(np.sum(A * A) * np.sum(B * B))
        inner_product = inner_product / norm

    return inner_product

@multi_dispatch(argnum=[0, 2])
def scatter_element_add(tensor, index, value, like=None):
    """In-place addition of a multidimensional value over various
    indices of a tensor.

    Args:
        tensor (tensor_like[float]): Tensor to add the value to
        index (tuple or list[tuple]): Indices to which to add the value
        value (float or tensor_like[float]): Value to add to ``tensor``
        like (str): Manually chosen interface to dispatch to.
    Returns:
        tensor_like[float]: The tensor with the value added at the given indices.

    **Example**

    >>> tensor = torch.tensor([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]])
    >>> index = (1, 2)
    >>> value = -3.1
    >>> qml.math.scatter_element_add(tensor, index, value)
    tensor([[ 0.1000,  0.2000,  0.3000],
            [ 0.4000,  0.5000, -2.5000]])

    If multiple indices are given, in the form of a list of tuples, the
    ``k`` th tuple is interpreted to contain the ``k`` th entry of all indices:

    >>> indices = [(1, 0), (2, 1)] # This will modify the entries (1, 2) and (0, 1)
    >>> values = torch.tensor([10, 20])
    >>> qml.math.scatter_element_add(tensor, indices, values)
    tensor([[ 0.1000, 20.2000,  0.3000],
            [ 0.4000,  0.5000, 10.6000]])
    """
    if len(np.shape(tensor)) == 0 and index == ():
        return tensor + value

    return np.scatter_element_add(tensor, index, value, like=like)

def unwrap(values, max_depth=None):
    """Unwrap a sequence of objects to NumPy arrays.

    Note that tensors on GPUs will automatically be copied
    to the CPU.

    Args:
        values (Sequence[tensor_like]): sequence of tensor-like objects to unwrap
        max_depth (int): Positive integer indicating the depth of unwrapping to perform
            for nested tensor-objects. This argument only applies when unwrapping
            Autograd ``ArrayBox`` objects.

    **Example**

    >>> values = [np.array([0.1, 0.2]), torch.tensor(0.1, dtype=torch.float64), torch.tensor([0.5, 0.2])]
    >>> math.unwrap(values)
    [array([0.1, 0.2]), 0.1, array([0.5, 0.2], dtype=float32)]

    This function will continue to work during backpropagation:

    >>> def cost_fn(params):
    ...     unwrapped_params = math.unwrap(params)
    ...     print("Unwrapped:", [(i, type(i)) for i in unwrapped_params])
    ...     return np.sum(np.sin(params))
    >>> params = np.array([0.1, 0.2, 0.3])
    >>> grad = autograd.grad(cost_fn)(params)
    Unwrapped: [(0.1, <class 'float'>), (0.2, <class 'float'>), (0.3, <class 'float'>)]
    >>> print(grad)
    [0.99500417 0.98006658 0.95533649]
    """
    res = []

    for t in values:
        if isinstance(t, ArrayBox):
            a = np.to_numpy(t, max_depth=max_depth)
        else:
            a = np.to_numpy(t)

        if isinstance(a, ndarray) and not a.shape:
            # if NumPy array is scalar, convert to a Python float
            res.append(a.tolist())
        else:
            res.append(a)

    return res


#!/usr/bin/env python
# encoding: utf-8

'''
Generation of waterfall figures
'''

__author__      = 'Jake Retallick'
__copyright__   = 'Apache License 2.0'
__version__     = '1.2'
__date__        = '2018-04-17'  # last update

from hopper import HoppingModel
import os

import numpy as np
from itertools import product

import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap

class Waterfall:

    log_fn = os.path.join('.', '.temp', 'waterfall.log')
    img_fn = os.path.join('.', 'img', 'waterfall.pdf')

    ppnm = 50   # pixels per nm
    sig = .8    # 'atom diameter'

    lo, hi = 1., 4.     # amplitudes of unoccupied/occupied DBs
    portrait = True     # portrait style plotting
    xticks = False

    # coloring for plots
    cdict = {'red':    ((0, .48, .48),
                        (.5, .14, .14),
                        (1., 0, 0)),
             'green':  ((0,.82,.82),
                        (.5,.55,.55),
                        (1., 0, 0)),
             'blue':   ((0,.89,.89),
                        (.5,.70,.70),
                        (1., 0, 0))
             }

    cm = LinearSegmentedColormap('cm', cdict)

    def __init__(self, hopper):
        '''Initialise a waterfall generator. For now, assumes the
        device being simulated is linear in the x direction'''

        self.hopper = hopper
        self.X = .1*self.hopper.a*self.hopper.X
        self.N = self.hopper.N

    def generate(self, nscans=100, srate=10.0, pad=1.0, mu=.25):
        '''Generate the waterfall image.

        inputs:
            nscans  : number of line scans
            srate   : tip scan rate, in nm/s
            pad     : padding on either side of device, in nm
            mu      : Ef-DB- difference, eV
        '''

        if not self.hopper.initialised:
            self.hopper.initialise()

        bulk = self.hopper.getChannel('bulk')
        if bulk is not None:
            bulk.mu = mu

        # burn
        self.hopper.burn(10, per=True)

        # determine run time of scan
        xlo, xhi = np.min(self.X)-pad, np.max(self.X)+pad

        self.T = nscans*(xhi-xlo)/srate

        # populate hopping log
        self.hopper.startLog(self.log_fn)
        self.hopper.run(self.T)
        self.hopper.endLog()

        # simulate scan process
        self._show(nscans, srate, xlo, xhi)

    def _parser(self):
        '''Parse the time and state information from the log file'''

        with open(self.log_fn, 'r') as fp:
            state = self.hopper.charge
            for line in fp:
                t, s = line.split(' :: ')
                t = float(t)
                s = format(int(s, base=16), '0{0}b'.format(self.N))
                state = [int(x) for x in s]
                yield t, state
            yield np.inf, state     # hold on last state

    def _show(self, nscans, srate, xlo, xhi):
        '''Generate the waterfall from the hopping log'''

        xx = np.linspace(xlo, xhi, 1+int((xhi-xlo)*self.ppnm))
        kernel = np.exp(-np.linspace(-2,2, 1+int(self.sig*self.ppnm))**2)
        nx = len(xx)

        # index array, nearest db at each position
        D = np.abs(xx.reshape(-1,1) - self.X.reshape(1,-1))
        ind = np.argmin(D, axis=1)

        # impulse array
        imp = np.zeros(xx.shape, dtype=float)
        for x in self.X:
            n, r = divmod((x-xlo)*self.ppnm, 1)
            imp[int(n):int(n)+2] = 1-r, r

        lo_val = np.convolve(imp, self.lo*kernel, 'same')
        hi_val = np.convolve(imp, self.hi*kernel, 'same')

        # generate waterfall
        self.state_gen = self._parser()
        self.t, self.state = 0., None   # cache

        n, dn, dt = 0, 1, (xhi-xlo)/(srate*nx)
        data = np.zeros([nscans, nx], dtype=float)
        for m in range(nscans):
            for _ in range(nx):
                data[m,n] = self._integrate(ind[n], dt)
                n += dn
            data[m,:] = lo_val + data[m,:]*(hi_val-lo_val)
            n, dn = n - dn, -dn

        self._plot_handler(xlo, xhi, data)

    def _plot_handler(self, xlo, xhi, data):
        '''Create the waterfall plot'''

        # formatting options
        FS = 20
        TFS = 18
        size, asp = 4, 1.9

        if self.portrait:
            fig, (ax, cax) = plt.subplots(nrows=2, figsize=(size, size*asp),
                        gridspec_kw={'height_ratios': [1, .02]})
            dat = data
            xext, yext = xhi-xlo, self.T/60
            xlab, tlab = ax.set_xlabel, ax.set_ylabel
            cor = 'horizontal'
            ax.xaxis.tick_top()
            ax.xaxis.set_label_position('top')
        else:
            fig, (ax, cax) = plt.subplots(ncols=2, figsize=(size*asp, size),
                        gridspec_kw={'width_ratios': [1, .05]})
            dat = data.T
            xext, yext = self.T/60, xhi-xlo
            xlab, tlab = ax.set_ylabel, ax.set_xlabel
            cor = 'vertical'

        im = ax.imshow(dat, interpolation='None', aspect='auto', cmap=self.cm,
                        extent=[0, xext, 0, yext])

        # ticks
        tick_fill = lambda f, ext: f([round(x*ext,1) for x in [0, .5, 1.]])
        if self.xticks:
            tick_fill(ax.set_xticks, xext)
            tick_fill(ax.set_yticks, yext)
            xlab('x (nm)', fontsize=FS)
            tlab('Time (min)', fontsize=FS)
        else:
            if self.portrait:
                ax.set_xticks([])
                tick_fill(ax.set_yticks, yext)
            else:
                tick_fill(ax.set_xticks, xext)
                ax.set_yticks([])
            tlab('Time (min)', fontsize=FS)
        ax.tick_params(axis='both', which='major', labelsize=TFS)

        # colorbar
        cbar = fig.colorbar(im, cax=cax, ticks=[0,self.lo,self.hi], orientation=cor)
        cbar.set_label('Charges', fontsize=FS)
        cbar.ax.tick_params(labelsize=TFS)

        if self.portrait:
            cbar.ax.set_xticklabels([0,1,2])
        else:
            cbar.ax.set_yticklabels([0,1,2])

        plt.tight_layout(pad=0, h_pad=0, w_pad=0)

        fname = self.img_fn
        direc = os.path.dirname(fname)
        if not os.path.isdir(direc):
            os.makedirs(direc)

        plt.savefig(fname, bbox_inches='tight')
        plt.show()

    def _integrate(self, n, dt):
        '''Integrate the charge of the n^th DB for dt seconds'''

        charge, norm = 0., 1./dt
        while dt > 0:
            if self.t > 0:
                t, state = self.t, self.state
                self.t = 0.
            else:
                t, state = next(self.state_gen)
            if t<dt:
                charge += t*state[n]
                dt -= t
            else:
                charge += dt*state[n]
                self.t, self.state = t-dt, list(state)
                dt = 0

        return charge*norm

if __name__ == '__main__':

    d1221 = [8,10,15,17]
    d1221.insert(0, d1221[0]-7)
    d1221.append(d1221[-1]+7)

    device = d1221

    model = HoppingModel(device, model='marcus')
    model.addChannel('bulk')
    waterfall = Waterfall(model)
    waterfall.generate(nscans=800, srate=8.9, pad = 1.6, mu=0.1)