# Copyright 2011-2013 Kwant authors.
#
# This file is part of Kwant.  It is subject to the license terms in the file
# LICENSE.rst found in the top-level directory of this distribution and at
# https://kwant-project.org/license.  A list of Kwant authors can be found in
# the file AUTHORS.rst at the top-level directory of this distribution and at
# https://kwant-project.org/authors.

import tempfile
import warnings
import itertools
import numpy as np
import tinyarray as ta
from math import cos, sin

try:
    from scipy.integrate import simpson
except ImportError:
    # TODO: remove this once we drop support for scipy < 1.6
    from scipy.integrate import simps as simpson

import scipy.stats
import pytest
import sys

import kwant
from kwant._common import ensure_rng

try:
    from mpl_toolkits import mplot3d
    import matplotlib

    # This check is the same as the one performed inside matplotlib.use.
    matplotlib_backend_chosen = 'matplotlib.backends' in sys.modules
    # If the user did not already choose a backend, then choose
    # the one with the least dependencies.
    if not matplotlib_backend_chosen:
        matplotlib.use('Agg')

    from matplotlib import pyplot  # pragma: no flakes
except ImportError:
    matplotlib_backend_chosen = False

from kwant import plotter, _plotter


@pytest.mark.skipif(not _plotter.mpl_available, reason="Matplotlib unavailable.")
def test_matplotlib_backend_unset():
    """Simply importing Kwant should not set the matplotlib backend."""
    assert matplotlib_backend_chosen is False


def test_importable_without_backends():
    prefix, sep, suffix = _plotter.__file__.rpartition('.')
    if suffix in ['pyc', 'pyo']:
        suffix = 'py'
    assert suffix == 'py'
    fname = sep.join((prefix, suffix))
    with open(fname, 'rb') as f:
        code = f.read()
    code = code.replace(b'from . import', b'from kwant import')
    code = code.replace(b'matplotlib', b'totalblimp')
    code = code.replace(b'plotly', b'plylot')

    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")
        exec(code)               # Trigger the warning.
        assert len(w) == 2
        assert issubclass(w[0].category, RuntimeWarning)
        assert issubclass(w[1].category, RuntimeWarning)
        assert "totalblimp is not available" in str(w[0].message)
        assert "plylot is not available" in str(w[1].message)


def syst_2d(W=3, r1=3, r2=8):
    a = 1
    t = 1.0
    lat = kwant.lattice.square(a, norbs=1)
    syst = kwant.Builder()

    def ring(pos):
        (x, y) = pos
        rsq = x ** 2 + y ** 2
        return r1 ** 2 < rsq < r2 ** 2

    syst[lat.shape(ring, (0, r1 + 1))] = 4 * t
    syst[lat.neighbors()] = -t
    sym_lead0 = kwant.TranslationalSymmetry(lat.vec((-1, 0)))
    lead0 = kwant.Builder(sym_lead0)

    lead_shape = lambda pos: (-W / 2 < pos[1] < W / 2)

    lead0[lat.shape(lead_shape, (0, 0))] = 4 * t
    lead0[lat.neighbors()] = - t
    lead1 = lead0.reversed()
    syst.attach_lead(lead0)
    syst.attach_lead(lead1)
    return syst


def syst_3d(W=3, r1=2, r2=4, a=1, t=1.0):
    lat = kwant.lattice.general(((a, 0, 0), (0, a, 0), (0, 0, a)), norbs=1)
    syst = kwant.Builder()

    def ring(pos):
        (x, y, z) = pos
        rsq = x ** 2 + y ** 2
        return (r1 ** 2 < rsq < r2 ** 2) and abs(z) < 2
    syst[lat.shape(ring, (0, -r2 + 1, 0))] = 4 * t
    syst[lat.neighbors()] = - t
    sym_lead0 = kwant.TranslationalSymmetry(lat.vec((-1, 0, 0)))
    lead0 = kwant.Builder(sym_lead0)

    lead_shape = lambda pos: (-W / 2 < pos[1] < W / 2) and abs(pos[2]) < 2

    lead0[lat.shape(lead_shape, (0, 0, 0))] = 4 * t
    lead0[lat.neighbors()] = - t
    lead1 = lead0.reversed()
    syst.attach_lead(lead0)
    syst.attach_lead(lead1)
    return syst


def plotter_file_suffix(engine):
    # We need this function so that we can add a .html suffix to the output filename.
    # This is required because plotly will throw an error if filename is without the suffix.
    if engine == "plotly":
        return ".html"
    else:
        return None


@pytest.mark.skipif(not _plotter.mpl_available, reason="Matplotlib unavailable.")
def test_matplotlib_plot():
    plotter.set_engine('matplotlib')
    plot = plotter.plot
    syst2d = syst_2d()
    syst3d = syst_3d()
    color_opts = ['k', (lambda site: site.tag[0]),
                  lambda site: (abs(site.tag[0] / 100),
                                abs(site.tag[1] / 100), 0)]
    engine = plotter.get_engine()
    with tempfile.NamedTemporaryFile('w+b', suffix=plotter_file_suffix(engine)) as out:
        out_filename = out.name
        for color in color_opts:
            for syst in (syst2d, syst3d):
                fig = plot(syst, site_color=color, cmap='binary', file=out_filename)
                if (color != 'k' and
                    isinstance(color(next(iter(syst2d.sites()))), float)):
                    assert fig.axes[0].collections[0].get_array() is not None
                assert len(fig.axes[0].collections) == 6
        color_opts = ['k', (lambda site, site2: site.tag[0]),
                      lambda site, site2: (abs(site.tag[0] / 100),
                                           abs(site.tag[1] / 100), 0)]
        for color in color_opts:
            for syst in (syst2d, syst3d):
                fig = plot(syst2d, hop_color=color, cmap='binary', file=out_filename,
                           fig_size=(2, 10), dpi=30)
                if color != 'k' and isinstance(color(next(iter(syst2d.sites())),
                                                          None), float):
                    assert fig.axes[0].collections[1].get_array() is not None

        assert isinstance(plot(syst3d, file=out_filename).axes[0], mplot3d.axes3d.Axes3D)

        syst2d.leads = []
        plot(syst2d, file=out_filename)
        del syst2d[list(syst2d.hoppings())]
        plot(syst2d, file=out_filename)

        plot(syst3d, file=out_filename)
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            plot(syst2d.finalized(), file=out_filename)

        # test 2D projections of 3D systems
        plot(syst3d, file=out_filename, pos_transform=lambda pos: pos[:2])


@pytest.mark.skipif(not _plotter.plotly_available, reason="Plotly unavailable.")
def test_plotly_plot():
    plotter.set_engine('plotly')
    plot = plotter.plot
    syst2d = syst_2d()
    syst3d = syst_3d()
    color_opts = ['black', (lambda site: site.tag[0]),
                  lambda site: (abs(site.tag[0] / 100),
                                abs(site.tag[1] / 100), 0)]
    engine = plotter.get_engine()
    with tempfile.NamedTemporaryFile('w+b', suffix=plotter_file_suffix(engine)) as out:
        out_filename = out.name
        for color in color_opts:
            for syst in (syst2d, syst3d):
                plot(syst, site_color=color, cmap='binary', file=out_filename, show=False)

        color_opts = ['black', (lambda site, site2: site.tag[0]),
                      lambda site, site2: (abs(site.tag[0] / 100),
                                           abs(site.tag[1] / 100), 0)]

        syst2d.leads = []
        plot(syst2d, file=out_filename, show=False)
        del syst2d[list(syst2d.hoppings())]
        plot(syst2d, file=out_filename, show=False)

        plot(syst3d, file=out_filename, show=False)
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            plot(syst2d.finalized(), file=out_filename, show=False)

        # test 2D projections of 3D systems
        plot(syst3d, file=out_filename, pos_transform=lambda pos: pos[:2], show=False)


@pytest.mark.parametrize("engine", _plotter.engines)
def test_plot_more_site_families_than_colors(engine):
    # test against regression reported in
    # https://gitlab.kwant-project.org/kwant/kwant/issues/257
    ncolors = len(pyplot.rcParams['axes.prop_cycle'])
    syst = kwant.Builder()
    lattices = [kwant.lattice.square(name=i, norbs=1)
                for i in range(ncolors + 1)]
    for i, lat in enumerate(lattices):
        syst[lat(i, 0)] = None

    plotter.set_engine(engine)
    with tempfile.NamedTemporaryFile('w+b', suffix=plotter_file_suffix(engine)) as out:
        out_filename = out.name
        plotter.plot(syst, file=out_filename, show=False)


@pytest.mark.parametrize("engine", _plotter.engines)
def test_plot_raises_on_bad_site_spec(engine):
    syst = kwant.Builder()
    lat = kwant.lattice.square(norbs=1)
    syst[(lat(i, j) for i in range(5) for j in range(5))] = None

    # Cannot provide site_size as an array when syst is a Builder
    plotter.set_engine(engine)
    with pytest.raises(TypeError):
        plotter.plot(syst, site_size=[1] * 25)

    # Cannot provide site_size as an array when syst is a Builder
    with pytest.raises(TypeError):
        plotter.plot(syst, site_symbol=['o'] * 25)


def good_transform(pos):
    x, y = pos
    return y, x

def bad_transform(pos):
    x, y = pos
    return x, y, 0

@pytest.mark.parametrize("engine", _plotter.engines)
def test_map(engine):
    plotter.set_engine(engine)
    syst = syst_2d()

    with tempfile.NamedTemporaryFile('w+b', suffix=plotter_file_suffix(engine)) as out:
        out_filename = out.name
        plotter.map(syst, lambda site: site.tag[0], pos_transform=good_transform,
                    file=out_filename, method='linear', a=4, oversampling=4, cmap='flag', show=False)
        pytest.raises(ValueError, plotter.map, syst,
                      lambda site: site.tag[0],
                      pos_transform=bad_transform, file=out_filename)
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            plotter.map(syst.finalized(), range(len(syst.sites())),
                        file=out_filename, show=False)
        pytest.raises(ValueError, plotter.map, syst,
                      range(len(syst.sites())), file=out_filename)


def test_mask_interpolate():
    # A coordinate array with coordinates of two points almost coinciding.
    coords = np.array([[0, 0], [1e-7, 1e-7], [1, 1], [1, 0]])

    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")
        plotter.mask_interpolate(coords, np.ones(len(coords)), a=1)
        assert len(w) == 1
        assert issubclass(w[-1].category, RuntimeWarning)
        assert "coinciding" in str(w[-1].message)

    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        pytest.raises(ValueError, plotter.mask_interpolate,
                      coords, np.ones(len(coords)))
        pytest.raises(ValueError, plotter.mask_interpolate,
                      coords, np.ones(2 * len(coords)))


@pytest.mark.parametrize("engine", _plotter.engines)
def test_bands(engine):

    plotter.set_engine(engine)

    syst = syst_2d().finalized().leads[0]

    with tempfile.NamedTemporaryFile('w+b', suffix=plotter_file_suffix(engine)) as out:
        out_filename = out.name
        plotter.bands(syst, show=False, file=out_filename)
        plotter.bands(syst, show=False, momenta=np.linspace(0, 2 * np.pi), file=out_filename)

        if engine == 'matplotlib':
            plotter.bands(syst, show=False, fig_size=(10, 10), file=out_filename)

            fig = pyplot.Figure()
            ax = fig.add_subplot(1, 1, 1)
            plotter.bands(syst, show=False, ax=ax, file=out_filename)



@pytest.mark.parametrize("engine", _plotter.engines)
def test_spectrum(engine):

    plotter.set_engine(engine)

    def ham_1d(a, b, c):
        return a**2 + b**2 + c**2

    def ham_2d(a, b, c):
        return np.eye(2) * (a**2 + b**2 + c**2)

    lat = kwant.lattice.chain(norbs=1)
    syst = kwant.Builder()
    syst[(lat(i) for i in range(3))] = lambda site, a, b: a + b
    syst[lat.neighbors()] = lambda site1, site2, c: c
    fsyst = syst.finalized()

    vals = np.linspace(0, 1, 3)

    with tempfile.NamedTemporaryFile('w+b', suffix=plotter_file_suffix(engine)) as out:
        out_filename = out.name

        for ham in (ham_1d, ham_2d, fsyst):
            plotter.spectrum(ham, ('a', vals), params=dict(b=1, c=1), file=out_filename, show=False)
            if engine == 'matplotlib':
                # test with explicit figsize
                plotter.spectrum(ham, ('a', vals), params=dict(b=1, c=1),
                                 fig_size=(10, 10), file=out_filename, show=False)

        for ham in (ham_1d, ham_2d, fsyst):
            plotter.spectrum(ham, ('a', vals), ('b', 2 * vals),
                             params=dict(c=1), file=out_filename, show=False)
            if engine == 'matplotlib':
                # test with explicit figsize
                plotter.spectrum(ham, ('a', vals), ('b', 2 * vals),
                                 params=dict(c=1), fig_size=(10, 10), file=out_filename, show=False)

        if engine == 'matplotlib':
            # test 2D plot and explicitly passing axis
            fig = pyplot.figure()
            ax = fig.add_subplot(1, 1, 1, projection='3d')
            plotter.spectrum(ham_1d, ('a', vals), ('b', 2 * vals),
                             params=dict(c=1), ax=ax, file=out_filename, show=False)
            # explicitly pass axis without 3D support
            ax = fig.add_subplot(1, 1, 1)
            with pytest.raises(TypeError):
                plotter.spectrum(ham_1d, ('a', vals), ('b', 2 * vals),
                                 params=dict(c=1), ax=ax, file=out_filename, show=False)

    def mask(a, b):
        return a > 0.5

    with tempfile.NamedTemporaryFile('w+b', suffix=plotter_file_suffix(engine)) as out:
        out_filename = out.name
        plotter.spectrum(ham, ('a', vals), ('b', 2 * vals), params=dict(c=1),
                         mask=mask, file=out_filename, show=False)


def syst_rect(lat, salt, W=3, L=50):
    syst = kwant.Builder()

    ll = L//2
    ww = W//2

    def onsite(site):
        return 4 + 0.1 * kwant.digest.gauss(repr(site.tag), salt=salt)

    syst[(lat(i, j) for i in range(-ll, ll+1)
         for j in range(-ww, ww+1))] = onsite
    syst[lat.neighbors()] = -1

    sym = kwant.TranslationalSymmetry(lat.vec((-1, 0)))
    lead = kwant.Builder(sym)
    lead[(lat(0, j) for j in range(-ww, ww + 1))] = 4
    lead[lat.neighbors()] = -1

    syst.attach_lead(lead)
    syst.attach_lead(lead.reversed())

    return syst


def div(F, h):
    """Calculate the divergence of a vector field F over a grid of spacing h."""
    assert len(F.shape[:-1]) == F.shape[-1]
    assert len(h) == F.shape[-1]
    return sum(np.gradient(F[..., i], h[i])[i] for i in range(F.shape[-1]))


def rotational_currents(g):
    """Return a basis of divergence-free currents for a closed graph.

    Given the graph 'g' of a Kwant system, returns a sequence of arrays
    which are linearly independent, divergence-free currents on the graph.
    """
    #'A' represents the set of expressions that give the net current flow
    # into the system sites. 'perm' is a map from the edges of a graph
    # with only 1 edge per hopping to the proper Kwant graph (2 edges
    # per hopping).
    A = np.zeros((g.num_nodes, g.num_edges // 2))
    hoppings = dict()
    perm_data = np.zeros(g.num_edges, dtype=int)
    perm_ij = np.zeros((2, g.num_edges), dtype=int)
    i = 0
    for k, (a, b) in enumerate(g):
        hop = frozenset((a, b))
        if hop not in hoppings:
            A[a, i] = 1
            A[b, i] = -1
            hoppings[hop] = i
            perm_data[k] = 1
            perm_ij[:, k] = (k, i)
            i += 1
        else:
            perm_data[k] = -1
            perm_ij[:, k] = (k, hoppings[hop])

    perm = scipy.sparse.coo_matrix((perm_data, perm_ij))

    # Get the row vectors of V with singular value 0. These form
    # a basis for the right null space of 'A'.
    U, S, V = np.linalg.svd(A)
    tol = S.max() * max(A.shape) * np.finfo(S.dtype).eps
    rank = sum(S > tol)
    # Transform null space basis into vectors defined over the full
    # hopping space (both hopping directions).
    null_space_basis = V[-(len(V) - rank):].transpose()
    null_space_basis = perm.dot(null_space_basis).transpose()
    return null_space_basis


def _border_is_0(field):
    borders = [(0, slice(None)), (-1, slice(None)),
               (slice(None), 0), (slice(None), -1)]
    return all(np.allclose(field[a, b], 0) for a, b in borders)


def _test_border_0(interpolator):
    ## Test that current is always identically zero at box boundaries
    syst = kwant.Builder()
    lat = kwant.lattice.square(norbs=1)
    syst[[lat(0, 0), lat(1, 0)]] = None
    syst[(lat(0, 0), lat(1, 0))] = None
    syst = syst.finalized()
    values = [1, -1]

    ns = [3, 4, 5, 10, 100]
    abswidths = [0.01, 0.1, 1, 10, 100]
    relwidths = [0.01, 0.1, 1, 10, 100]
    for n, abswidth in itertools.product(ns, abswidths):
        field, _ = interpolator(syst, values, abswidth=abswidth, n=n)
        assert _border_is_0(field)
    for n, relwidth in itertools.product(ns, relwidths):
        field, _ = interpolator(syst, values, relwidth=relwidth, n=n)
        assert _border_is_0(field)


def test_density_interpolation():
    ## Passing a Builder will raise an error
    pytest.raises(TypeError, plotter.interpolate_density, syst_2d(), None)

    # Test that the density is always identically zero at the box boundaries
    # as the bump function has finite support and we add a padding
    _test_border_0(kwant.plotter.interpolate_density)

    def R(theta):
        return ta.array([[cos(theta), -sin(theta)], [sin(theta), cos(theta)]])

    # Make lattice with lattice vectors perturbed from x and y directions
    def make_lattice(a, salt='0'):
        theta_x = kwant.digest.uniform('x', salt=salt) * np.pi / 6
        theta_y = kwant.digest.uniform('y', salt=salt) * np.pi / 6
        x = ta.dot(R(theta_x), (a, 0))
        y = ta.dot(R(theta_y), (0, a))
        return kwant.lattice.general([x, y], norbs=1)

    # Check that integrating the interpolated density gives the same result
    # as summing the densities on all sites. We check this for several lattice
    # widths, lattice orientations and bump widths.
    for a, width in itertools.product((1, 2), (1, 0.5)):
        lat = make_lattice(a)
        syst = syst_rect(lat, salt='0').finalized()

        psi = kwant.wave_function(syst, energy=3)(0)[0]
        density = kwant.operator.Density(syst)(psi)
        exact_charge = sum(density)
        # We verify that the result is good by interpolating for
        # various numbers of points-per-bump and verifying that
        # the error falls of as 1/n.
        data = []
        for n in [4, 6, 8, 11, 16]:
            rho, box = plotter.interpolate_density(syst, density,
                                                   n=n, abswidth=width)
            (xmin, xmax), (ymin, ymax) = box
            area =  xmax - xmin * (ymax - ymin)
            N = rho.shape[0] * rho.shape[1]
            charge = np.sum(rho) * area / N
            data.append((n, abs(charge - exact_charge)))
        _, _, rvalue, *_ = scipy.stats.linregress(np.log(data))
        # Gradient of -1 on log-log plot means error falls off as 1/n
        # TODO: review this value once #280 has been dealt with.
        assert rvalue < -0.7

    # Test that the interpolation is linear in the input.
    rng = ensure_rng(1)
    lat = make_lattice(1, '1')
    syst = syst_rect(lat, salt='1').finalized()
    rho_0 = rng.rand(len(syst.sites))
    rho_1 = rng.rand(len(syst.sites))

    irho_0, _ = plotter.interpolate_density(syst, rho_0)
    irho_1, _ = plotter.interpolate_density(syst, rho_1)

    rho_tot, _ = plotter.interpolate_density(syst, rho_0 + 2 * rho_1)
    assert np.allclose(rho_tot, irho_0 + 2 * irho_1)


def test_current_interpolation():

    ## Passing a Builder will raise an error
    pytest.raises(TypeError, plotter.interpolate_current, syst_2d(), None)

    def R(theta):
        return ta.array([[cos(theta), -sin(theta)], [sin(theta), cos(theta)]])

    def make_lattice(a, theta):
        x = ta.dot(R(theta), (a, 0))
        y = ta.dot(R(theta), (0, a))
        return kwant.lattice.general([x, y], norbs=1)

    _test_border_0(plotter.interpolate_current)

    ## Check current through cross section is same for different lattice
    ## parameters and orientations of the system wrt. the discretization grid
    for a, theta, width in [(1, 0, 1),
                            (1, 0, 0.5),
                            (2, 0, 1),
                            (1, 0.2, 1),
                            (2, 0.4, 1)]:
        lat = make_lattice(a, theta)
        syst = syst_rect(lat, salt='0').finalized()
        psi = kwant.wave_function(syst, energy=3)(0)

        def cut(a, b):
            return b.tag[0] < 0 and a.tag[0] >= 0

        J = kwant.operator.Current(syst).bind()
        J_cut = kwant.operator.Current(syst, where=cut, sum=True).bind()
        J_exact = J_cut(psi[0])

        data = []
        for n in [4, 6, 8, 11, 16]:
            j0, box = plotter.interpolate_current(syst, J(psi[0]),
                                                  n=n, abswidth=width)
            x, y = (np.linspace(mn, mx, shape)
                    for (mn, mx), shape in zip(box, j0.shape))
            # slice field perpendicular to a cut along the y axis
            y_axis = (np.argmin(np.abs(x)), slice(None), 0)
            J_interp = simpson(j0[y_axis], x=y)
            data.append((n, abs(J_interp - J_exact)))
        # 3rd value returned from 'linregress' is 'rvalue'
        # TODO: review this value once #280 has been dealt with.
        assert scipy.stats.linregress(np.log(data))[2] < -0.7


    ### Tests on a divergence-free current (closed system)

    lat = kwant.lattice.general([(1, 0), (0.5, np.sqrt(3) / 2)], norbs=1)
    syst = kwant.Builder()
    sites = [lat(0, 0), lat(1, 0), lat(0, 1), lat(2, 2)]
    syst[sites] = None
    syst[((s, t) for s, t in itertools.product(sites, sites) if s != t)] = None
    del syst[lat(0, 0), lat(2, 2)]
    syst = syst.finalized()

    # generate random divergence-free currents
    Js = rotational_currents(syst.graph)
    rng = ensure_rng(3)
    J0 = sum(rng.rand(len(Js))[:, None] * Js)
    J1 = sum(rng.rand(len(Js))[:, None] * Js)

    # Sanity check that diverence on the graph is 0
    divergence = np.zeros(len(syst.sites))
    for (a, _), current in zip(syst.graph, J0):
        divergence[a] += current
    assert np.allclose(divergence, 0)

    j0, _ = plotter.interpolate_current(syst, J0)
    j1, _ = plotter.interpolate_current(syst, J1)

    ## Test linearity of interpolation.
    j_tot, _ = plotter.interpolate_current(syst, J0 + 2 * J1)
    assert np.allclose(j_tot, j0 + 2 * j1)

    ## Test that divergence of interpolated current approaches zero as we make
    ## the interpolation finer.
    data = []
    for n in [4, 6, 8, 11, 16]:
        j, box = plotter.interpolate_current(syst, J0, n=n)
        dx = [(mx - mn) / (shape - 1) for (mn, mx), shape in zip(box, j.shape)]
        div_j = np.max(np.abs(div(j, dx)))
        data.append((n, div_j))

    # 3rd value returned from 'linregress' is 'rvalue'
    # TODO: review this value once #280 has been dealt with.
    assert scipy.stats.linregress(np.log(data))[2] < -0.7


@pytest.mark.skipif(not _plotter.mpl_available, reason="Matplotlib unavailable.")
def test_current():
    plotter.set_engine('matplotlib')
    syst = syst_2d().finalized()
    J = kwant.operator.Current(syst)
    current = J(kwant.wave_function(syst, energy=1)(1)[0])

    # Test good codepath
    with tempfile.NamedTemporaryFile('w+b') as out:
        plotter.current(syst, current, file=out)

        fig = pyplot.Figure()
        ax = fig.add_subplot(1, 1, 1)
        plotter.current(syst, current, ax=ax, file=out)


@pytest.mark.parametrize("engine", _plotter.engines)
def test_density_works(engine):
    # This mainly checks that density() works, which was only caught by docs
    # before.
    plotter.set_engine(engine)
    syst = syst_2d().finalized()
    psi = kwant.wave_function(syst, energy=1)(1)[0]
    density = kwant.operator.Density(syst)(psi)

    suffix = plotter_file_suffix(engine)
    with tempfile.NamedTemporaryFile('w+b', suffix=suffix) as out:
        out_filename = out.name
        plotter.density(syst, density, file=out_filename, show=False)
