from typing import Optional, Sequence, Literal
from copy import deepcopy

import numpy as np
import xarray as xr
import json_tricks

from questplus import psychometric_function


class QuestPlus:
    def __init__(
        self,
        *,
        stim_domain: dict,
        param_domain: dict,
        outcome_domain: dict,
        prior: Optional[dict] = None,
        func: Literal["weibull", "csf", "norm_cdf", "norm_cdf_2", "thurstone_scaling"],
        stim_scale: Optional[Literal["log10", "dB", "linear"]],
        stim_selection_method: str = "min_entropy",
        stim_selection_options: Optional[dict] = None,
        param_estimation_method: str = "mean",
    ):
        """
        A QUEST+ staircase procedure.

        Parameters
        ----------
        stim_domain
            Specification of the stimulus domain: dictionary keys correspond to
            the names of the stimulus dimensions, and  values describe the
            respective possible stimulus values (e.g., intensities, contrasts,
            or orientations).

        param_domain
            Specification of the parameter domain: dictionary keys correspond
            to the names of the parameter dimensions, and  values describe the
            respective possible parameter values (e.g., threshold, slope,
            lapse rate.

        outcome_domain
            Specification of the outcome domain: dictionary keys correspond
            to the names of the outcome dimensions, and  values describe the
            respective possible outcome values (e.g., "Yes", "No", "Correct",
            "Incorrect"). This argument typically describes the responses a
            participant can provide.

        prior
            A-priori probabilities of parameter values.

        func
            The psychometric function whose parameters to estimate.

        stim_scale
            The scale on which the stimuli are provided. Has no effect for the
            Thurstonian scaling function.

        stim_selection_method
            How to select the next stimulus. `min_entropy` picks the stimulus
            that will minimize the expected entropy. `min_n_entropy` randomly
            selects a stimulus from the set of stimuli that will yield the `n`
            smallest entropies. `n` has to be specified via the
            `stim_selection_options` keyword argument.

        stim_selection_options
            Use this argument to specify options for the stimulus selection
            method specified via `stim_selection_method`. Currently, this can
            be used to specify the number of `n` stimuli that will yield the
            `n` smallest entropies if `stim_selection_method=min_n_entropy`,
            and `max_consecutive_reps`, the number of times the same stimulus
            can be presented consecutively. A random number generator seed
            may be passed via `random_seed=12345`.

        param_estimation_method
            The method to use when deriving the final parameter estimate.
            Possible values are `mean` (mean of each parameter, weighted by the
            posterior probabilities) and `mode` (the parameters at the peak of
            the posterior distribution).

        """
        if func == "thurstone_scaling" and stim_scale is not None:
            raise ValueError(
                "The Thurstonian scaling function cannot be used with "
                "a stim_scale parameter."
            )

        self.func = func
        self.stim_scale = stim_scale
        self.stim_domain = self._ensure_ndarray(stim_domain)
        self.param_domain = self._ensure_ndarray(param_domain)
        self.outcome_domain = self._ensure_ndarray(outcome_domain)

        self.prior = self._gen_prior(prior=prior)
        self.posterior = deepcopy(self.prior)
        self.likelihoods = self._gen_likelihoods()

        self.stim_selection = stim_selection_method

        if self.stim_selection == "min_n_entropy":
            from ._constants import (
                DEFAULT_N,
                DEFAULT_RANDOM_SEED,
                DEFAULT_MAX_CONSECUTIVE_REPS,
            )

            if stim_selection_options is None:
                self.stim_selection_options = dict(
                    n=DEFAULT_N,
                    max_consecutive_reps=DEFAULT_MAX_CONSECUTIVE_REPS,
                    random_seed=DEFAULT_RANDOM_SEED,
                )
            else:
                self.stim_selection_options = stim_selection_options.copy()

                if "n" not in stim_selection_options:
                    self.stim_selection_options["n"] = DEFAULT_N
                if "max_consecutive_reps" not in stim_selection_options:
                    self.stim_selection_options[
                        "max_consecutive_reps"
                    ] = DEFAULT_MAX_CONSECUTIVE_REPS
                if "random_seed" not in stim_selection_options:
                    self.stim_selection_options["random_seed"] = DEFAULT_RANDOM_SEED

            del DEFAULT_N, DEFAULT_MAX_CONSECUTIVE_REPS, DEFAULT_RANDOM_SEED

            seed = self.stim_selection_options["random_seed"]
            self._rng = np.random.RandomState(seed=seed)
            del seed
        else:
            self.stim_selection_options = stim_selection_options
            self._rng = None

        self.param_estimation_method = param_estimation_method

        self.resp_history = list()
        self.stim_history = list()
        self.entropy = np.nan

    @staticmethod
    def _ensure_ndarray(x: dict) -> dict:
        x = deepcopy(x)
        for k, v in x.items():
            x[k] = np.atleast_1d(v)

        return x

    def _gen_prior(self, *, prior: dict) -> xr.DataArray:
        """
        Raises
        ------
        ValueError
            If the user specifies priors for parameters that do not appear in
            the parameter domain.

        """
        prior_orig = deepcopy(prior)

        if prior_orig is None:
            # Uninformative prior.
            prior = np.ones([len(x) for x in self.param_domain.values()])
        elif set(prior_orig.keys()) - set(self.param_domain.keys()):
            msg = (
                f"Mismatch between specified parameter domain and supplied "
                f"prior.\n"
                f"You specified priors for the following parameters that "
                f"do not appear in the parameter domain: "
                f"{set(prior_orig.keys()) - set(self.param_domain.keys())}"
            )
            raise ValueError(msg)
        elif set(self.param_domain.keys()) - set(prior_orig.keys()):
            # The user specified prior probabilities for only a subset
            # of the parameters. We use those, obviously; and fill the
            # remaining prior distributions with uninformative priors.
            grid_dims = []
            for param_name, param_vals in self.param_domain.items():
                if param_name in prior_orig.keys():
                    prior_vals = np.atleast_1d(prior_orig[param_name])
                else:
                    prior_vals = np.ones(len(param_vals))

                grid_dims.append(prior_vals)

            prior_grid = np.meshgrid(*grid_dims, sparse=True, indexing="ij")
            prior = np.prod(
                np.array(prior_grid, dtype="object")  # avoid warning re "ragged" array
            )
        else:
            # A "proper" prior was specified (i.e., prior probabilities for
            # all parameters.)
            prior_grid = np.meshgrid(
                *list(prior_orig.values()), sparse=True, indexing="ij"
            )
            prior = np.prod(
                np.array(prior_grid, dtype="object")  # avoid warning re "ragged" array
            )

        # Normalize.
        prior /= prior.sum()

        # Create the prior object we are actually going to use.
        dims = (*self.param_domain.keys(),)
        coords = dict(**self.param_domain)
        prior_ = xr.DataArray(data=prior, dims=dims, coords=coords)

        return prior_

    def _gen_likelihoods(self) -> xr.DataArray:
        outcome_dim_name = list(self.outcome_domain.keys())[0]
        outcome_values = list(self.outcome_domain.values())[0]

        if self.func not in [
            "weibull",
            "csf",
            "norm_cdf",
            "norm_cdf_2",
            "thurstone_scaling",
        ]:
            raise ValueError(
                f"Unknown psychometric function name specified: {self.func}"
            )

        if self.func == "weibull":
            f = psychometric_function.weibull
        elif self.func == "csf":
            f = psychometric_function.csf
        elif self.func == "norm_cdf":
            f = psychometric_function.norm_cdf
        elif self.func == "norm_cdf_2":
            f = psychometric_function.norm_cdf_2
        elif self.func == "thurstone_scaling":
            f = psychometric_function.thurstone_scaling_function

        if self.func == "thurstone_scaling":
            prop_correct = f(**self.stim_domain, **self.param_domain)
        else:
            prop_correct = f(
                **self.stim_domain, **self.param_domain, scale=self.stim_scale
            )
        prop_incorrect = 1 - prop_correct

        # Now this is a bit awkward. We concatenate the psychometric
        # functions for the different responses. To do that, we first have
        # to add an additional dimension.
        # TODO: There's got to be a neater way to do this?!
        corr_resp_dim = {outcome_dim_name: [outcome_values[0]]}
        inccorr_resp_dim = {outcome_dim_name: [outcome_values[1]]}

        prop_correct = prop_correct.expand_dims(corr_resp_dim)
        prop_incorrect = prop_incorrect.expand_dims(inccorr_resp_dim)

        pf_values = xr.concat(
            [prop_correct, prop_incorrect],
            dim=outcome_dim_name,
            coords=self.outcome_domain,
        )
        return pf_values

    def update(self, *, stim: dict, outcome: dict) -> None:
        """
        Inform QUEST+ about a newly gathered measurement outcome for a given
        stimulus parameter set, and update the posterior accordingly.

        Parameters
        ----------
        stim
            The stimulus that was used to generate the given outcome.

        outcome
            The observed outcome.

        """
        likelihood = self.likelihoods.sel(**stim, **outcome)

        self.posterior = self.posterior * likelihood
        self.posterior /= self.posterior.sum()

        # Log the results, too.
        self.stim_history.append(stim)
        self.resp_history.append(outcome)

    @property
    def next_stim(self) -> dict:
        """
        Retrieve the stimulus to present next.

        The stimulus will be selected based on the method in
        ``self.stim_selection``.

        """
        new_posterior = self.posterior * self.likelihoods

        # Probability.
        pk = new_posterior.sum(dim=self.param_domain.keys())
        new_posterior /= pk

        # Entropies.
        #
        # Note:
        #   - np.log(0) returns -inf (division by zero)
        #   - the multiplcation of new_posterior with -inf values generates
        #     NaN's
        #   - xr.DataArray.sum() has special handling for NaN's.
        #
        # NumPy also emits a warning, which we suppress here.
        with np.errstate(divide="ignore"):
            H = -(
                (new_posterior * np.log(new_posterior)).sum(
                    dim=self.param_domain.keys()
                )
            )

        # Expected entropies for all possible stimulus parameters.
        EH = (pk * H).sum(dim=list(self.outcome_domain.keys()))

        if self.stim_selection == "min_entropy":
            # Get the stimulus properties that minimize entropy.
            indices = EH.argmin(dim=...)
            stim = dict()
            for stim_property, index in indices.items():
                stim_val = EH[stim_property][index].item()
                stim[stim_property] = stim_val

            self.entropy = EH.min().item()
        elif self.stim_selection == "min_n_entropy":
            # Number of stimuli to include (the n stimuli that yield the lowest
            # entropies)
            n_stim = self.stim_selection_options["n"]

            indices = np.unravel_index(EH.argsort(), EH.shape)[0]
            indices = indices[:n_stim]

            while True:
                # Randomly pick one index and retrieve its coordinates
                # (stimulus parameters).
                candidate_index = self._rng.choice(indices)
                coords = EH[candidate_index].coords
                stim = {
                    stim_property: stim_val.item()
                    for stim_property, stim_val in coords.items()
                }

                max_reps = self.stim_selection_options["max_consecutive_reps"]

                if len(self.stim_history) < 2:
                    break
                elif all(
                    [stim == prev_stim for prev_stim in self.stim_history[-max_reps:]]
                ):
                    # Shuffle again.
                    continue
                else:
                    break
        else:
            raise ValueError("Unknown stim_selection supplied.")

        return stim

    @property
    def param_estimate(self) -> dict:
        """
        Retrieve the final parameter estimates after the QUEST+  run.

        The parameters will be calculated according to
        ``self.param_estimation_method``.

        This returns a dictionary of parameter estimates, where the dictionary
        keys correspond to the parameter names.

        """
        method = self.param_estimation_method
        param_estimates = dict()
        for param_name in self.param_domain.keys():
            params = list(self.param_domain.keys())
            params.remove(param_name)

            if method == "mean":
                param_estimates[param_name] = (
                    (self.posterior.sum(dim=params) * self.param_domain[param_name])
                    .sum()
                    .item()
                )
            elif method == "mode":
                indices = self.posterior.argmax(dim=...)
                coords = self.posterior[indices]
                param_estimates[param_name] = coords[param_name].item()
            else:
                raise ValueError("Unknown method parameter.")

        return param_estimates

    @property
    def marginal_posterior(self) -> dict:
        """
        Retrieve the a dictionary of marginal posterior probability
        density functions (PDFs).

        This returns a  dictionary of marginal PDFs, where the dictionary keys
        correspond to the parameter names.

        """
        marginal_posterior = dict()
        for param_name in self.param_domain.keys():
            marginalized_out_params = list(self.param_domain.keys())
            marginalized_out_params.remove(param_name)
            marginal_posterior[param_name] = self.posterior.sum(
                dim=marginalized_out_params
            ).values

        return marginal_posterior

    def to_json(self) -> str:
        """
        Dump this `QuestPlus` instance as a JSON string which can be loaded
        again later.

        Returns
        -------
        str
            A JSON dump of the current `QuestPlus` instance.

        See Also
        --------
        from_json

        """
        self_copy = deepcopy(self)
        self_copy.prior = self_copy.prior.to_dict()
        self_copy.posterior = self_copy.posterior.to_dict()
        self_copy.likelihoods = self_copy.likelihoods.to_dict()

        if self_copy._rng is not None:  # NumPy RandomState cannot be serialized.
            self_copy._rng = self_copy._rng.get_state()

        return json_tricks.dumps(self_copy, allow_nan=True)

    @staticmethod
    def from_json(data: str):
        """
        Load and recreate a ``QuestPlus`` instance from a JSON string.

        Parameters
        ----------
        data
            The JSON string, generated via :meth:`to_json`.

        Returns
        -------
        QuestPlus
            A ``QuestPlus`` instance, generated from the JSON string.

        See Also
        --------
        to_json

        """
        loaded = json_tricks.loads(data)
        loaded.prior = xr.DataArray.from_dict(loaded.prior)
        loaded.posterior = xr.DataArray.from_dict(loaded.posterior)
        loaded.likelihoods = xr.DataArray.from_dict(loaded.likelihoods)

        if loaded._rng is not None:
            state = deepcopy(loaded._rng)
            loaded._rng = np.random.RandomState()
            loaded._rng.set_state(state)

        return loaded

    def __eq__(self, other):
        if not self.likelihoods.equals(other.likelihoods):
            return False

        if not self.prior.equals(other.prior):
            return False

        if not self.posterior.equals(other.posterior):
            return False

        for param_name in self.param_domain.keys():
            if not np.array_equal(
                self.param_domain[param_name], other.param_domain[param_name]
            ):
                return False

        for stim_property in self.stim_domain.keys():
            if not np.array_equal(
                self.stim_domain[stim_property], other.stim_domain[stim_property]
            ):
                return False

        for outcome_name in self.outcome_domain.keys():
            if not np.array_equal(
                self.outcome_domain[outcome_name], other.outcome_domain[outcome_name]
            ):
                return False

        if self.stim_selection != other.stim_selection:
            return False

        if self.stim_selection_options != other.stim_selection_options:
            return False

        if self.stim_scale != other.stim_scale:
            return False

        if self.stim_history != other.stim_history:
            return False

        if self.resp_history != other.resp_history:
            return False

        if self.param_estimation_method != other.param_estimation_method:
            return False

        if self.func != other.func:
            return False

        return True


class QuestPlusWeibull(QuestPlus):
    def __init__(
        self,
        *,
        intensities: Sequence,
        thresholds: Sequence,
        slopes: Sequence,
        lower_asymptotes: Sequence,
        lapse_rates: Sequence,
        prior: Optional[dict] = None,
        responses: Sequence = ("Yes", "No"),
        stim_scale: str = "log10",
        stim_selection_method: str = "min_entropy",
        stim_selection_options: Optional[dict] = None,
        param_estimation_method: str = "mean",
    ):
        """QUEST+ using the Weibull distribution function.

        This is a convenience class that wraps `QuestPlus`.
        """
        super().__init__(
            stim_domain=dict(intensity=intensities),
            param_domain=dict(
                threshold=thresholds,
                slope=slopes,
                lower_asymptote=lower_asymptotes,
                lapse_rate=lapse_rates,
            ),
            outcome_domain=dict(response=responses),
            prior=prior,
            stim_scale=stim_scale,
            stim_selection_method=stim_selection_method,
            stim_selection_options=stim_selection_options,
            param_estimation_method=param_estimation_method,
            func="weibull",
        )

    @property
    def intensities(self) -> np.ndarray:
        """
        Stimulus intensity or contrast domain.
        """
        return self.stim_domain["intensity"]

    @property
    def thresholds(self) -> np.ndarray:
        """
        The threshold parameter domain.
        """
        return self.param_domain["threshold"]

    @property
    def slopes(self) -> np.ndarray:
        """
        The slope parameter domain.
        """
        return self.param_domain["slope"]

    @property
    def lower_asymptotes(self) -> np.ndarray:
        """
        The lower asymptote parameter domain.
        """
        return self.param_domain["lower_asymptote"]

    @property
    def lapse_rates(self) -> np.ndarray:
        """
        The lapse rate parameter domain.
        """
        return self.param_domain["lapse_rate"]

    @property
    def responses(self) -> np.ndarray:
        """
        The response (outcome) domain.
        """
        return self.outcome_domain["response"]

    @property
    def next_intensity(self) -> float:
        """
        The intensity or contrast to present next.
        """
        return super().next_stim["intensity"]

    def update(self, *, intensity: float, response: str) -> None:
        """
        Inform QUEST+ about a newly gathered measurement outcome for a given
        stimulus intensity or contrast, and update the posterior accordingly.

        Parameters
        ----------
        intensity
            The intensity or contrast of the presented stimulus.

        response
            The observed response.

        """
        super().update(stim=dict(intensity=intensity), outcome=dict(response=response))


class QuestPlusThurstone(QuestPlus):
    def __init__(
        self,
        *,
        physical_magnitudes_stim_1: Sequence,
        physical_magnitudes_stim_2: Sequence,
        thresholds: Sequence,
        powers: Sequence,
        perceptual_scale_maxs: Sequence,
        prior: Optional[dict] = None,
        responses: Sequence = ("First", "Second"),
        stim_selection_method: str = "min_entropy",
        stim_selection_options: Optional[dict] = None,
        param_estimation_method: str = "mean",
    ):
        """QUEST+ for Thurstonian scaling.

        This is a convenience class that wraps `QuestPlus`.
        """
        super().__init__(
            stim_domain={
                "physical_magnitudes_stim_1": physical_magnitudes_stim_1,
                "physical_magnitudes_stim_2": physical_magnitudes_stim_2,
            },
            param_domain={
                "threshold": thresholds,
                "power": powers,
                "perceptual_scale_max": perceptual_scale_maxs,
            },
            outcome_domain={"response": responses},
            prior=prior,
            stim_scale=None,
            stim_selection_method=stim_selection_method,
            stim_selection_options=stim_selection_options,
            param_estimation_method=param_estimation_method,
            func="thurstone_scaling",
        )

    @property
    def physical_magnitudes_stim_1(self) -> np.ndarray:
        """
        Physical magnitudes of the first stimulus.
        """
        return self.stim_domain["physical_magnitudes_stim_1"]

    @property
    def physical_magnitudes_stim_2(self) -> np.ndarray:
        """
        Physical magnitudes of the second stimulus.
        """
        return self.stim_domain["physical_magnitudes_stim_2"]

    @property
    def thresholds(self) -> np.ndarray:
        """
        The threshold parameter domain.
        """
        return self.param_domain["threshold"]

    @property
    def powers(self) -> np.ndarray:
        """
        The power parameter domain.
        """
        return self.param_domain["power"]

    @property
    def perceptual_scale_maxss(self) -> np.ndarray:
        """
        The "maximum value of the subjective perceptual scale" parameter domain.
        """
        return self.param_domain["perceptual_scale_max"]

    @property
    def responses(self) -> np.ndarray:
        """
        The response (outcome) domain.
        """
        return self.outcome_domain["response"]

    def update(self, *, stim: dict, response: str) -> None:
        """
        Inform QUEST+ about a newly gathered measurement outcome for a given
        stimulus parameter set, and update the posterior accordingly.

        Parameters
        ----------
        stim
            The stimulus that was used to generate the given outcome.

        outcome
            The observed outcome.

        """
        super().update(stim=stim, outcome=dict(response=response))