from collections.abc import Callable, Sequence, Mapping

import matplotlib.pyplot as plt
import numpy as np

from bayesflow.utils.dict_utils import make_variable_array, dicts_to_arrays, filter_kwargs, compute_test_quantities
from bayesflow.utils.plot_utils import (
    add_titles_and_labels,
    make_figure,
    set_layout,
    prettify_subplots,
)
from bayesflow.utils.validators import check_estimates_prior_shapes


def plot_quantity(
    values: Mapping[str, np.ndarray] | np.ndarray | Callable,
    targets: Mapping[str, np.ndarray] | np.ndarray,
    *,
    variable_keys: Sequence[str] = None,
    variable_names: Sequence[str] = None,
    estimates: Mapping[str, np.ndarray] | np.ndarray | None = None,
    test_quantities: dict[str, Callable] = None,
    figsize: Sequence[int] = None,
    label_fontsize: int = 16,
    title_fontsize: int = 18,
    tick_fontsize: int = 12,
    color: str = "#132a70",
    markersize: float = 25.0,
    marker: str = "o",
    alpha: float = 0.5,
    xlabel: str = "Ground truth",
    ylabel: str = "",
    num_col: int = None,
    num_row: int = None,
    default_name: str = "v",
) -> plt.Figure:
    """
    Plot a quantity as a function of a variable for each variable key.

    The function supports the following different combinations to pass
    or compute the values:

    1. pass `values` as an array of shape (num_datasets,) or (num_datasets, num_variables)
    2. pass `values` as a dictionary with the keys 'values', 'metric_name' and 'variable_names'
       as provided by the metrics functions. Note that the functions have to be called
       without aggregation to obtain value per dataset.
    3. pass a function to `values`, as well as `estimates`. The function should have the
       signature fn(estimates, targets, [aggregation]) and return an object like the
       `values` described in the previous options.

    Parameters
    ----------
    values      : dict[str, np.ndarray] | np.ndarray | Callable,
        The value of the quantity to plot. One of the following:

        1. an array of shape (num_datasets,) or (num_datasets, num_variables)
        2. a dictionary with the keys 'values', 'metric_name' and 'variable_names'
           as provided by the metrics functions. Note that the functions have to be called
           without aggregation to obtain value per dataset.
        3. a callable, requires passing `estimates` as well. The function should have the
           signature fn(estimates, targets, [aggregation]) and return an object like the
           ones described in the previous options.
    targets     : dict[str, np.ndarray] | np.ndarray,
        The parameter values plotted on the axis.
    variable_keys       : list or None, optional, default: None
       Select keys from the dictionary provided in samples.
       By default, select all keys.
    variable_names    : list or None, optional, default: None
        The parameter names for nice plot titles. Inferred if None
    estimates      : np.ndarray of shape (n_data_sets, n_post_draws, n_params), optional, default:  None
        The posterior draws obtained from n_data_sets. Can only be supplied if
        `values` is of type Callable.
    test_quantities   : dict or None, optional, default: None
        A dict that maps plot titles to functions that compute
        test quantities based on estimate/target draws.
        Can only be supplied if `values` is a function.

        The dict keys are automatically added to ``variable_keys``
        and ``variable_names``.
        Test quantity functions are expected to accept a dict of draws with
        shape ``(batch_size, ...)`` as the first (typically only)
        positional argument and return an NumPy array of shape
        ``(batch_size,)``.
        The functions do not have to deal with an additional
        sample dimension, as appropriate reshaping is done internally.
    figsize           : tuple or None, optional, default : None
        The figure size passed to the matplotlib constructor. Inferred if None.
    label_fontsize    : int, optional, default: 16
        The font size of the y-label text
    title_fontsize    : int, optional, default: 18
        The font size of the title text
    tick_fontsize     : int, optional, default: 12
        The font size of the axis ticklabels
    color             : str, optional, default: '#8f2727'
        The color for the true vs. estimated scatter points and error bars
    markersize        : float, optional, default: 25.0
        The marker size in points**2 for the scatter plot.
    marker            : str, optional, default: 'o'
        The marker for the scatter plot.
    alpha             : float, default: 0.5
        The opacity for the scatter plot
    num_row           : int, optional, default: None
        The number of rows for the subplots. Dynamically determined if None.
    num_col           : int, optional, default: None
        The number of columns for the subplots. Dynamically determined if None.
    default_name      : str, optional (default = "v")
        The default name to use for estimates if None provided

    Returns
    -------
    f : plt.Figure - the figure instance for optional saving

    Raises
    ------
    ShapeError
        If there is a deviation from the expected shapes of ``estimates`` and ``targets``.
    """

    if isinstance(values, Callable) and estimates is None:
        raise ValueError("Supplied a callable as `values`, but no `estimates`.")
    if not isinstance(values, Callable) and test_quantities is not None:
        raise ValueError(
            "Supplied `test_quantities`, but `values` is not a function. "
            "As the values have to be calculated for the test quantities, "
            "passing a function is required."
        )

    d = _prepare_values(
        values=values,
        targets=targets,
        estimates=estimates,
        variable_keys=variable_keys,
        variable_names=variable_names,
        test_quantities=test_quantities,
        label=None,
        default_name=default_name,
    )
    (values, targets, variable_keys, variable_names, test_quantities, _) = (
        d["values"],
        d["targets"],
        d["variable_keys"],
        d["variable_names"],
        d["test_quantities"],
        d["label"],
    )

    # store variable information at the top level for easy access
    num_variables = len(variable_names)

    # Configure layout
    num_row, num_col = set_layout(num_variables, num_row, num_col)

    # Initialize figure
    fig, axes = make_figure(num_row, num_col, figsize=figsize)

    # Loop and plot
    for i, ax in enumerate(axes.flat):
        if i >= num_variables:
            break

        ax.scatter(targets[:, i], values[:, i], color=color, alpha=alpha, s=markersize, marker=marker)

    prettify_subplots(axes, num_subplots=num_variables, tick_fontsize=tick_fontsize)

    # Add labels, titles, and set font sizes
    add_titles_and_labels(
        axes=axes,
        num_row=num_row,
        num_col=num_col,
        title=variable_names,
        xlabel=xlabel,
        ylabel=ylabel,
        title_fontsize=title_fontsize,
        label_fontsize=label_fontsize,
    )

    fig.tight_layout()
    return fig


def _prepare_values(
    *,
    values: Mapping[str, np.ndarray] | np.ndarray | Callable,
    targets: Mapping[str, np.ndarray] | np.ndarray,
    estimates: Mapping[str, np.ndarray] | np.ndarray | None,
    variable_keys: Sequence[str],
    variable_names: Sequence[str],
    test_quantities: dict[str, Callable],
    label: str | None,
    default_name: str,
):
    """
    Private helper function to compute/extract the values required for plotting
    a quantity.

    Refer to pairs_quantity and plot_quantity for details.
    """
    is_values_callable = isinstance(values, Callable)
    # Optionally, compute and prepend test quantities from draws
    if test_quantities is not None:
        updated_data = compute_test_quantities(
            targets=targets,
            estimates=estimates,
            variable_keys=variable_keys,
            variable_names=variable_names,
            test_quantities=test_quantities,
        )
        variable_names = updated_data["variable_names"]
        variable_keys = updated_data["variable_keys"]
        estimates = updated_data["estimates"]
        targets = updated_data["targets"]

    if estimates is not None:
        if is_values_callable:
            values = values(
                estimates=estimates,
                targets=targets,
                variable_keys=variable_keys,
                **filter_kwargs({"aggregation": None}, values),
            )

        data = dicts_to_arrays(
            estimates=estimates,
            targets=targets,
            variable_keys=variable_keys,
            variable_names=variable_names,
            default_name=default_name,
        )
        check_estimates_prior_shapes(data["estimates"], data["targets"])
        estimates = data["estimates"]
        targets = data["targets"]

        variable_keys = variable_keys or estimates.variable_keys
        if test_quantities is None:
            variable_names = variable_names or estimates.variable_names

    if all([key in values for key in ["values", "metric_name", "variable_names"]]):
        # output of a metric function
        label = values["metric_name"] if label is None else label
        variable_names = variable_names or values["variable_names"]
        values = values["values"]

    if hasattr(values, "variable_keys"):
        variable_keys = variable_keys or values.variable_keys
    if hasattr(values, "variable_names") and test_quantities is None:
        variable_names = variable_names or values.variable_names

    try:
        targets = make_variable_array(
            targets,
            variable_keys=variable_keys,
            variable_names=variable_names,
            default_name=default_name,
        )
    except ValueError:
        raise ValueError(
            "Length of 'variable_names' and number of variables do not match. "
            "Did you forget to specify `variable_keys`?"
        )
    variable_names = targets.variable_names
    variable_keys = targets.variable_keys

    if values.ndim == 1:
        values = values[:, None].repeat(len(variable_names), axis=-1)

    try:
        values = make_variable_array(
            values,
            variable_keys=variable_keys,
            variable_names=variable_names,
            default_name=default_name,
        )
    except ValueError:
        raise ValueError(
            "Length of 'variable_names' and number of variables do not match. "
            "Did you forget to specify `variable_keys`?"
        )

    return {
        "values": values,
        "targets": targets,
        "variable_keys": variable_keys,
        "variable_names": variable_names,
        "test_quantities": test_quantities,
        "label": label,
    }
