import numpy as np
import matplotlib.pyplot as plt

from .style import apply_dark_style, apply_dark_style_colorbar


def snapshot2d(
    snapshot,
    ax=None,
    vmin=None,
    vmax=None,
    figsize=(5, 5),
    cmap="RdBu_r",
    colorbar=True,
    dark=False,
):
    if ax:
        fig = ax.get_figure()
    else:
        # Create figure
        fig, ax = plt.subplots(figsize=figsize)

    # Initialize snapshot plot
    extent = [0, snapshot.length[0], 0, snapshot.length[1]]
    im = ax.imshow(
        np.asarray(snapshot.field).T,
        origin="lower",
        cmap=cmap,
        extent=extent,
        vmin=vmin,
        vmax=vmax,
        animated=True,
    )
    ax.set_xlabel(r"$x$", fontsize=16)
    ax.set_ylabel(r"$y$", fontsize=16)
    if colorbar:
        cbar = plt.colorbar(im, ax=ax, fraction=0.047, pad=0.01)
        cbar.ax.set_ylabel(r"$\varphi(x,y)$", fontsize=16)
        cbar.ax.tick_params(labelsize=12)
        if dark:
            apply_dark_style_colorbar(cbar)

    if dark:
        apply_dark_style(ax)

    fig.tight_layout()
    return im
