from collections.abc import Mapping, Sequence

import numpy as np
from scipy.special import logsumexp
import keras

from bayesflow.adapters import Adapter
from bayesflow.networks import ScoringRuleNetwork, SummaryNetwork
from bayesflow.types import Tensor
from bayesflow.utils import (
    logging,
    split_arrays,
    squeeze_inner_estimates_dict,
)
from bayesflow.utils.serialization import serializable

from .continuous_approximator import ContinuousApproximator


@serializable("bayesflow.approximators")
class ScoringRuleApproximator(ContinuousApproximator):
    """
    A workflow for fast amortized Bayes risk minimization for arbitrary scoring rules.

    Inherits from :class:`ContinuousApproximator` and adapts the sample, log_prob, and estimate
    interfaces for the nested output structure of :class:`~bayesflow.networks.ScoringRuleNetwork`.

    Parameters
    ----------
    inference_network : InferenceNetwork
        The inference network used for point estimation.
    adapter : bayesflow.adapters.Adapter, optional
        Adapter for data processing. You can use :py:meth:`build_adapter`
        to create it. If ``None`` (default), an identity adapter is used
        that makes a shallow copy and passes data through unchanged.
    summary_network : SummaryNetwork, optional
        The summary network used for data summarization (default is None).
    standardize : str | Sequence[str] | None
        The variables to standardize before passing to the networks. Can be either
        "all" or any subset of ["inference_variables", "summary_variables", "inference_conditions"].
        (default is "inference_variables").
    **kwargs : dict, optional
        Additional arguments passed to the :py:class:`ContinuousApproximator` class.
    """

    def __init__(
        self,
        *,
        inference_network: ScoringRuleNetwork,
        adapter: Adapter = None,
        summary_network: SummaryNetwork | None = None,
        standardize: str | Sequence[str] | None = "inference_variables",
        **kwargs,
    ):
        super().__init__(
            adapter=adapter,
            inference_network=inference_network,
            summary_network=summary_network,
            standardize=standardize,
            **kwargs,
        )

        self.distribution_keys = []
        for score_key, score in self.inference_network.scoring_rules.items():
            has_sample = callable(getattr(score, "sample", None))
            has_log_prob = callable(getattr(score, "log_prob", None))

            if has_sample and has_log_prob:
                self.distribution_keys.append(score_key)

        self.has_distribution = len(self.distribution_keys) > 0

    def estimate(
        self,
        conditions: Mapping[str, np.ndarray] | None = None,
        split: bool = False,
        groupby: str = "variable",
        **kwargs,
    ) -> dict[str, dict[str, np.ndarray | dict[str, np.ndarray]]]:
        """
        Estimate point summaries and distributional parameters of inference variables induced by
        scoring rules.

        Parameters
        ----------
        conditions : Mapping[str, np.ndarray]
            A batch of conditions to estimate. None for unconditional distributions.
        split : bool, optional
            If True, split estimated arrays along the last axis, by default False.
        groupby : {"variable", "score"}, default "variable"
            Controls the top-level nesting of the returned dictionary.

            - "variable": return estimates as ``variable -> score -> head -> array``.
              If a score has just one head, which is called value, squeeze the unnecessary nesting.
            - "score": return estimates as ``score -> head -> variable -> array``.
        **kwargs
            Additional keyword arguments passed to underlying processing functions.

        Returns
        -------
        dict
            Estimates in the requested layout. Leafs have shape
            ``(num_conditions, point_estimate_size, variable_block_size)`` or
            ``(point_estimate_size, variable_block_size)`` if ``conditions=None``.
        """
        estimates = self._estimate_byscore(conditions=conditions, split=split, **kwargs)

        if groupby == "score":
            return estimates

        # Reorder the nested dictionary so that original variable names are at the top.
        estimates = self._reorder_estimates(estimates)
        estimates = self._squeeze_estimates(estimates)

        return estimates

    def _estimate_byscore(
        self,
        conditions: Mapping[str, np.ndarray],
        split: bool = False,
        **kwargs,
    ) -> dict[str, dict[str, dict[str, np.ndarray]]]:
        resolved_conditions, adapted, _ = self._prepare_conditions(conditions)

        inference_kwargs = kwargs | self._collect_mask_kwargs(self._INFERENCE_MASK_KEYS, adapted)

        estimates = self.inference_network(
            conditions=resolved_conditions,
            **inference_kwargs,
        )

        # Unstandardize the network outputs
        for score_key, score in self.inference_network.scoring_rules.items():
            for head_key in estimates[score_key].keys():
                transformation_type = score.TRANSFORMATION_TYPE.get(head_key, "location_scale")
                estimates[score_key][head_key] = self.standardizer.maybe_standardize(
                    estimates[score_key][head_key],
                    key="inference_variables",
                    stage="inference",
                    forward=False,
                    transformation_type=transformation_type,
                )

        estimates = self._apply_inverse_adapter_to_estimates(estimates, **kwargs)

        if split:
            estimates = split_arrays(estimates, axis=-1)

        return estimates

    def sample(
        self,
        *,
        num_samples: int,
        conditions: Mapping[str, np.ndarray] | None = None,
        split: bool = False,
        score_weights: Mapping[str, float] | None = None,
        merge_scores: bool = True,
        **kwargs,
    ) -> dict[str, np.ndarray] | dict[str, dict[str, np.ndarray]]:
        """
        Draw samples from the parametric distributions induced by the configured scoring rules.

        This method supports two modes:

        1) ``merge_scores=True`` (default):
        Samples are drawn from the *mixture* over scoring rules.
        The requested ``num_samples`` are allocated across scores according to ``score_weights``
        (uniform if None), drawn for each score, and merged into a single shuffled set.

        2) ``merge_scores=False``:
        Samples are drawn *separately* for each scoring rule, returning a nested dictionary
        keyed by score name. In this mode, ``num_samples`` samples are generated per score.

        Parameters
        ----------
        num_samples : int
            Number of samples to draw. If ``merge_scores=True``, this is the total number of
            mixture samples returned. If ``merge_scores=False``, this is the number of samples
            generated per scoring rule.
        conditions : Mapping[str, np.ndarray] or None,
            A batch of conditions to sample for. None for unconditional distributions.
        split : bool, optional
            Whether to split the output arrays along the last axis. Delegated to :meth:`sample_separate`.
        score_weights : Mapping[str, float] or None, default None
            Probability weights for each scoring rule. If ``None``, uniform weights are assumed.
            Must be positive, will be normalized to sum to 1. Only used when ``merge_scores=True``.
        merge_scores : bool, default True
            If True, return samples aggregated across scoring rules as a mixture.
            If False, return samples separately for each scoring rule.
        **kwargs
            Additional keyword arguments such as ``batch_size``.

        Returns
        -------
        samples : dict[str, np.ndarray] or dict[str, dict[str, np.ndarray]]
            If ``merge_scores=True``:
                A dictionary keyed by inference variable name. Entries have shape
                ``(num_datasets, num_samples, ...)``.
            If ``merge_scores=False``:
                A nested dictionary where the first-level key is the score name and the second-level
                key is the inference variable name. Each leaf array has shape
                ``(num_datasets, num_samples, ...)``.

        Raises
        ------
        NotImplementedError
            If ``split=True`` is not supported by the approximator implementation.
        """
        self._check_has_distribution()

        if not merge_scores:
            if score_weights is not None:
                logging.warning(
                    "`score_weights` is ignored when `merge_scores=False`. "
                    "Set `merge_scores=True` to sample from the weighted mixture."
                )
            return self._sample_separate(num_samples=num_samples, conditions=conditions, split=split, **kwargs)

        score_weights = self._resolve_score_weights(score_weights)

        # Single score: _sample_separate already squeezed to a plain result,
        # and mixing with uniform weight is an identity operation.
        if len(self.distribution_keys) == 1:
            return self._sample_separate(num_samples=num_samples, conditions=conditions, split=split, **kwargs)

        # Allocate samples per score and draw only as many as needed (max over scores).
        num_samples_per_score = np.random.multinomial(num_samples, list(score_weights.values()))
        max_k = int(np.max(num_samples_per_score))
        num_samples_per_score = {k: num_samples_per_score[i] for i, k in enumerate(score_weights.keys())}

        samples_by_score = self._sample_separate(num_samples=max_k, conditions=conditions, split=split, **kwargs)

        # Crop each score's samples down to its allocated k
        cropped_list = []
        for score_key, k in num_samples_per_score.items():
            if k == 0:
                continue
            cropped = keras.tree.map_structure(lambda arr: arr[:, :k], samples_by_score[score_key])
            cropped_list.append(cropped)

        # Concatenate across scores along the sample axis.
        concatenated = keras.tree.map_structure(
            lambda *arrays: np.concatenate(arrays, axis=1),
            *cropped_list,
        )

        # Shuffle along the sample axis (1) to form the mixture samples.
        shuffle_idx = np.random.permutation(num_samples)
        shuffled = keras.tree.map_structure(lambda arr: np.take(arr, shuffle_idx, axis=1), concatenated)
        return shuffled

    def _sample_separate(
        self,
        *,
        num_samples: int,
        conditions: Mapping[str, np.ndarray] | None = None,
        split: bool = False,
        batch_size: int | None = None,
        **kwargs,
    ) -> dict[str, np.ndarray | dict[str, np.ndarray]]:
        """
        Draws samples from a parametric distribution based on point estimates.

        Uses :meth:`ContinuousApproximator.sample` for condition resolution, sampling,
        unstandardization, and inverse adapter, then squeezes the nested score-major
        output structure.

        Parameters
        ----------
        num_samples : int
            The number of samples to generate.
        conditions : Mapping[str, np.ndarray]
            A dictionary mapping variable names to arrays representing the conditions.
        split : bool, optional
            If True, the sampled arrays are split along the last axis, by default False.
            Currently not supported for :py:class:`ScoringRuleApproximator`.
        batch_size : int or None, optional
            If provided, the conditions are split into batches of size `batch_size`,
            for which samples are generated sequentially.
        **kwargs
            Additional keyword arguments passed to underlying processing functions.

        Returns
        -------
        samples : dict[str, dict[str, np.ndarray]]
            Samples for all inference variables and all parametric scoring rules in a nested dictionary.
            Shape: (num_datasets, num_samples, variable_block_size).
        """

        if split:
            raise NotImplementedError("split=True is currently not supported for `ScoringRuleApproximator`.")

        self._check_has_distribution()

        # Delegate to parent for condition resolution, sampling, unstandardization, and inverse adapter
        samples = super().sample(
            num_samples=num_samples,
            conditions=conditions,
            split=False,
            batch_size=batch_size,
            **kwargs,
        )

        return self._squeeze_parametric_score_major_dict(samples)

    def log_prob(
        self,
        data: Mapping[str, np.ndarray],
        score_weights: Mapping[str, float] | None = None,
        merge_scores: bool = True,
        **kwargs,
    ) -> np.ndarray | dict[str, np.ndarray]:
        """
        Compute log-probabilities under the parametric distribution(s) induced by the scoring rules.

        This method supports two modes:

        1) ``merge_scores=True`` (default):
        Return the marginalized (mixture) log-probability across scoring rules using ``score_weights``:
        ``log p(x) = log sum_s w_s p_s(x)``.

        2) ``merge_scores=False``:
        Return log-probabilities separately for each scoring rule as a dictionary keyed by score name.
        In this mode, ``score_weights`` is ignored.

        Parameters
        ----------
        data : Mapping[str, np.ndarray]
            Dictionary containing inference variables and conditions.
        score_weights : Mapping[str, float] or None, default None
            Probability weights for each scoring rule. If ``None``, uniform weights are assumed.
            Must be positive, will be normalized to sum to 1. Only used when ``merge_scores=True``.
        merge_scores : bool, default True
            If True, return marginalized log-probabilities across scoring rules.
            If False, return per-score log-probabilities.
        **kwargs
            Additional keyword arguments.

        Returns
        -------
        np.ndarray or dict[str, np.ndarray]
            If ``merge_scores=True``: array of shape ``(num_datasets,)`` with marginalized log-probabilities.
            If ``merge_scores=False``: dictionary mapping score name to an array of shape ``(num_datasets,)``.

        """
        log_probs = self._log_prob_separate(data=data, **kwargs)

        if not merge_scores:
            if score_weights is not None:
                logging.warning(
                    "`score_weights` is ignored when `merge_scores=False`. "
                    "Set `merge_scores=True` to compute the weighted mixture log-probability."
                )
            return log_probs

        # Single score: _log_prob_separate already squeezed to a plain array,
        # and merging with uniform weight is an identity operation.
        if len(self.distribution_keys) == 1:
            return np.asarray(log_probs)

        score_weights = self._resolve_score_weights(score_weights)

        stacked = np.stack([np.asarray(log_probs[score_key]) for score_key in self.distribution_keys], axis=-1)
        log_weights = np.log(list(score_weights.values()))

        # stacked: (num_datasets, num_scores), log_weights: (num_scores,)
        z = stacked + log_weights  # broadcasted to (num_datasets, num_scores)

        # stable logsumexp over last axis
        return logsumexp(z, axis=-1)

    def _log_prob_separate(self, data: Mapping[str, np.ndarray], **kwargs) -> dict[str, np.ndarray]:
        """
        Computes the log-probability of given data under the parametric distribution(s).

        Parameters
        ----------
        data : dict[str, np.ndarray]
            A dictionary mapping variable names to arrays representing the data.
        **kwargs
            Additional keyword arguments passed to underlying processing functions.

        Returns
        -------
        log_prob : np.ndarray or dict[str, np.ndarray]
            Log-probabilities of the distribution for all parametric scoring rules.
            If only one parametric score is available, returns an array.
            Otherwise, returns a dictionary with score names as keys.
            Shape: (num_datasets,)
        """
        log_prob = super().log_prob(data, **kwargs)
        return self._squeeze_parametric_score_major_dict(log_prob)

    def _check_has_distribution(self):
        if not self.has_distribution:
            raise ValueError("No parametric distribution scores available for sample/log_prob.")

    def _resolve_score_weights(
        self,
        score_weights: Mapping[str, float] | None,
    ) -> np.ndarray:
        if score_weights is None:
            score_weights = {k: 1.0 for k in self.distribution_keys}

        for key, weight in score_weights.items():
            if key not in self.distribution_keys:
                raise ValueError(
                    "Score weights must be subset of self.distribution_keys. "
                    f"Unknown keys: {set(score_weights) - set(self.distribution_keys)}"
                )
            if weight < 0:
                raise ValueError(f"All score_weights must be positive. Received {key}: {weight}.")

        # Normalize weights to 1
        sum = np.sum(list(score_weights.values()))
        score_weights = {k: v / sum for k, v in score_weights.items()}
        return score_weights

    def _apply_inverse_adapter_to_estimates(
        self, estimates: Mapping[str, Mapping[str, Tensor]], **kwargs
    ) -> dict[str, dict[str, dict[str, np.ndarray]]]:
        """Applies the inverse adapter on each inner element of the _estimate output dictionary."""
        estimates = keras.tree.map_structure(keras.ops.convert_to_numpy, estimates)
        processed = {}
        for score_key, score_val in estimates.items():
            processed[score_key] = {}
            for head_key, estimate in score_val.items():
                if head_key in self.inference_network.scoring_rules[score_key].NOT_TRANSFORMING_LIKE_VECTOR_WARNING:
                    logging.warning(
                        f"Estimate '{score_key}.{head_key}' is marked to not transform like a vector. "
                        f"It was treated like a vector by the adapter. Handle '{head_key}' estimates with care."
                    )

                adapted = self.adapter(
                    {"inference_variables": estimate},
                    inverse=True,
                    strict=False,
                    **kwargs,
                )
                processed[score_key][head_key] = adapted
        return processed

    @staticmethod
    def _reorder_estimates(
        estimates: Mapping[str, Mapping[str, Mapping[str, np.ndarray]]],
    ) -> dict[str, dict[str, dict[str, np.ndarray]]]:
        """Reorders the nested dictionary so that the inference variable names become the top-level keys."""
        sample_inner = next(iter(next(iter(estimates.values())).values()))
        variable_names = sample_inner.keys()
        reordered = {}
        for variable in variable_names:
            reordered[variable] = {}
            for score_key, inner_dict in estimates.items():
                reordered[variable][score_key] = {inner_key: value[variable] for inner_key, value in inner_dict.items()}
        return reordered

    @staticmethod
    def _squeeze_estimates(
        estimates: Mapping[str, Mapping[str, Mapping[str, np.ndarray]]],
    ) -> dict[str, dict[str, np.ndarray]]:
        """Squeezes each inner estimate dictionary to remove unnecessary nesting."""
        squeezed = {}
        for variable, variable_estimates in estimates.items():
            squeezed[variable] = {
                score_key: squeeze_inner_estimates_dict(inner_estimate)
                for score_key, inner_estimate in variable_estimates.items()
            }
        return squeezed

    @staticmethod
    def _squeeze_parametric_score_major_dict(samples: Mapping[str, np.ndarray]) -> np.ndarray | dict[str, np.ndarray]:
        """Squeezes the dictionary to just the value if there is only one key-value pair."""
        if len(samples) == 1:
            return next(iter(samples.values()))
        return samples
