"""Utilities for plotting parameter data out of bulk fits."""
import matplotlib.pyplot as plt
from arpes.plotting.utils import latex_escape

import xarray as xr
from arpes.provenance import save_plot_provenance

__all__ = ("plot_parameter",)


@save_plot_provenance
def plot_parameter(
    fit_data: xr.DataArray,
    param_name: str,
    ax=None,
    fillstyle="none",
    shift=0,
    x_shift=0,
    markersize=8,
    title=None,
    out=None,
    two_sigma=False,
    **kwargs
):
    """Makes a simple scatter plot of a parameter from an `broadcast_fit` result."""
    fig = None

    if ax is None:
        fig, ax = plt.subplots(figsize=kwargs.get("figsize", (7, 5)))

    ds = fit_data.F.param_as_dataset(param_name)
    x_name = ds.value.dims[0]
    x = ds.coords[x_name].values

    color = kwargs.get("color")
    e_width = None
    l_width = None
    if two_sigma:
        line, markers, lines = ax.errorbar(
            x + x_shift,
            ds.value.values + shift,
            yerr=2 * ds.error.values,
            fmt="",
            elinewidth=1,
            linewidth=0,
            c=color,
        )
        color = lines[0].get_color()[0]
        e_width = 2
        l_width = 0

    ax.errorbar(
        x + x_shift,
        ds.value.values + shift,
        yerr=ds.error.values,
        fmt="s",
        color=color,
        elinewidth=e_width,
        linewidth=l_width,
        markeredgewidth=e_width or 2,
        fillstyle=fillstyle,
        markersize=markersize,
    )

    ax.set_xlabel(latex_escape(x_name))
    ax.set_ylabel(latex_escape(param_name))
    return ax
