import itertools
import warnings

import numpy as np
import pytest

import xarray as xr
from xarray import Dataset
from xarray.testing import _assert_internal_invariants

pytest.importorskip("hypothesis")
pytestmark = pytest.mark.slow_hypothesis

import hypothesis.extra.numpy as npst
import hypothesis.strategies as st
from hypothesis import note, settings
from hypothesis.stateful import (
    RuleBasedStateMachine,
    initialize,
    invariant,
    precondition,
    rule,
)

import xarray.testing.strategies as xrst


@st.composite
def unique(draw, strategy):
    # https://stackoverflow.com/questions/73737073/create-hypothesis-strategy-that-returns-unique-values
    seen = draw(st.shared(st.builds(set), key="key-for-unique-elems"))
    return draw(
        strategy.filter(lambda x: x not in seen).map(lambda x: seen.add(x) or x)
    )


# Share to ensure we get unique names on each draw,
# so we don't try to add two variables with the same name
# or stack to a dimension with a name that already exists in the Dataset.
UNIQUE_NAME = unique(strategy=xrst.names())
DIM_NAME = xrst.dimension_names(name_strategy=UNIQUE_NAME, min_dims=1, max_dims=1)
index_variables = st.builds(
    xr.Variable,
    data=npst.arrays(
        dtype=xrst.pandas_index_dtypes(),
        shape=npst.array_shapes(min_dims=1, max_dims=1),
        elements=dict(allow_nan=False, allow_infinity=False, allow_subnormal=False),
        unique=True,
    ),
    dims=DIM_NAME,
    attrs=xrst.attrs(),
)


def add_dim_coord_and_data_var(ds, var):
    (name,) = var.dims
    # dim coord
    ds[name] = var
    # non-dim coord of same size; this allows renaming
    ds[name + "_"] = var


class DatasetStateMachine(RuleBasedStateMachine):
    # Can't use bundles because we'd need pre-conditions on consumes(bundle)
    # indexed_dims = Bundle("indexed_dims")
    # multi_indexed_dims = Bundle("multi_indexed_dims")

    def __init__(self):
        super().__init__()
        self.dataset = Dataset()
        self.check_default_indexes = True

        # We track these separately as lists so we can guarantee order of iteration over them.
        # Order of iteration over Dataset.dims is not guaranteed
        self.indexed_dims = []
        self.multi_indexed_dims = []

    @initialize(var=index_variables)
    def init_ds(self, var):
        """Initialize the Dataset so that at least one rule will always fire."""
        (name,) = var.dims
        add_dim_coord_and_data_var(self.dataset, var)

        self.indexed_dims.append(name)

    # TODO: stacking with a timedelta64 index and unstacking converts it to object
    @rule(var=index_variables)
    def add_dim_coord(self, var):
        (name,) = var.dims
        note(f"adding dimension coordinate {name}")
        add_dim_coord_and_data_var(self.dataset, var)

        self.indexed_dims.append(name)

    @rule(var=index_variables)
    def assign_coords(self, var):
        (name,) = var.dims
        note(f"assign_coords: {name}")
        self.dataset = self.dataset.assign_coords({name: var})

        self.indexed_dims.append(name)

    @property
    def has_indexed_dims(self) -> bool:
        return bool(self.indexed_dims + self.multi_indexed_dims)

    @rule(data=st.data())
    @precondition(lambda self: self.has_indexed_dims)
    def reset_index(self, data):
        dim = data.draw(st.sampled_from(self.indexed_dims + self.multi_indexed_dims))
        self.check_default_indexes = False
        note(f"> resetting {dim}")
        self.dataset = self.dataset.reset_index(dim)

        if dim in self.indexed_dims:
            del self.indexed_dims[self.indexed_dims.index(dim)]
        elif dim in self.multi_indexed_dims:
            del self.multi_indexed_dims[self.multi_indexed_dims.index(dim)]

    @rule(newname=UNIQUE_NAME, data=st.data(), create_index=st.booleans())
    @precondition(lambda self: bool(self.indexed_dims))
    def stack(self, newname, data, create_index):
        oldnames = data.draw(
            st.lists(
                st.sampled_from(self.indexed_dims),
                min_size=1,
                max_size=3 if create_index else None,
                unique=True,
            )
        )
        note(f"> stacking {oldnames} as {newname}")
        self.dataset = self.dataset.stack(
            {newname: oldnames}, create_index=create_index
        )

        if create_index:
            self.multi_indexed_dims += [newname]

        # if create_index is False, then we just drop these
        for dim in oldnames:
            del self.indexed_dims[self.indexed_dims.index(dim)]

    @rule(data=st.data())
    @precondition(lambda self: bool(self.multi_indexed_dims))
    def unstack(self, data):
        # TODO: add None
        dim = data.draw(st.sampled_from(self.multi_indexed_dims))
        note(f"> unstacking {dim}")
        if dim is not None:
            pd_index = self.dataset.xindexes[dim].index
        self.dataset = self.dataset.unstack(dim)

        del self.multi_indexed_dims[self.multi_indexed_dims.index(dim)]

        if dim is not None:
            self.indexed_dims.extend(pd_index.names)
        else:
            # TODO: fix this
            pass

    @rule(newname=UNIQUE_NAME, data=st.data())
    @precondition(lambda self: bool(self.dataset.variables))
    def rename_vars(self, newname, data):
        dim = data.draw(st.sampled_from(sorted(self.dataset.variables)))
        # benbovy: "skip the default indexes invariant test when the name of an
        # existing dimension coordinate is passed as input kwarg or dict key
        # to .rename_vars()."
        self.check_default_indexes = False
        note(f"> renaming {dim} to {newname}")
        self.dataset = self.dataset.rename_vars({dim: newname})

        if dim in self.indexed_dims:
            del self.indexed_dims[self.indexed_dims.index(dim)]
        elif dim in self.multi_indexed_dims:
            del self.multi_indexed_dims[self.multi_indexed_dims.index(dim)]

    @precondition(lambda self: bool(self.dataset.dims))
    @rule(data=st.data())
    def drop_dims(self, data):
        dims = data.draw(
            st.lists(
                st.sampled_from(sorted(self.dataset.dims)),
                min_size=1,
                unique=True,
            )
        )
        note(f"> drop_dims: {dims}")
        # TODO: dropping a multi-index dimension raises a DeprecationWarning
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", category=DeprecationWarning)
            self.dataset = self.dataset.drop_dims(dims)

        for dim in dims:
            if dim in self.indexed_dims:
                del self.indexed_dims[self.indexed_dims.index(dim)]
            elif dim in self.multi_indexed_dims:
                del self.multi_indexed_dims[self.multi_indexed_dims.index(dim)]

    @precondition(lambda self: bool(self.indexed_dims))
    @rule(data=st.data())
    def drop_indexes(self, data):
        self.check_default_indexes = False

        dims = data.draw(
            st.lists(st.sampled_from(self.indexed_dims), min_size=1, unique=True)
        )
        note(f"> drop_indexes: {dims}")
        self.dataset = self.dataset.drop_indexes(dims)

        for dim in dims:
            if dim in self.indexed_dims:
                del self.indexed_dims[self.indexed_dims.index(dim)]
            elif dim in self.multi_indexed_dims:
                del self.multi_indexed_dims[self.multi_indexed_dims.index(dim)]

    @property
    def swappable_dims(self):
        ds = self.dataset
        options = []
        for dim in self.indexed_dims:
            choices = [
                name
                for name, var in ds._variables.items()
                if var.dims == (dim,)
                # TODO: Avoid swapping a dimension to itself
                and name != dim
            ]
            options.extend(
                (a, b) for a, b in itertools.zip_longest((dim,), choices, fillvalue=dim)
            )
        return options

    @rule(data=st.data())
    # TODO: swap_dims is basically all broken if a multiindex is present
    # TODO: Avoid swapping from Index to a MultiIndex level
    # TODO: Avoid swapping from MultiIndex to a level of the same MultiIndex
    # TODO: Avoid swapping when a MultiIndex is present
    @precondition(lambda self: not bool(self.multi_indexed_dims))
    @precondition(lambda self: bool(self.swappable_dims))
    def swap_dims(self, data):
        ds = self.dataset
        options = self.swappable_dims
        dim, to = data.draw(st.sampled_from(options))
        note(
            f"> swapping {dim} to {to}, found swappable dims: {options}, all_dims: {tuple(self.dataset.dims)}"
        )
        self.dataset = ds.swap_dims({dim: to})

        del self.indexed_dims[self.indexed_dims.index(dim)]
        self.indexed_dims += [to]

    @invariant()
    def assert_invariants(self):
        # note(f"> ===\n\n {self.dataset!r} \n===\n\n")
        _assert_internal_invariants(self.dataset, self.check_default_indexes)


DatasetStateMachine.TestCase.settings = settings(max_examples=300, deadline=None)
DatasetTest = DatasetStateMachine.TestCase


@pytest.mark.skip(reason="failure detected by hypothesis")
def test_unstack_object():
    ds = xr.Dataset()
    ds["0"] = np.array(["", "\x000"], dtype=object)
    ds.stack({"1": ["0"]}).unstack()


@pytest.mark.skip(reason="failure detected by hypothesis")
def test_unstack_timedelta_index():
    ds = xr.Dataset()
    ds["0"] = np.array([0, 1, 2, 3], dtype="timedelta64[ns]")
    ds.stack({"1": ["0"]}).unstack()