from .constants import GROUPBY_SEASONALITIES

OPTIONS = {
    "seasonality": "month",
    "warn_for_failed_PredictionEnsemble_xr_call": True,
    "warn_for_rename_to_climpred_dims": True,
    "warn_for_init_coords_int_to_annual": True,
    "climpred_warnings": True,
}  # defaults

_SEASONALITY_OPTIONS = frozenset(GROUPBY_SEASONALITIES)

_VALIDATORS = {
    "seasonality": _SEASONALITY_OPTIONS.__contains__,
    "warn_for_PredictionEnsemble_xr_call": lambda choice: choice
    in [True, False, "default"],
    "warn_for_rename_to_climpred_dims": lambda choice: choice
    in [True, False, "default"],
    "warn_for_init_coords_int_to_annual": lambda choice: choice
    in [True, False, "default"],
    "climpred_warnings": lambda choice: choice in [True, False, "default"],
}


class set_options:
    """Set options for climpred in a controlled context. Analogous to
    `xarray.set_options(**option) <http://xarray.pydata.org/en/stable/generated/xarray.set_options.html>`_.

    Currently supported options:

    - ``seasonality``
        - Attribute to group dimension ``groupby(f"{dim}.{seasonality}"")``.
            Used in ``reference=climatology`` and
            :py:meth:`~climpred.classes.HindcastEnsemble.remove_bias`.
        - Allowed: [``"dayofyear"``, ``"weekofyear"``, ``"month"``, ``"season"``]
        - Default: ``dayofyear``.
    - ``warn_for_failed_PredictionEnsemble_xr_call``
        - Raise UserWarning when PredictionEnsemble.xr_call,
            e.g. ``.sel(lead=[1])`` fails on one of the datasets.
        - Allowed: [True, False]
        - Default: True
    - ``warn_for_rename_to_climpred_dims``
        - Raise UserWarning when dimensions are renamed to ``CLIMPRED_DIMS`` when
            PredictionEnsemble is instantiated.
        - Allowed: [True, False]
        - Default: True
    - ``warn_for_init_coords_int_to_annual``
        - Raise UserWarning when ``init`` coordinate is of type integer and gets
            converted to annual cftime_range when PredictionEnsemble is instantiated.
        - Allowed: [True, False]
        - Default: True
    - ``climpred_warnings``
        - Overwrites all options containing ``"*warn*"``.
        - Allowed: [True, False]
        - Default: True

    Examples:
        You can use ``set_options`` either as a context manager:

        >>> kw = dict(metric='mse', comparison='e2o', dim='init',
        ...           alignment='same_verifs', reference='climatology')
        >>> with climpred.set_options(seasonality='month'):
        ...     HindcastEnsemble.verify(**kw).SST.sel(skill='climatology')
        <xarray.DataArray 'SST' (lead: 10)>
        array([0.03712573, 0.03712573, 0.03712573, 0.03712573, 0.03712573,
               0.03712573, 0.03712573, 0.03712573, 0.03712573, 0.03712573])
        Coordinates:
          * lead     (lead) int32 1 2 3 4 5 6 7 8 9 10
            skill    <U11 'climatology'

        Or to set global options:

        >>> climpred.set_options(seasonality='month')  # doctest: +ELLIPSIS
        <climpred.options.set_options object at 0x...>
    """

    def __init__(self, **kwargs):
        self.old = {}
        for k, v in kwargs.items():
            if k not in OPTIONS:
                raise ValueError(
                    "argument name %r is not in the set of valid options %r"
                    % (k, set(OPTIONS))
                )
            if k in _VALIDATORS and not _VALIDATORS[k](v):
                if k == "seasonality":
                    expected = f"Expected one of {_SEASONALITY_OPTIONS!r}"
                else:
                    expected = ""
                raise ValueError(
                    f"option {k!r} given an invalid value: {v!r}. " + expected
                )
            self.old[k] = OPTIONS[k]
        self._apply_update(kwargs)

    def _apply_update(self, options_dict):
        if (
            "climpred_warnings" in options_dict
        ):  # climpred_warnings == False overwrites all warnings options
            if not options_dict["climpred_warnings"]:
                for k in [o for o in OPTIONS.keys() if "warn" in o]:
                    options_dict[k] = False
        OPTIONS.update(options_dict)

    def __enter__(self):
        return

    def __exit__(self, type, value, traceback):
        self._apply_update(self.old)
