"""
Includes some basic plotting templates
"""
import matplotlib.pyplot as plt
from matplotlib.ticker import ScalarFormatter
from matplotlib import cm
from numpy import array, max, min, arange

try:
    from seaborn import set_style

    set_style("whitegrid")
except ModuleNotFoundError as e:
    print("Seaborn will not be used for plotting: ", e)
    pass


def show():
    """Show created plots by using this wrapper around matplotlib's show."""
    plt.show()


def general_plot(
    x,
    matrix,
    range_x=None,
    range_y=None,
    log: bool = True,
    labels=None,
    xlabel: str = "Frequency / Hz",
    ylabel: str | None = None,
    info_box: str | None = None,
    tight_layout: bool = True,
    returns: bool = False,
):
    """Generic plot template.

    Parameters
    ----------
    x : array-like
        Vector for x axis. Pass `None` to ignore.
    matrix : `np.ndarray`
        Matrix with data to plot.
    range_x : array-like, optional
        Range to show for x axis. Default: None.
    range_y : array-like, optional
        Range to show for y axis. Default: None.
    log : bool, optional
        Show x axis as logarithmic. Default: `True`.
    labels : list or str, optional
        Labels for the drawn lines as list of strings. Default: `None`.
    xlabel : str, optional
        Label for x axis. Default: None.
    ylabel : str, optional
        Label for y axis. Default: None.
    info_box : str, optional
        String containing extra information to be shown in a info box on the
        plot. Default: None.
    tight_layout: bool, optional
        When `True`, tight layout is activated. Default: `True`.
    returns : bool, optional
        When `True`, the figure and axis are returned. Default: `False`.

    Returns
    -------
    fig, ax
        Returned only when `returns=True`.

    """
    fig, ax = plt.subplots(1, 1, figsize=(8, 5))
    if matrix.ndim == 1:
        matrix = matrix[..., None]
    elif matrix.ndim > 2:
        raise ValueError("Only 1D and 2D-arrays are supported")
    if x is None:
        x = arange(matrix.shape[0])
    if labels is not None:
        if type(labels) not in (list, tuple):
            assert type(labels) is str, "labels should be a list or a string"
            labels = [labels]
    for n in range(matrix.shape[1]):
        if labels is not None:
            ax.plot(x, matrix[:, n], label=labels[n])
        else:
            ax.plot(x, matrix[:, n])
    if log:
        ax.set_xscale("log")
        ticks = array([20, 50, 100, 200, 500, 1000, 2000, 5000, 10000, 20000])
        if range_x is not None:
            ticks = ticks[(ticks > range_x[0]) & (ticks < range_x[-1])]
        ax.set_xticks(ticks)
        ax.get_xaxis().set_major_formatter(ScalarFormatter())
    ax.xaxis.grid(True, which="minor")
    if range_x is not None:
        ax.set_xlim(range_x)
    if range_y is not None:
        ax.set_ylim(range_y)
    if labels is not None:
        ax.legend()
    if ylabel is not None:
        ax.set_ylabel(ylabel)
    if xlabel is not None:
        ax.set_xlabel(xlabel)
    if info_box is not None:
        ax.text(
            0.1,
            0.5,
            info_box,
            transform=ax.transAxes,
            verticalalignment="top",
            bbox=dict(boxstyle="round", facecolor="grey", alpha=0.75),
        )
    if tight_layout:
        fig.tight_layout()
    if returns:
        return fig, ax


def general_subplots_line(
    x,
    matrix,
    column: bool = True,
    sharex: bool = True,
    sharey: bool = False,
    log: bool = False,
    xlabels=None,
    ylabels=None,
    range_x=None,
    range_y=None,
    returns: bool = False,
):
    """Generic plot template with subplots in one column or row.

    Parameters
    ----------
    x : array-like
        Vector for x axis.
    matrix : `np.ndarray`
        Matrix with data to plot.
    column : bool, optional
        When `True`, the subplots are organized in one column. Default: `True`.
    sharex : bool, optional
        When `True`, all subplots share the same values for the x axis.
        Default: `True`.
    sharey : bool, optional
        When `True`, all subplots share the same values for the y axis.
        Default: `False`.
    log : bool, optional
        Show x axis as logarithmic. Default: `False`.
    xlabels : array_like, optional
        Labels for x axis. Default: None.
    ylabels : array_like, optional
        Labels for y axis. Default: None.
    range_x : array-like, optional
        Range to show for x axis. Default: None.
    range_y : array-like, optional
        Range to show for y axis. Default: None.
    returns : bool, optional
        When `True`, the figure and axis are returned. Default: `False`.

    Returns
    -------
    fig, ax
        Returned only when `returns=True`.

    """
    if matrix.ndim == 1:
        matrix = matrix[..., None]
    elif matrix.ndim > 2:
        raise ValueError("Unsupported dimension. Matrix must be a 2D-array")
    number_of_channels = matrix.shape[1]
    if column:
        fig, ax = plt.subplots(
            number_of_channels,
            1,
            sharex=sharex,
            figsize=(8, 2 * number_of_channels),
            sharey=sharey,
        )
    else:
        fig, ax = plt.subplots(
            1,
            number_of_channels,
            sharex=sharex,
            figsize=(2 * number_of_channels, 8),
            sharey=sharey,
        )
    if number_of_channels == 1:
        ax = [ax]
    for n in range(number_of_channels):
        ax[n].plot(x, matrix[:, n])
        if log:
            ax[n].set_xscale("log")
            ticks = array(
                [20, 50, 100, 200, 500, 1000, 2000, 5000, 10000, 20000]
            )
            if range_x is not None:
                ticks = ticks[(ticks > range_x[0]) & (ticks < range_x[-1])]
            ax[n].set_xticks(ticks)
            ax[n].get_xaxis().set_major_formatter(ScalarFormatter())
        if ylabels is not None:
            ax[n].set_ylabel(ylabels[n])
        if xlabels is not None:
            if not type(xlabels) is str and len(xlabels) > 1:
                ax[n].set_xlabel(xlabels[n])
        if range_x is not None:
            ax[n].set_xlim(range_x)
        if range_y is not None:
            ax[n].set_ylim(range_y)
    if type(xlabels) is str or len(xlabels) == 1:
        ax[-1].set_xlabel(xlabels)
    fig.tight_layout()
    if returns:
        return fig, ax


def general_matrix_plot(
    matrix,
    range_x=None,
    range_y=None,
    range_z: float | None = None,
    xlabel: str | None = None,
    ylabel: str | None = None,
    zlabel: str | None = None,
    xlog: bool = False,
    ylog: bool = False,
    colorbar: bool = True,
    cmap: str = "magma",
    lower_origin: bool = True,
    returns: bool = False,
):
    """Generic plot template for a matrix's heatmap.

    Parameters
    ----------
    matrix : `np.ndarray`
        Matrix with data to plot.
    range_x : array-like, optional
        Range to show for x axis. Default: `None`.
    range_y : array-like, optional
        Range to show for y axis. Default: `None`.
    range_z : float, optional
        Dynamic range to show. Default: `None`.
    xlabel : str, optional
        Label for x axis. Default: `None`.
    ylabel : str, optional
        Label for y axis. Default: `None`.
    zlabel : str, optional
        Label for z axis. Default: `None`.
    xlog : bool, optional
        Show x axis as logarithmic. Default: `False`.
    ylog : bool, optional
        Show y axis as logarithmic. Default: `False`.
    colorbar : bool, optional
        When `True`, a colorbar for zaxis is shown. Default: `True`.
    cmap : str, optional
        Type of colormap to use from matplotlib.
        See https://matplotlib.org/stable/tutorials/colors/colormaps.html.
        Default: `'magma'`.
    lower_origin : bool, optional
        When `True`, the origin of the vertical axis of the matrix is put
        below. Default: `True`.
    returns : bool, optional
        When `True`, the figure and axis are returned. Default: `False`.

    Returns
    -------
    fig, ax
        Returned only when `returns=True`.

    """
    assert matrix.ndim == 2, "Only 2D-arrays are supported for this plot type"
    extent = None
    if range_x is not None:
        assert range_y is not None, (
            "When x range is given, y range is " + "also necessary"
        )
        assert len(range_x) == 2 and len(range_y) == 2, (
            "xrange and or yrange are invalid. Please give a list "
            + "containing (min, max) values"
        )
        extent = (range_x[0], range_x[1], range_y[0], range_y[1])

    fig, ax = plt.subplots(1, 1, figsize=(7, 5))
    cmap = cm.get_cmap(cmap)
    # cmap._init()
    # cmap._lut[-3, -1]
    if range_z is not None:
        max_val = max(matrix)
        min_val = max_val - range_z
    else:
        max_val = max(matrix)
        min_val = min(matrix)

    if lower_origin:
        origin = "lower"
    else:
        origin = "upper"

    if extent is None:
        col = ax.imshow(
            matrix,
            alpha=0.95,
            cmap=cmap,
            vmin=min_val,
            vmax=max_val,
            origin=origin,
            aspect="auto",
        )
    else:
        col = ax.imshow(
            matrix,
            extent=extent,
            alpha=0.95,
            cmap=cmap,
            vmin=min_val,
            vmax=max_val,
            origin=origin,
            aspect="auto",
        )
    if colorbar:
        if zlabel is not None:
            fig.colorbar(col, ax=ax, label=zlabel)
        else:
            fig.colorbar(col, ax=ax)
    if xlabel is not None:
        ax.set_xlabel(xlabel)
    if ylabel is not None:
        ax.set_ylabel(ylabel)
    if xlog:
        ax.set_xscale("log")
    if ylog:
        ax.set_yscale("log")
        ticks = array([20, 50, 100, 200, 500, 1000, 2000, 5000, 10000, 20000])
        if range_y is not None:
            ticks = ticks[(ticks > range_y[0]) & (ticks < range_y[-1])]
        ax.set_yticks(ticks)
        ax.get_yaxis().set_major_formatter(ScalarFormatter())
    fig.tight_layout()
    if returns:
        return fig, ax
